bytestring-trie 0.1.3 → 0.1.4
raw patch · 2 files changed
+167/−142 lines, 2 filesPVP ok
version bump matches the API change (PVP)
API changes (from Hackage documentation)
Files
- bytestring-trie.cabal +1/−1
- src/Data/Trie/Internal.hs +166/−141
bytestring-trie.cabal view
@@ -3,7 +3,7 @@ ---------------------------------------------------------------- Name: bytestring-trie-Version: 0.1.3+Version: 0.1.4 Cabal-Version: >= 1.2 Build-Type: Simple Stability: beta
src/Data/Trie/Internal.hs view
@@ -165,6 +165,25 @@ -- Trie instances ---------------------------------------------------------------} +instance Binary a => Binary (Trie a) where+ put Empty = put (0 :: Word8)+ put (Arc k m t) = do put (1 :: Word8)+ put k+ put m+ put t+ put (Branch p m l r) = do put (2 :: Word8)+ put p+ put m+ put l+ put r+ + get = do tag <- get :: Get Word8+ case tag of+ 0 -> return Empty+ 1 -> liftM3 Arc get get get+ _ -> liftM4 Branch get get get get++ instance Functor Trie where fmap _ Empty = Empty fmap f (Arc k Nothing t) = Arc k Nothing (fmap f t)@@ -198,6 +217,10 @@ mappend = mergeBy $ \x y -> Just (x `mappend` y) +-- Not a MonadPlus for any definition I can think of.+++-- TODO: cf toListBy. We should provide foldr and foldl directly instance Foldable Trie where foldMap _ Empty = mempty foldMap f (Arc _ Nothing t) = foldMap f t@@ -212,25 +235,29 @@ traverse f (Branch p m l r) = Branch p m <$> traverse f l <*> traverse f r -instance Binary a => Binary (Trie a) where- put Empty = put (0 :: Word8)- put (Arc k m t) = do put (1 :: Word8)- put k- put m- put t- put (Branch p m l r) = do put (2 :: Word8)- put p- put m- put l- put r- - get = do tag <- get :: Get Word8- case tag of- 0 -> return Empty- 1 -> liftM3 Arc get get get- _ -> liftM4 Branch get get get get+{---------------------------------------------------------------+-- Mapping functions+---------------------------------------------------------------} +-- | Apply a function to all values, potentially removing them.+filterMap :: (a -> Maybe b) -> Trie a -> Trie b+filterMap _ Empty = empty+filterMap f (Arc k Nothing t) = arc k Nothing (filterMap f t)+filterMap f (Arc k (Just v) t) = arc k (f v) (filterMap f t)+filterMap f (Branch p m l r) = branch p m (filterMap f l) (filterMap f r) +-- | Generic version of 'fmap'. This function is notably more+-- expensive than 'fmap' or 'filterMap' because we have to reconstruct+-- the keys.+mapBy :: (KeyString -> a -> Maybe b) -> Trie a -> Trie b+mapBy f = go S.empty+ where+ go _ Empty = empty+ go q (Arc k Nothing t) = arc k Nothing (go q' t) where q' = S.append q k+ go q (Arc k (Just v) t) = arc k (f q' v) (go q' t) where q' = S.append q k+ go q (Branch p m l r) = branch p m (go q l) (go q r)++ {--------------------------------------------------------------- -- Smart constructors and helper functions for building tries ---------------------------------------------------------------}@@ -261,8 +288,8 @@ branchMerge _ Empty _ t2 = t2 branchMerge _ t1 _ Empty = t1 branchMerge p1 t1 p2 t2- | zero p1 m = Branch p m t1 t2- | otherwise = Branch p m t2 t1+ | zero p1 m = Branch p m t1 t2+ | otherwise = Branch p m t2 t1 where m = branchMask p1 p2 p = mask p1 m@@ -273,10 +300,10 @@ -- we can see 4/8/?*Word8 at a time instead of just one. -- But that makes maintaining invariants ...difficult :( getPrefix :: Trie a -> Prefix-getPrefix (Branch p _ _ _) = p-getPrefix (Arc k _ _) | S.null k = 0 -- for lack of a better- | otherwise = S.head k-getPrefix Empty = error "getPrefix: no Prefix of Empty"+getPrefix (Branch p _ _ _) = p+getPrefix (Arc k _ _) | S.null k = 0 -- for lack of a better value+ | otherwise = S.head k+getPrefix Empty = error "getPrefix: no Prefix of Empty" {---------------------------------------------------------------@@ -335,8 +362,10 @@ -- TODO: rewrite list-catenation to be lazier (real CPS instead of -- function building? is the function building really better than -- (++) anyways?)+-- -- TODO: the @q@ accumulator should be lazy ByteString and only--- forced by @f@+-- forced by @f@. It's already non-strict, but we should ensure+-- O(n) not O(n^2) when it's forced. -- -- | Convert a trie into a list using a function. Resulting values -- are in sorted order according to the keys.@@ -367,22 +396,22 @@ -- version, see @lookupBy@ in "Data.Trie". lookupBy_ :: (Maybe a -> Trie a -> b) -> b -> (Trie a -> b) -> KeyString -> Trie a -> b-lookupBy_ f z a = let - isBranch (Branch _ _ _ _) = True- isBranch _ = False- in \q t -> if S.null q && isBranch t then f Nothing t else go q t+lookupBy_ f z a = lookupBy_' where- - go _ Empty = z+ -- | Deal with epsilon query (when there is no epsilon value)+ lookupBy_' q t@(Branch _ _ _ _) | S.null q = f Nothing t+ lookupBy_' q t = go q t + -- | The main recursion+ go _ Empty = z+ go q (Arc k mv t) =- case (not $ S.null k', S.null q') of- (True, True) -> a (Arc k' mv t)- (True, False) -> z- (False, True) -> f mv t- (False, False) -> go q' t- where- (_,k',q') = splitMaximalPrefix k q+ let (_,k',q') = splitMaximalPrefix k q+ in case (not $ S.null k', S.null q') of+ (True, True) -> a (Arc k' mv t)+ (True, False) -> z+ (False, True) -> f mv t+ (False, False) -> go q' t go q t_@(Branch _ _ _ _) = findArc t_ where@@ -404,14 +433,14 @@ -- arc k Nothing t === singleton k () >> t -- arc k (Just v) t === singleton k v >>= unionR t . singleton S.empty -- (...except 'arc' doesn't do the invariant correction--- of (>>=) for t=epsilon)+-- of (>>=) for epsilon`elem`t) -- -- | Return the subtrie containing all keys beginning with a prefix. {-# INLINE submap #-} submap :: KeyString -> Trie a -> Trie a submap q = lookupBy_ (arc q) empty (arc q Nothing) q-{- -- Disable superfluous error checking.- -- @submap'@ would replace the first argument to @lookupBy_@+{- -- Disable superfluous error checking.+ -- @submap'@ would replace the first argument to @lookupBy_@ where submap' Nothing Empty = errorEmptyAfterNothing "submap" submap' Nothing (Arc _ _ _) = errorArcAfterNothing "submap"@@ -427,7 +456,7 @@ errorEmptyAfterNothing :: String -> a errorEmptyAfterNothing s = errorInvariantBroken s "Empty after Nothing"--}+-- -} {---------------------------------------------------------------@@ -445,12 +474,18 @@ alterBy :: (KeyString -> a -> Maybe a -> Maybe a) -> KeyString -> a -> Trie a -> Trie a alterBy f_ q_ x_- | S.null q_ = mergeBy (\x y -> f_ q_ x (Just y)) (singleton q_ x_) + | S.null q_ = alterEpsilon | otherwise = go q_ where f = f_ q_ x_ nothing q = arc q (f Nothing) Empty + alterEpsilon t_@Empty = arc q_ (f Nothing) t_+ alterEpsilon t_@(Branch _ _ _ _) = arc q_ (f Nothing) t_+ alterEpsilon t_@(Arc k mv t) | S.null k = arc q_ (f mv) t+ | otherwise = arc q_ (f Nothing) t_+ + go q Empty = nothing q go q t@(Branch p m l r)@@ -461,25 +496,24 @@ qh = errorLogHead "alterBy" q go q t_@(Arc k mv t) =- case (not $ S.null k', S.null q') of- (True, True) -> -- add node to middle of arc- arc p (f Nothing) (Arc k' mv t)- (True, False) ->- case nothing q' of- Empty -> t_ -- Nothing to add, reuse old arc- l -> arc' (branchMerge (getPrefix l) l- (getPrefix r) r)- where- r = Arc k' mv t- - -- inlined version of 'arc'- arc' | S.null p = id- | otherwise = Arc p Nothing- - (False, True) -> arc k (f mv) t- (False, False) -> arc k mv (go q' t)- where- (p,k',q') = splitMaximalPrefix k q+ let (p,k',q') = splitMaximalPrefix k q+ in case (not $ S.null k', S.null q') of+ (True, True) -> -- add node to middle of arc+ arc p (f Nothing) (Arc k' mv t)+ (True, False) ->+ case nothing q' of+ Empty -> t_ -- Nothing to add, reuse old arc+ l -> arc' (branchMerge (getPrefix l) l+ (getPrefix r) r)+ where+ r = Arc k' mv t+ + -- inlined version of 'arc'+ arc' | S.null p = id+ | otherwise = Arc p Nothing+ + (False, True) -> arc k (f mv) t+ (False, False) -> arc k mv (go q' t) {---------------------------------------------------------------@@ -491,102 +525,93 @@ -- where t = map (\s -> (pk s, 0)) -- ["heat","hello","hoi","apple","appa","hell","appb","appc"] ----- TODO: switch to 'go', closing over @f@.--- -- | Combine two tries, using a function to resolve collisions. -- This can only define the space of functions between union and -- symmetric difference but, with those two, all set operations can -- be defined (albeit inefficiently). mergeBy :: (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a-mergeBy _ Empty t1 = t1-mergeBy _ t0 Empty = t0---- /O(n+m)/ for this part where /n/ and /m/ are sizes of the branchings-mergeBy f t0@(Branch p0 m0 l0 r0) t1@(Branch p1 m1 l1 r1)- | shorter m0 m1 = union0- | shorter m1 m0 = union1- | p0 == p1 = branch p0 m0 (mergeBy f l0 l1) (mergeBy f r0 r1)- | otherwise = branchMerge p0 t0 p1 t1+mergeBy f = mergeBy' where- union0 | nomatch p1 p0 m0 = branchMerge p0 t0 p1 t1- | zero p1 m0 = branch p0 m0 (mergeBy f l0 t1) r0- | otherwise = branch p0 m0 l0 (mergeBy f r0 t1)+ -- | Deals with epsilon entries, before recursing into @go@+ mergeBy'+ t0_@(Arc k0 mv0 t0)+ t1_@(Arc k1 mv1 t1)+ | S.null k0 && S.null k1 = arc k0 (mergeMaybe f mv0 mv1) (go t0 t1)+ | S.null k0 = arc k0 mv0 (go t0 t1_)+ | S.null k1 = arc k1 mv1 (go t1 t0_)+ mergeBy'+ (Arc k0 mv0@(Just _) t0)+ t1_@(Branch _ _ _ _)+ | S.null k0 = arc k0 mv0 (go t0 t1_)+ mergeBy'+ t0_@(Branch _ _ _ _)+ (Arc k1 mv1@(Just _) t1)+ | S.null k1 = arc k1 mv1 (go t1 t0_)+ mergeBy' t0_ t1_ = go t0_ t1_ - union1 | nomatch p0 p1 m1 = branchMerge p0 t0 p1 t1- | zero p0 m1 = branch p1 m1 (mergeBy f t0 l1) r1- | otherwise = branch p1 m1 l1 (mergeBy f t0 r1)--mergeBy f t0_ t1_ =- case (t0_,t1_) of- (Arc k0 mv0 t0, Arc k1 mv1 t1)- | S.null k0 && S.null k1 -> arc k0 (mergeMaybe f mv0 mv1)- (mergeBy f t0 t1)- | S.null k0 -> arc k0 mv0 (mergeBy f t0 t1_)- | S.null k1 -> arc k1 mv1 (mergeBy f t1 t0_)- | m' == 0 ->- let (pk,k0',k1') = splitMaximalPrefix k0 k1- in if S.null pk- then error "mergeBy: no mask, but no prefix string"- else let arcMerge mv' t1' t2' = arc pk mv' (mergeBy f t1' t2')- in case (S.null k0', S.null k1') of- (True, True) -> arcMerge (mergeMaybe f mv0 mv1) t0 t1- (True, False) -> arcMerge mv0 t0 (Arc k1' mv1 t1)- (False,True) -> arcMerge mv1 t1 (Arc k0' mv0 t0)- (False,False) -> arcMerge Nothing (Arc k0' mv0 t0)- (Arc k1' mv1 t1) - -- Deal with epsilons. Could be hoisted if we use @go@ style- (Arc k0 mv0@(Just _) t0, Branch _ _ _ _)- | S.null k0 -> arc k0 mv0 (mergeBy f t0 t1_)- (Branch _ _ _ _, Arc k1 mv1@(Just _) t1)- | S.null k1 -> arc k1 mv1 (mergeBy f t1 t0_)+ -- | The main recursion+ go Empty t1 = t1+ go t0 Empty = t0+ + -- /O(n+m)/ for this part where /n/ and /m/ are sizes of the branchings+ go t0@(Branch p0 m0 l0 r0)+ t1@(Branch p1 m1 l1 r1)+ | shorter m0 m1 = union0+ | shorter m1 m0 = union1+ | p0 == p1 = branch p0 m0 (go l0 l1) (go r0 r1)+ | otherwise = branchMerge p0 t0 p1 t1+ where+ union0 | nomatch p1 p0 m0 = branchMerge p0 t0 p1 t1+ | zero p1 m0 = branch p0 m0 (go l0 t1) r0+ | otherwise = branch p0 m0 l0 (go r0 t1) - (Arc _ _ _, Branch _p1 m1 l r)- | nomatch p0 p1 m1 -> branchMerge p1 t1_ p0 t0_- | zero p0 m1 -> branch p1 m1 (mergeBy f t0_ l) r- | otherwise -> branch p1 m1 l (mergeBy f t0_ r)- (Branch _p0 m0 l r, Arc _ _ _)- | nomatch p1 p0 m0 -> branchMerge p0 t0_ p1 t1_- | zero p1 m0 -> branch p0 m0 (mergeBy f t1_ l) r- | otherwise -> branch p0 m0 l (mergeBy f t1_ r)+ union1 | nomatch p0 p1 m1 = branchMerge p0 t0 p1 t1+ | zero p0 m1 = branch p1 m1 (go t0 l1) r1+ | otherwise = branch p1 m1 l1 (go t0 r1) - -- Inlined branchMerge. Both tries are disjoint @Arc@s now.- _ | zero p0 m' -> Branch p' m' t0_ t1_- _ -> Branch p' m' t1_ t0_- where- p0 = getPrefix t0_- p1 = getPrefix t1_- m' = branchMask p0 p1- p' = mask p0 m'+ go t0_ t1_ = go' t0_ t1_+ where+ p0 = getPrefix t0_+ p1 = getPrefix t1_+ m' = branchMask p0 p1+ p' = mask p0 m'+ + go' (Arc k0 mv0 t0)+ (Arc k1 mv1 t1)+ | m' == 0 =+ let (pre,k0',k1') = splitMaximalPrefix k0 k1+ in if S.null pre+ then error "mergeBy: no mask, but no prefix string"+ else let {-# INLINE arcMerge #-}+ arcMerge mv' t1' t2' = arc pre mv' (go t1' t2')+ in case (S.null k0', S.null k1') of+ (True, True) -> arcMerge (mergeMaybe f mv0 mv1) t0 t1+ (True, False) -> arcMerge mv0 t0 (Arc k1' mv1 t1)+ (False,True) -> arcMerge mv1 t1 (Arc k0' mv0 t0)+ (False,False) -> arcMerge Nothing (Arc k0' mv0 t0)+ (Arc k1' mv1 t1)+ go' (Arc _ _ _)+ (Branch _p1 m1 l r)+ | nomatch p0 p1 m1 = branchMerge p1 t1_ p0 t0_+ | zero p0 m1 = branch p1 m1 (go t0_ l) r+ | otherwise = branch p1 m1 l (go t0_ r)+ go' (Branch _p0 m0 l r)+ (Arc _ _ _)+ | nomatch p1 p0 m0 = branchMerge p0 t0_ p1 t1_+ | zero p1 m0 = branch p0 m0 (go t1_ l) r+ | otherwise = branch p0 m0 l (go t1_ r)+ + -- Inlined branchMerge. Both tries are disjoint @Arc@s now.+ go' _ _ | zero p0 m' = Branch p' m' t0_ t1_+ go' _ _ = Branch p' m' t1_ t0_ +{-# INLINE mergeMaybe #-} mergeMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a mergeMaybe _ Nothing Nothing = Nothing mergeMaybe _ Nothing mv1@(Just _) = mv1 mergeMaybe _ mv0@(Just _) Nothing = mv0 mergeMaybe f (Just v0) (Just v1) = f v0 v1--{------------------------------------------------------------------ Mapping functions----------------------------------------------------------------}---- | Generic version of 'fmap'. This function is notably more--- expensive than 'fmap' or 'filterMap' because we have to reconstruct--- the keys.-mapBy :: (KeyString -> a -> Maybe b) -> Trie a -> Trie b-mapBy f = go S.empty- where- go _ Empty = empty- go q (Arc k Nothing t) = arc k Nothing (go q' t) where q' = S.append q k- go q (Arc k (Just v) t) = arc k (f q' v) (go q' t) where q' = S.append q k- go q (Branch p m l r) = branch p m (go q l) (go q r)----- | Apply a function to all values, potentially removing them.-filterMap :: (a -> Maybe b) -> Trie a -> Trie b-filterMap _ Empty = empty-filterMap f (Arc k Nothing t) = arc k Nothing (filterMap f t)-filterMap f (Arc k (Just v) t) = arc k (f v) (filterMap f t)-filterMap f (Branch p m l r) = branch p m (filterMap f l) (filterMap f r) ---------------------------------------------------------------- ----------------------------------------------------------- fin.