{-# LANGUAGE BangPatterns, UnboxedTuples, TupleSections, TypeFamilies, PatternGuards, UnboxedTuples #-} module Data.TrieMap.RadixTrie () where import Data.TrieMap.TrieKey import Data.TrieMap.Sized import Data.TrieMap.Applicative import Control.Applicative import Control.Monad import Data.Maybe import Data.Foldable import Data.Traversable import Prelude hiding (lookup, foldr, foldl) data Edge k m a = Edge {-# UNPACK #-} !Int [k] (Maybe a) (m (Edge k m a)) type Edge' k a = Edge k (TrieMap k) a type MEdge' k a = Maybe (Edge' k a) edgeSize :: Edge k m a -> Int edgeSize (Edge sz _ _ _) = sz instance TrieKey k => TrieKey [k] where newtype TrieMap [k] a = Radix (MEdge' k a) emptyM = Radix Nothing singletonM s ks a = Radix (Just (Edge (s a) ks (Just a) emptyM)) nullM (Radix m) = isNothing m sizeM _ (Radix m) = maybe 0 edgeSize m lookupM ks (Radix m) = m >>= lookup ks alterM s f ks (Radix m) = Radix (alter s f ks m) alterLookupM s f ks (Radix m) = onUnboxed Radix (alterLookupE s f ks) m traverseWithKeyM s f (Radix m) = Radix <$> traverse (traverseE s f) m extractM s f (Radix m) = maybe empty (fmap Radix <.> extractE s f) m foldWithKeyM f (Radix m) z = foldr (foldE f) z m foldlWithKeyM f (Radix m) z = foldl (foldlE f) z m mapMaybeM s f (Radix m) = Radix (m >>= mapMaybeE s f) mapEitherM _ _ _ (Radix Nothing) = (# emptyM, emptyM #) mapEitherM s1 s2 f (Radix (Just m)) = both Radix Radix (mapEitherE s1 s2 f) m unionM s f (Radix m1) (Radix m2) = Radix (unionMaybe (unionE s f) m1 m2) isectM s f (Radix m1) (Radix m2) = Radix (isectMaybe (isectE s f) m1 m2) diffM s f (Radix m1) (Radix m2) = Radix (diffMaybe (diffE s f) m1 m2) -- lookupIxM s ks (Radix m) = maybe (empty, empty, empty) (lookupIxE s 0 ks) m isSubmapM (<=) (Radix m1) (Radix m2) = subMaybe (isSubmapE (<=)) m1 m2 splitLookupM _ _ _ (Radix Nothing) = (# emptyM, Nothing, emptyM #) splitLookupM s f ks (Radix (Just e)) = sides Radix (splitLookupE s f ks) e -- assocAtM s i (Radix m) = maybe (empty, empty, empty) (assocAtE s 0 i) m cat :: [k] -> Edge' k a -> Edge' k a ks `cat` Edge sz ls v ts = Edge sz (ks ++ ls) v ts cons :: k -> Edge' k a -> Edge' k a k `cons` Edge sz ks v ts = Edge sz (k:ks) v ts edge :: TrieKey k => Sized a -> [k] -> Maybe a -> TrieMap k (Edge' k a) -> Edge' k a edge s ks v ts = Edge (maybe 0 s v + sizeM edgeSize ts) ks v ts singleMaybe :: TrieKey k => Sized a -> [k] -> Maybe a -> MEdge' k a singleMaybe s ks v = do v <- v return (edge s ks (Just v) emptyM) compact :: TrieKey k => Edge' k a -> MEdge' k a compact e@(Edge _ ks Nothing ts) = case assocsM ts of [] -> Nothing [(l, e')] -> compact (ks `cat` (l `cons` e')) _ -> Just e compact e = Just e lookup :: (Eq k, TrieKey k) => [k] -> Edge' k a -> Maybe a lookup ks (Edge _ ls v ts) = match ks ls where match (k:ks) (l:ls) | k == l = match ks ls match (k:ks) [] = lookupM k ts >>= lookup ks match [] [] = v match _ _ = Nothing alter :: TrieKey k => Sized a -> (Maybe a -> Maybe a) -> [k] -> MEdge' k a -> MEdge' k a alter s f ks0 Nothing = singleMaybe s ks0 (f Nothing) alter s f ks0 (Just e@(Edge sz ls0 v ts)) = match 0 ks0 ls0 where match !i (k:ks) (l:ls) = case compare k l of LT | Just v' <- f Nothing -> Just $ let sv = s v' in Edge (sv + sz) (take i ls0) Nothing (fromDistAscListM edgeSize [(k, Edge sv ks (Just v') emptyM), (l, Edge sz ls v ts)]) EQ -> match (i+1) ks ls GT | Just v' <- f Nothing -> Just $ let sv = s v' in Edge (sv + sz) (take i ls0) Nothing (fromDistAscListM edgeSize [(l, Edge sz ls v ts), (k, Edge sv ks (Just v') emptyM)]) _ -> Just e match _ (k:ks) [] = compact $ edge s ls0 v (alterM edgeSize g k ts) where g = alter s f ks match _ [] (l:ls) | Just v' <- f Nothing = Just (Edge (s v' + sz) ks0 (Just v') (singletonM edgeSize l (Edge sz ls v ts))) match _ [] [] = compact (edge s ls0 (f v) ts) match _ _ _ = Just e alterLookupE :: TrieKey k => Sized a -> (Maybe a -> (# z, Maybe a #)) -> [k] -> MEdge' k a -> (# z, MEdge' k a #) alterLookupE s f ks Nothing = onUnboxed (singleMaybe s ks) f Nothing alterLookupE s f ks0 (Just e@(Edge sz ls0 v0 ts0)) = match 0 ks0 ls0 where match !i (k:ks) (l:ls) = case compare k l of LT -> onUnboxed (Just . maybe e (\ v' -> let sv = s v' in Edge (sz + sv) (take i ls0) Nothing $ fromDistAscListM edgeSize [(k, Edge sv ks (Just v') emptyM), (l, Edge sz ls v0 ts0)])) f Nothing GT -> onUnboxed (Just . maybe e (\ v' -> let sv = s v' in Edge (sz + sv) (take i ls0) Nothing $ fromDistAscListM edgeSize [(l, Edge sz ls v0 ts0), (k, Edge sv ks (Just v') emptyM)])) f Nothing EQ -> match (i+1) ks ls match _ (k:ks) [] = onUnboxed (compact . edge s ls0 v0) (alterLookupM edgeSize g k) ts0 where g = alterLookupE s f ks match _ [] (l:ls) = onUnboxed (Just . maybe e (\ v' -> let sv = s v' in Edge (sv + sz) ks0 (Just v') (singletonM edgeSize l (Edge sz ls v0 ts0)))) f Nothing match _ [] [] = onUnboxed (\ v' -> compact $ edge s ls0 v' ts0) f v0 traverseE :: (Applicative f, TrieKey k) => Sized b -> ([k] -> a -> f b) -> Edge' k a -> f (Edge' k b) traverseE s f (Edge _ ks v ts) = edge s ks <$> traverse (f ks) v <*> traverseWithKeyM edgeSize g ts where g l = traverseE s (\ ls -> f (ks ++ l:ls)) extractE :: (Alternative f, TrieKey k) => Sized a -> ([k] -> a -> f (x, Maybe a)) -> Edge' k a -> f (x, MEdge' k a) extractE s f (Edge _ ks v ts) = case v of Nothing -> rest Just v -> fmap (\ v' -> compact (edge s ks v' ts)) <$> f ks v <|> rest where rest = fmap (compact . edge s ks v) <$> extractM edgeSize g ts g l = extractE s (\ ls -> f (ks ++ l:ls)) foldE :: TrieKey k => ([k] -> a -> b -> b) -> Edge' k a -> b -> b foldE f (Edge _ ks v ts) z = foldr (f ks) (foldWithKeyM g ts z) v where g l = foldE (\ ls -> f (ks ++ l:ls)) foldlE :: TrieKey k => ([k] -> b -> a -> b) -> b -> Edge' k a -> b foldlE f z (Edge _ ks v ts) = foldlWithKeyM g ts (foldl (f ks) z v) where g l = foldlE (\ ls -> f (ks ++ l:ls)) mapMaybeE :: TrieKey k => Sized b -> ([k] -> a -> Maybe b) -> Edge' k a -> MEdge' k b mapMaybeE s f (Edge _ ks v ts) = compact (edge s ks (v >>= f ks) (mapMaybeM edgeSize (\ l -> mapMaybeE s (\ ls -> f (ks ++ l:ls))) ts)) mapEitherE :: TrieKey k => Sized b -> Sized c -> ([k] -> a -> (# Maybe b, Maybe c #)) -> Edge' k a -> (# MEdge' k b, MEdge' k c #) mapEitherE s1 s2 f (Edge _ ks v ts) = case mapEitherM edgeSize edgeSize (\ l -> mapEitherE s1 s2 (\ ls -> f (ks ++ l:ls))) ts of (# tsL, tsR #) -> case v of Nothing -> (# compact (edge s1 ks Nothing tsL), compact (edge s2 ks Nothing tsR) #) Just v -> case f ks v of (# vL, vR #) -> (# compact (edge s1 ks vL tsL), compact (edge s2 ks vR tsR) #) unionE :: TrieKey k => Sized a -> ([k] -> a -> a -> Maybe a) -> Edge' k a -> Edge' k a -> MEdge' k a unionE s f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match 0 ks0 ls0 where match !i (k:ks) (l:ls) = case compare k l of EQ -> match (i+1) ks ls LT -> Just $ Edge (szK + szL) (take i ks0) Nothing (fromDistAscListM edgeSize [(k, Edge szK ks vK tsK), (l, Edge szL ls vL tsL)]) GT -> Just $ Edge (szK + szL) (take i ks0) Nothing (fromDistAscListM edgeSize [(l, Edge szL ls vL tsL), (k, Edge szK ks vK tsK)]) match _ [] (l:ls) = compact (edge s ks0 vK (alterM edgeSize g l tsK)) where g (Just eK') = unionE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL) g Nothing = Just (Edge szL ls vL tsL) match _ (k:ks) [] = compact (edge s ls0 vL (alterM edgeSize g k tsL)) where g Nothing = Just (Edge szK ks vK tsK) g (Just eL') = unionE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL' match _ [] [] = compact (edge s ls0 (unionMaybe (f ls0) vK vL) (unionM edgeSize g tsK tsL)) where g x = unionE s (\ xs -> f (ls0 ++ x:xs)) isectE :: TrieKey k => Sized c -> ([k] -> a -> b -> Maybe c) -> Edge' k a -> Edge' k b -> MEdge' k c isectE s f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where match (k:ks) (l:ls) | k == l = match ks ls match (k:ks) [] = do eL' <- lookupM k tsL cat ls0 <$> cons k <$> isectE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL' match [] (l:ls) = do eK' <- lookupM l tsK cat ks0 <$> cons l <$> isectE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL) match [] [] = compact (edge s ks0 (isectMaybe (f ks0) vK vL) (isectM edgeSize g tsK tsL)) where g x = isectE s (\ xs -> f (ks0 ++ x:xs)) match _ _ = Nothing diffE :: TrieKey k => Sized a -> ([k] -> a -> b -> Maybe a) -> Edge' k a -> Edge' k b -> MEdge' k a diffE s f eK@(Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where match (k:ks) (l:ls) | k == l = match ks ls match (k:ks) [] | Just eL' <- lookupM k tsL = cat ls0 . cons k <$> diffE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL' match [] (l:ls) = compact (edge s ks0 vK (alterM edgeSize (>>= g) l tsK)) where g eK' = diffE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL) match [] [] = compact (edge s ks0 (diffMaybe (f ks0) vK vL) (diffM edgeSize g tsK tsL)) where g x = diffE s (\ xs -> f (ks0 ++ x:xs)) match _ _ = Just eK isSubmapE :: TrieKey k => LEq a b -> LEq (Edge' k a) (Edge' k b) isSubmapE (<=) (Edge szK ks vK tsK) (Edge _ ls vL tsL) = match ks ls where match (k:ks) (l:ls) | k == l = match ks ls match (k:ks) [] | Just eL' <- lookupM k tsL = isSubmapE (<=) (Edge szK ks vK tsK) eL' match [] [] = subMaybe (<=) vK vL && isSubmapM (isSubmapE (<=)) tsK tsL match _ _ = False splitLookupE :: TrieKey k => Sized a -> (a -> (# Maybe a, Maybe x, Maybe a #)) -> [k] -> Edge' k a -> (# MEdge' k a, Maybe x, MEdge' k a #) splitLookupE s f ks e@(Edge _ ls v ts) = match ks ls where match (k:ks) (l:ls) = case compare k l of LT -> (# Nothing, Nothing, Just e #) GT -> (# Just e, Nothing, Nothing #) EQ -> match ks ls match (k:ks) [] = case splitLookupM edgeSize g k ts of (# tsL, x, tsR #) -> (# compact (edge s ls v tsL), x, compact (edge s ls Nothing tsR) #) where g = splitLookupE s f ks match [] (_:_) = (# Nothing, Nothing, Just e #) match [] [] = case v of Nothing -> (# Nothing, Nothing, compact (edge s ls Nothing ts) #) Just v -> case f v of (# vL, x, vR #) -> (# singleMaybe s ls vL, x, compact (edge s ls vR ts) #)