module Data.TrieMap.Regular.RadixTrie () where
import Data.TrieMap.Regular.Class
import Data.TrieMap.Regular.Base
import Data.TrieMap.Regular.Ord
import Data.TrieMap.Regular.Eq
import Data.TrieMap.Sized
import Data.TrieMap.TrieKey
import Data.TrieMap.Applicative
import Data.TrieMap.CPair
import Control.Arrow
import Control.Applicative
import Control.Monad
import Data.Maybe
import Data.Monoid
import Data.Foldable
import Data.Traversable
import Prelude hiding (foldr, foldl)
data Edge f (m :: * -> * -> *) k a = Edge !Int [f k] (Maybe (a)) (m k (Edge f m k a))
type Edge' f k a = Edge f (TrieMapT f) k a
type MEdge f k m a = Maybe (Edge f m k a)
type MEdge' f k a = Maybe (Edge' f k a)
newtype RadixTrie f k a = Radix (MEdge' f k a)
type instance TrieMapT (L f) = RadixTrie f
type instance TrieMap (L f r) = RadixTrie f r
edgeSize :: Sized (Edge f m k a)
edgeSize (Edge s _ _ _) = s
edge :: (TrieKeyT f m, m ~ TrieMapT f, TrieKey k (TrieMap k)) => Sized a -> [f k] -> Maybe (a) -> m k (Edge f m k a) -> Edge f m k a
edge s ks v ts = Edge (maybe 0 s v + sizeT edgeSize ts) ks v ts
instance (OrdT f, TrieKeyT f m, m ~ TrieMapT f) => TrieKeyT (L f) (RadixTrie f) where
emptyT = Radix Nothing
nullT (Radix m) = isNothing m
sizeT _ (Radix m) = maybe 0 edgeSize m
lookupT (List ks) (Radix m) = m >>= lookupE ks
lookupIxT s (List ks) (Radix m) = maybe (mzero, mzero, mzero) (onKey List . lookupIxE s 0 ks) m
assocAtT s i (Radix m) = maybe (mzero, mzero, mzero) (onKey List . assocAtE s 0 i) m
alterT s f (List ks) (Radix m) = Radix (maybe (singletonME s ks (f Nothing)) (alterE s f ks) m)
alterLookupT s f (List ks) (Radix m) = Radix <$> maybe (singletonME s ks <$> f Nothing) (alterLookupE s f ks) m
traverseWithKeyT s f (Radix m) = Radix <$> traverse (traverseE s (f . List)) m
foldWithKeyT f (Radix m) z = foldr (foldE (f . List)) z m
foldlWithKeyT f (Radix m) z = foldr (foldlE (f . List)) z m
mapEitherT s1 s2 f (Radix m) = (Radix *** Radix) (maybe (Nothing, Nothing) (mapEitherE s1 s2 (f . List)) m)
splitLookupT s f (List ks) (Radix m) = Radix `sides` maybe (Nothing, Nothing, Nothing) (splitLookupE s f ks) m
unionT s f (Radix m1) (Radix m2) = Radix (unionMaybe (unionE s (f . List)) m1 m2)
isectT s f (Radix m1) (Radix m2) = Radix (isectMaybe (isectE s (f . List)) m1 m2)
diffT s f (Radix m1) (Radix m2) = Radix (diffMaybe (diffE s (f . List)) m1 m2)
extractT s f (Radix m) = maybe empty (fmap Radix <.> extractE s (f . List)) m
isSubmapT (<=) (Radix m1) (Radix m2) = subMaybe (isSubEdge (<=)) m1 m2
fromListT s f xs = Radix (fromListE s (f . List) [(ks, a) | (List ks, a) <- xs])
fromAscListT s f xs = Radix (fromAscListE s (f . List) [(ks, a) | (List ks, a) <- xs])
instance (OrdT f, TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => TrieKey (L f k) (RadixTrie f k) where
emptyM = emptyT
nullM = nullT
sizeM = sizeT
lookupM = lookupT
lookupIxM = lookupIxT
assocAtM = assocAtT
alterM = alterT
alterLookupM = alterLookupT
traverseWithKeyM = traverseWithKeyT
foldWithKeyM = foldWithKeyT
foldlWithKeyM = foldlWithKeyT
mapEitherM = mapEitherT
splitLookupM = splitLookupT
unionM = unionT
isectM = isectT
diffM = diffT
extractM = extractT
isSubmapM = isSubmapT
fromListM = fromListT
fromAscListM = fromAscListT
fromDistAscListM = fromDistAscListT
compact :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Edge' f k a -> MEdge' f k a
compact e@(Edge s ks Nothing ts) = case assocsT ts of
[] -> Nothing
[~(k, e'@(Edge s' ls v ts'))]
-> e' `seq` compact (Edge s' (ks ++ k:ls) v ts')
_ -> Just e
compact e = Just e
cons :: f k -> Edge' f k a -> Edge' f k a
l `cons` Edge s ls v ts = Edge s (l:ls) v ts
cat :: [f k] -> Edge' f k a -> Edge' f k a
ks `cat` Edge s ls v ts = Edge s (ks ++ ls) v ts
singletonME :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Sized a -> [f k] -> Maybe (a) -> MEdge' f k a
singletonME s ks = fmap (\ v -> Edge (s v) ks (Just v) emptyT)
lookupE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => [f k] -> Edge' f k a -> Maybe (a)
lookupE ks (Edge _ ls v ts) = match ks ls where
match (k:ks) (l:ls)
| k `eqT` l = match ks ls
match (k:ks) [] = do e' <- lookupT k ts
lookupE ks e'
match [] [] = v
match _ _ = Nothing
alterE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) =>
Sized a -> (Maybe (a) -> Maybe (a)) -> [f k] -> Edge' f k a -> MEdge' f k a
alterE s f ks0 e@(Edge sz ls0 v0 ts0) = match 0 ks0 ls0 where
match i _ _ | i `seq` False = undefined
match i (k:ks) (l:ls)
| k `eqT` l = match (i+1) ks ls
| Just v <- f Nothing
= Just (Edge (sz + s v) (take i ls0) Nothing
(fromListT edgeSize (const const) [(k, Edge (s v) ks (Just v) emptyT),
(l, Edge sz ls v0 ts0)]))
match _ (k:ks) [] = compact $ edge s ls0 v0 $ alterT edgeSize g k ts0 where
g = maybe (singletonME s ks (f Nothing)) (alterE s f ks)
match _ [] (l:ls)
| Just v <- f Nothing
= Just (Edge (sz + s v) ks0 (Just v) (singletonT edgeSize l (Edge sz ls v0 ts0)))
match _ [] [] = compact (edge s ls0 (f v0) ts0)
match _ _ _ = Just e
alterLookupE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) =>
Sized a -> (Maybe a -> CPair x (Maybe a)) -> [f k] -> Edge' f k a -> CPair x (MEdge' f k a)
alterLookupE s f ks0 e@(Edge sz ls0 v0 ts0) = match 0 ks0 ls0 where
match i _ _ | i `seq` False = undefined
match i (k:ks) (l:ls) = case compareT k l of
LT -> fmap (Just . maybe e (\ v' -> let sv = s v' in Edge (sz + sv) (take i ls0) Nothing $
fromDistAscListT edgeSize [(k, Edge sv ks (Just v') emptyT), (l, Edge sz ls v0 ts0)]))
(f Nothing)
GT -> fmap (Just . maybe e (\ v' -> let sv = s v' in Edge (sz + sv) (take i ls0) Nothing $
fromDistAscListT edgeSize [(l, Edge sz ls v0 ts0), (k, Edge sv ks (Just v') emptyT)]))
(f Nothing)
EQ -> match (i+1) ks ls
match _ (k:ks) [] = fmap (compact . edge s ls0 v0) (alterLookupT edgeSize g k ts0) where
g = maybe (singletonME s ks <$> f Nothing) (alterLookupE s f ks)
match _ [] (l:ls) = fmap (Just . maybe e (\ v' -> Edge (sz + s v') ks0 (Just v') (singletonT edgeSize l (Edge sz ls v0 ts0))))
(f Nothing)
match _ [] [] = fmap (\ v' -> compact (edge s ls0 v' ts0)) (f v0)
traverseE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k), Applicative t) =>
Sized b -> ([f k] -> a -> t (b)) -> Edge' f k a -> t (Edge' f k b)
traverseE s f (Edge _ ks v ts) =
edge s ks <$> traverse (f ks) v <*> traverseWithKeyT edgeSize (\ l -> traverseE s (\ ls -> f (ks ++ l:ls))) ts
foldE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => ([f k] -> a -> b -> b) -> Edge' f k a -> b -> b
foldE f (Edge _ ks v ts) z = foldr (f ks) (foldWithKeyT (\ l -> foldE (\ ls -> f (ks ++ l:ls))) ts z) v
foldlE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => ([f k] -> b -> a -> b) -> Edge' f k a -> b -> b
foldlE f (Edge _ ks v ts) z = foldlWithKeyT (\ l z m -> foldlE (\ ls -> f (ks ++ l:ls)) m z) ts (foldl (f ks) z v)
mapEitherE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Sized b -> Sized c ->
EitherMap (EitherMap [f k] (a) (b) (c)) (Edge' f k a) (Edge' f k b) (Edge' f k c)
mapEitherE s1 s2 f (Edge _ ks v ts) = case (maybe (Nothing, Nothing) (f ks) v, mapEitherT edgeSize edgeSize
(\ l -> mapEitherE s1 s2 (\ ls -> f (ks ++ l:ls))) ts) of
((vL, vR), (tsL, tsR)) -> (compact (edge s1 ks vL tsL), compact (edge s2 ks vR tsR))
splitLookupE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Sized a -> SplitMap (a) x -> [f k] -> SplitMap (Edge' f k a) x
splitLookupE s f ks e@(Edge _ ls v ts) = match ks ls where
match (k:ks) (l:ls) = case compareT k l of
LT -> (Nothing, Nothing, Just e)
EQ -> match ks ls
GT -> (Just e, Nothing, Nothing)
match [] [] = case v of
Nothing -> (Nothing, Nothing, Just e)
Just v -> compact `sides` case f v of
(vL, x, vR) -> (edge s ls vL emptyT, x, edge s ls vR ts)
match [] (l:ls) = (Just e, Nothing, Nothing)
match (k:ks) [] = compact `sides` case splitLookupT edgeSize g k ts of
(tsL, x, tsR) -> (edge s ls v tsL, x, edge s ls Nothing tsR)
where g = splitLookupE s f ks
unionE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Sized a -> UnionFunc (UnionFunc [f k] (a)) (Edge' f k a)
unionE s f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match 0 ks0 ls0 where
match i _ _ | i `seq` False = undefined
match i (k:ks) (l:ls)
| k `eqT` l = match (i+1) ks ls
| otherwise = Just (Edge (szK + szL) (take i ks0) Nothing
(fromListT edgeSize (const const) [(k, Edge szK ks vK tsK), (l, Edge szL ls vL tsL)]))
match _ (k:ks) [] = compact (edge s ls0 vL $ alterT edgeSize g k tsL) where
g Nothing = Just (Edge szK ks vK tsK)
g (Just e) = unionE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) e
match _ [] (l:ls) = compact (edge s ks0 vK $ alterT edgeSize g l tsK) where
g Nothing = Just (Edge szL ls vL tsL)
g (Just e) = unionE s (\ ls' -> f (ks0 ++ l:ls')) e (Edge szL ls vL tsL)
match _ [] [] = compact (edge s ks0 (unionMaybe (f ks0) vK vL) (unionT edgeSize g tsK tsL)) where
g x = unionE s (\ xs -> f (ks0 ++ x:xs))
extractE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k), Alternative t) => Sized a -> ([f k] -> a -> t (CPair x (Maybe a))) ->
Edge' f k a -> t (CPair x (MEdge' f k a))
extractE s f (Edge _ ks v ts) = (maybe empty (fmap (\ v' -> compact (edge s ks v' ts)) <.> f ks) v) <|>
(fmap (compact . edge s ks Nothing) <$> extractT edgeSize g ts)
where g l = extractE s (\ ls -> f (ks ++ l:ls))
aboutE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k), Alternative t) => ([f k] -> a -> t x) ->
Edge' f k a -> t x
aboutE f = cpFst <.> extractE (const 0) (\ k a -> fmap (flip cP Nothing) (f k a))
isectE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Sized c ->
IsectFunc (IsectFunc [f k] (a) (b) (c)) (Edge' f k a) (Edge' f k b) (Edge' f k c)
isectE s f (Edge szK ks vK tsK) (Edge szL ls vL tsL) = match ks ls where
match (k:ks) (l:ls)
| k `eqT` l = match ks ls
match (k:ks) [] = do e' <- lookupT k tsL
liftM (cat ls . cons k) (isectE s (\ ks' -> f (ls ++ k:ks')) (Edge szK ks vK tsK) e')
match [] (l:ls) = do e' <- lookupT l tsK
liftM (cat ks . cons l) (isectE s (\ ls' -> f (ks ++ l:ls')) e' (Edge szL ls vL tsL))
match [] [] = compact (edge s ks (isectMaybe (f ks) vK vL) (isectT edgeSize g tsK tsL)) where
g x = isectE s (\ xs -> f (ks ++ x:xs))
diffE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Sized a ->
DiffFunc (DiffFunc [f k] (a) (b)) (Edge' f k a) (Edge' f k b)
diffE s f e@(Edge szK ks vK tsK) (Edge szL ls vL tsL) = match ks ls where
match (k:ks) (l:ls)
| k `eqT` l = match ks ls
match (k:ks) []
| Just e' <- lookupT k tsL
= fmap (cat ls . cons k) (diffE s (\ ks' -> f (ls ++ k:ks')) (Edge szK ks vK tsK) e')
match [] (l:ls) = compact (edge s ks vK (alterT edgeSize (>>= g) l tsK)) where
g e' = diffE s (\ ls' -> f (ks ++ l:ls')) e' (Edge szL ls vL tsL)
match [] [] = compact (edge s ks (diffMaybe (f ks) vK vL) (diffT edgeSize g tsK tsL)) where
g x = diffE s (\ xs -> f (ks ++ x:xs))
isSubEdge :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => LEq (a) (b) -> LEq (Edge' f k a) (Edge' f k b)
isSubEdge (<=) (Edge szK ks vK tsK) (Edge szL ls vL tsL) = match ks ls where
match (k:ks) (l:ls)
| k `eqT` l = match ks ls
match (k:ks) []
| Just e' <- lookupT k tsL
= isSubEdge (<=) (Edge szK ks vK tsK) e'
match [] []
= subMaybe (<=) vK vL && isSubmapT (isSubEdge (<=)) tsK tsL
match _ _ = False
filterer :: (k -> k -> Bool) -> (a -> a -> a) -> [([k], a)] -> (Maybe a, [(k, [([k], a)])])
filterer (==) f = filterer' where
filterer' (([], a):xs) = first (Just . maybe a (flip f a)) (filterer' xs)
filterer' ((k:ks, a):xs) = second (cons k ks a) (filterer' xs)
filterer' [] = (Nothing, [])
cons k ks a [] = [(k, [(ks, a)])]
cons k ks a ys0@((k', xs):ys)
| k == k' = (k', (ks,a):xs):ys
| otherwise = (k, [(ks, a)]):ys0
fromListE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Sized a -> ([f k] -> a -> a -> a) -> [([f k], a)] -> MEdge' f k a
fromListE _ _ [] = Nothing
fromListE s f xs = case filterer eqT (f []) xs of
(Nothing, [(k, xs)]) -> cons k <$> fromListE s (f . (k:)) xs
(v, xss) -> Just (edge s [] v (mapWithKeyT edgeSize (\ k (K0 xs) -> fromJust (fromListE s (f . (k:)) xs))
(fromListT (const 1) (\ _ (K0 xs) (K0 ys) -> K0 (ys ++ xs)) [(k, K0 xs) | (k, xs) <- xss])))
fromAscListE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) =>
Sized a -> ([f k] -> a -> a -> a) -> [([f k], a)] -> MEdge' f k a
fromAscListE _ _ [] = Nothing
fromAscListE s f xs = case filterer eqT (f []) xs of
(Nothing, [(k, xs)]) -> cons k <$> fromAscListE s (f . (k:)) xs
(v, xss) -> Just (edge s [] v (fromDistAscListT edgeSize [(k, fromJust (fromAscListE s (f . (k:)) xs)) | (k, xs) <- xss]))
lookupIxE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) =>
Sized a -> Int -> [f k] -> Edge' f k a -> IndexPos [f k] a
lookupIxE _ i _ _ | i `seq` False = undefined
lookupIxE s i ks e@(Edge _ ls v ts) = match ks ls where
match (k:ks) (l:ls) = case compareT k l of
LT -> (mzero, mzero, getMin (Asc i) e)
EQ -> match ks ls
GT -> (getMax (Asc i) e, mzero, mzero)
match (k:ks) [] = case lookupIxT edgeSize k ts of
(lb, x, ub) -> let lookupX = do Asc iK k' e' <- x
let (lb', x', ub') = lookupIxE s (i + iK) ks e'
let f = onKeyA ((ls ++) . (k' :))
return (f <$> lb', f <$> x', f <$> ub')
in ((do Asc iL kL eL <- lb
getMax (\ ksL -> Asc (i + iL) (ls ++ kL:ksL)) eL) <|>
(do (lb', _, _) <- Last lookupX
lb'),
(do (_, x', _) <- lookupX
x'),
(do (_, _, ub') <- First lookupX
ub') <|>
(do Asc iR kR eR <- ub
getMin (\ ksR -> Asc (i + iR) (ls ++ kR:ksR)) eR))
match [] [] = (mzero, Asc i ls <$> v, aboutT
(\ x -> aboutE (\ xs v' -> return (Asc (i + maybe 0 s v) (ls ++ x:xs) v'))) ts)
match [] _ = (mzero, mzero, getMin (Asc i) e)
getMin f = aboutE (return .: f)
getMax f = aboutE (return .: f)
assocAtE :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) =>
Sized a -> Int -> Int -> Edge' f k a -> IndexPos [f k] a
assocAtE s i0 i (Edge _ ks Nothing ts) = case assocAtT edgeSize i ts of
(lb, x, ub) -> let lookupX = do Asc i' l e' <- x
return (onKey ((ks ++) . (l:)) (assocAtE s (i0 + i') (i i') e'))
in ((do Asc iL lL eL <- lb
getMax (\ ls -> Asc (i0 + iL) (ks ++ lL:ls)) eL) <|>
(do (lb', _, _) <- Last lookupX
lb'),
(do (_, x', _) <- lookupX
x'),
(do (_, _, ub') <- First lookupX
ub') <|>
(do Asc iR lR eR <- ub
getMin (\ ls -> Asc (i0 + iR) (ks ++ lR:ls)) eR))
where getMin f e = aboutE (return .: f) e
getMax f e = aboutE (return .: f) e
assocAtE s i0 i (Edge _ ks (Just v) ts)
| i < sv = (mzero, return (Asc i ks v), aboutT (\ l -> aboutE (\ ls v' -> return (Asc (i0 + sv) (ks ++ l:ls) v'))) ts)
| (lb, x, ub) <- assocAtT edgeSize (i sv) ts
= let lookupX = do Asc i' l e' <- x
return (onKey ((ks ++) . (l:)) (assocAtE s (i0 + i' + sv) (i sv i') e'))
in ((do Asc iL lL eL <- lb
getMax (\ ls -> Asc (i0 + iL + sv) (ks ++ lL:ls)) eL) <|>
(do (lb', _, _) <- Last lookupX
lb'),
(do (_, x', _) <- lookupX
x'),
(do (_, _, ub') <- First lookupX
ub') <|>
(do Asc iR lR eR <- ub
getMin (\ ls -> Asc (i0 + iR + sv) (ks ++ lR:ls)) eR))
where getMin f = aboutE (return .: f)
getMax f = aboutE (return .: f)
sv = s v