diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright (c) 2008--2009, wren ng thornton. +Copyright (c) 2008, 2009, 2010, 2011 wren ng thornton. ALL RIGHTS RESERVED. Redistribution and use in source and binary forms, with or without diff --git a/bytestring-trie.cabal b/bytestring-trie.cabal --- a/bytestring-trie.cabal +++ b/bytestring-trie.cabal @@ -1,13 +1,14 @@ ---------------------------------------------------------------- --- wren ng thornton ~ 2010.06.10 +-- wren ng thornton ~ 2010.11.12 ---------------------------------------------------------------- Name: bytestring-trie -Version: 0.2.2 -Cabal-Version: >= 1.2 +Version: 0.2.3 +-- Source-Repository requires version 1.6 +Cabal-Version: >= 1.6 Build-Type: Simple Stability: provisional -Copyright: Copyright (c) 2008--2010 wren ng thornton +Copyright: Copyright (c) 2008--2011 wren ng thornton License: BSD3 License-File: LICENSE Author: wren ng thornton @@ -24,27 +25,45 @@ Patricia trees have efficient algorithms for union and other merging operations, but they're also quick for lookups and insertions. + . + If you are only interested in being able to associate + strings to values, then you may prefer the @hashmap@ + package which is faster for those only needing a + map-like structure. This package is intended for + those who need the extra capabilities that a trie-like + structure can offer (e.g., structure sharing to + reduce memory costs for highly redundant keys, + taking the submap of all keys with a given prefix, + contextual mapping, extracting the minimum and + maximum keys, etc.) +Source-Repository head + Type: darcs + Location: http://community.haskell.org/~wren/bytestring-trie + +---------------------------------------------------------------- Flag base4 - Description: base-4.0 deprecated Prelude which is imported qualified Default: True + Description: base-4.0 emits "Prelude deprecated" messages in + order to get people to be explicit about which + version of base they use. Flag useCinternal + Default: False Description: Use optimized C implementation for indexOfDifference. See notes in Data.Trie.ByteStringInternal. - Default: False Flag applicativeInBase - Description: Applicative functors were added in base-2.0 Default: True + Description: Applicative functors were added in base-2.0 Flag bytestringInBase + Default: False Description: The bytestring library was included in base-2.0 and base-2.1.1, but for base-1.0 and base-3.0 it was a separate package - Default: False - +---------------------------------------------------------------- Library Hs-Source-Dirs: src Exposed-Modules: Data.Trie @@ -52,7 +71,11 @@ , Data.Trie.Convenience Other-Modules: Data.Trie.BitTwiddle , Data.Trie.ByteStringInternal + , Data.Trie.Errors + Build-Depends: binary + -- I think this is all that needs doing to get rid of the warnings? + -- BUG: looks like it's not enough if flag(base4) Build-Depends: base >= 4 && < 5 else @@ -60,10 +83,6 @@ if flag(bytestringInBase) Build-Depends: base >= 2.0 && < 2.2 - Cpp-Options: -DBYTESTRING_IN_BASE - -- BUG (Cabal 1.2 + Haddock): enable for Haddock, disable - -- for Hackage. Fixed in Cabal 1.6 - --Ghc-Options: -DBYTESTRING_IN_BASE else Build-Depends: base < 2.0 || >= 3, bytestring @@ -75,8 +94,6 @@ --Ghc-Options: -DAPPLICATIVE_IN_BASE else Build-Depends: base < 2.0 - - Build-Depends: binary if flag(useCinternal) C-Sources: src/Data/Trie/ByteStringInternal/indexOfDifference.c diff --git a/src/Data/Trie.hs b/src/Data/Trie.hs --- a/src/Data/Trie.hs +++ b/src/Data/Trie.hs @@ -1,10 +1,11 @@ -{-# OPTIONS_GHC -Wall -fwarn-tabs #-} - +-- To make GHC stop warning about the Prelude +{-# OPTIONS_GHC -Wall -fwarn-tabs -fno-warn-unused-imports #-} +{-# LANGUAGE NoImplicitPrelude #-} ---------------------------------------------------------------- --- ~ 2009.01.20 +-- ~ 2011.02.12 -- | -- Module : Data.Trie --- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- Copyright : Copyright (c) 2008--2011 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org -- Stability : experimental @@ -17,16 +18,19 @@ -- elements. For further details on the latter, see -- -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\", --- Workshop on ML, September 1998, pages 77-86, --- +-- Workshop on ML, September 1998, pages 77-86, +-- -- -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve/ --- /Information Coded In Alphanumeric/\", Journal of the ACM, 15(4), --- October 1968, pages 514-534. +-- /Information Coded In Alphanumeric/\", Journal of the ACM, 15(4), +-- October 1968, pages 514-534. -- -- This module aims to provide an austere interface, while being -- detailed enough for most users. For an extended interface with --- many additional functions, see "Data.Trie.Convenience". +-- many additional functions, see "Data.Trie.Convenience". For +-- functions that give more detailed (potentially abstraction-breaking) +-- access to the data strucuture, or for experimental functions +-- which aren't quite ready for the public API, see "Data.Trie.Internal". ---------------------------------------------------------------- module Data.Trie @@ -53,14 +57,14 @@ , mapBy, filterMap ) where -import Prelude hiding (null, lookup) -import qualified Prelude +import Prelude hiding (null, lookup) +import qualified Prelude (null, lookup) import Data.Trie.Internal -import Data.ByteString (ByteString) - -import Data.Maybe (isJust) -import Control.Monad (liftM) +import Data.Trie.Errors (impossible) +import Data.ByteString (ByteString) +import Data.Maybe (isJust) +import Control.Monad (liftM) ---------------------------------------------------------------- ---------------------------------------------------------------- @@ -121,20 +125,20 @@ -- | Insert a new key. If the key is already present, overrides the -- old value -insert :: ByteString -> a -> Trie a -> Trie a +insert :: ByteString -> a -> Trie a -> Trie a {-# INLINE insert #-} -insert = alterBy (\_ x _ -> Just x) - +insert = alterBy (\_ x _ -> Just x) + -- | Apply a function to the value at a key. -adjust :: (a -> a) -> ByteString -> Trie a -> Trie a +adjust :: (a -> a) -> ByteString -> Trie a -> Trie a {-# INLINE adjust #-} -adjust f q = alterBy (\_ _ -> liftM f) q undefined --- TODO: use adjustBy, and benchmark differences +adjust f q = adjustBy (\_ _ -> f) q (impossible "adjust") +-- TODO: benchmark vs the definition with alterBy/liftM -- | Remove the value stored at a key. -delete :: ByteString -> Trie a -> Trie a +delete :: ByteString -> Trie a -> Trie a {-# INLINE delete #-} -delete q = alterBy (\_ _ _ -> Nothing) q undefined +delete q = alterBy (\_ _ _ -> Nothing) q (impossible "delete") {--------------------------------------------------------------- diff --git a/src/Data/Trie/BitTwiddle.hs b/src/Data/Trie/BitTwiddle.hs --- a/src/Data/Trie/BitTwiddle.hs +++ b/src/Data/Trie/BitTwiddle.hs @@ -7,7 +7,7 @@ -- ~ 2009.01.05 -- | -- Module : Data.Trie.BitTwiddle --- Copyright : Copyright (c) Daan Leijen 2002 +-- Copyright : Copyright (c) 2002 Daan Leijen -- License : BSD3 -- Maintainer : libraries@haskell.org, wren@community.haskell.org -- Stability : stable @@ -47,11 +47,11 @@ elemToNat :: KeyElem -> Word {-# INLINE elemToNat #-} -elemToNat i = fromIntegral i +elemToNat = fromIntegral natToElem :: Word -> KeyElem {-# INLINE natToElem #-} -natToElem w = fromIntegral w +natToElem = fromIntegral shiftRL :: Word -> Int -> Word {-# INLINE shiftRL #-} @@ -67,6 +67,8 @@ -- Endian independent bit twiddling (Trie endianness, not architecture) ---------------------------------------------------------------} +-- TODO: should we use the (Bits Word8) instance instead of 'elemToNat' and (Bits Nat)? We need to compare Core, C--, or ASM in order to decide this. The choice will apply to 'zero', 'mask', 'maskW',... If we shouldn't, then we should probably send a patch upstream to fix the (Bits Word8) instance. + -- | Is the value under the mask zero? zero :: KeyElem -> Mask -> Bool {-# INLINE zero #-} @@ -92,6 +94,10 @@ maskW :: Word -> Word -> Prefix {-# INLINE maskW #-} maskW i m = natToElem (i .&. (complement (m-1) `xor` m)) +-- TODO: try the alternatives mentioned in the Containers paper: +-- \i m -> natToElem (i .&. (negate m - m)) +-- \i m -> natToElem (i .&. (m * complement 1)) +-- N.B. these return /all/ the low bits, and therefore they are not equal functions for all m. They are, however, equal when only one bit of m is set. -- | Determine whether the first mask denotes a shorter prefix than -- the second. diff --git a/src/Data/Trie/ByteStringInternal.hs b/src/Data/Trie/ByteStringInternal.hs --- a/src/Data/Trie/ByteStringInternal.hs +++ b/src/Data/Trie/ByteStringInternal.hs @@ -15,7 +15,7 @@ -- ~ 2009.02.06 -- | -- Module : Data.Trie.ByteStringInternal --- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- Copyright : Copyright (c) 2008--2011 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org -- Stability : experimental @@ -117,13 +117,14 @@ -- Use the naive algorithm which doesn't depend on architecture details indexOfDifference p1 p2 limit = goByte 0 where - goByte n = if n >= limit - then return limit - else do c1 <- peekElemOff p1 n - c2 <- peekElemOff p2 n - if c1 == c2 - then goByte $! n+1 - else return n + goByte n = + if n >= limit + then return limit + else do c1 <- peekElemOff p1 n + c2 <- peekElemOff p2 n + if c1 == c2 + then goByte $! n+1 + else return n {- #endif -} diff --git a/src/Data/Trie/ByteStringInternal/indexOfDifference.c b/src/Data/Trie/ByteStringInternal/indexOfDifference.c --- a/src/Data/Trie/ByteStringInternal/indexOfDifference.c +++ b/src/Data/Trie/ByteStringInternal/indexOfDifference.c @@ -2,7 +2,7 @@ -- ~ 2009.01.07 -- | -- Module : Data.Trie.ByteStringInternal.indexOfDifference --- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- Copyright : Copyright (c) 2008--2011 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org -- Stability : beta diff --git a/src/Data/Trie/Convenience.hs b/src/Data/Trie/Convenience.hs --- a/src/Data/Trie/Convenience.hs +++ b/src/Data/Trie/Convenience.hs @@ -1,13 +1,13 @@ {-# OPTIONS_GHC -Wall -fwarn-tabs #-} ---------------------------------------------------------------- --- ~ 2009.01.20 +-- ~ 2011.02.12 -- | -- Module : Data.Trie.Convenience --- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- Copyright : Copyright (c) 2008--2011 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org --- Stability : provisional +-- Stability : experimental -- Portability : portable -- -- Additional convenience functions. In order to keep "Data.Trie" @@ -19,28 +19,35 @@ module Data.Trie.Convenience ( - -- * Conversion functions + -- * Conversion functions ('fromList' variants) -- $fromList - fromListL, fromListR, fromListS, fromListWith + fromListL, fromListR, fromListS + , fromListWith, fromListWith' + , fromListWithL, fromListWithL' - -- * 'lookupBy' variants + -- * Query functions ('lookupBy' variants) , lookupWithDefault - -- * 'alterBy' variants - , insertIfAbsent, insertWith, insertWithKey + -- * Inserting values ('alterBy' variants) + , insertIfAbsent + , insertWith, insertWith' + , insertWithKey, insertWithKey' + + -- * Updating and adjusting values ('alterBy' and 'adjustBy' variants) , adjustWithKey , update, updateWithKey - -- * 'mergeBy' variants - , disunion, unionWith + -- * Combining tries ('mergeBy' variants) + , disunion + , unionWith, unionWith' ) where import Data.Trie -import Data.Trie.Internal (lookupBy_) +import Data.Trie.Internal (lookupBy_, adjustBy) +import Data.Trie.Errors (impossible) import Data.ByteString (ByteString) import Data.List (foldl', sortBy) import Data.Ord (comparing) -import Control.Monad (liftM) ---------------------------------------------------------------- ---------------------------------------------------------------- @@ -54,20 +61,27 @@ -- which is swapped when reversing the list or changing which -- function is used). + +-- | A left-fold version of 'fromList'. If you run into issues with +-- stack overflows when using 'fromList' or 'fromListR', then you +-- should use this function instead. fromListL :: [(ByteString,a)] -> Trie a {-# INLINE fromListL #-} -fromListL = foldl' (flip $ uncurry $ insertIfAbsent) empty +fromListL = foldl' (flip . uncurry $ insertIfAbsent) empty --- | This version is just an alias for 'fromList'. It is a good --- producer for list fusion. Worst-case behavior is somewhat worse --- than worst-case for 'fromListL'. + +-- | An explicitly right-fold variant of 'fromList'. It is a good +-- consumer for list fusion. Worst-case behavior is somewhat worse +-- than worst-case for 'fromListL'. The 'fromList' function is +-- currently just an alias for 'fromListR'. fromListR :: [(ByteString,a)] -> Trie a {-# INLINE fromListR #-} -fromListR = fromList +fromListR = fromList -- == foldr (uncurry insert) empty --- TODO: compare performance against a fromListL definition, adjusting the sort + +-- TODO: compare performance against a fromListL variant, adjusting the sort appropriately -- --- | This version sorts the list before folding over it. This adds +-- | This variant sorts the list before folding over it. This adds -- /O(n log n)/ overhead and requires the whole list be in memory -- at once, but it ensures that the list is in best-case order. The -- benefits generally outweigh the costs. @@ -75,7 +89,13 @@ {-# INLINE fromListS #-} fromListS = fromListR . sortBy (comparing fst) --- | A variant of 'fromListR' that takes a function for combining values on conflict. + +-- | A variant of 'fromListR' that takes a function for combining +-- values on conflict. The first argument to the combining function +-- is the ``new'' value from the initial portion of the list; the +-- second argument is the value that has been accumulated into the +-- trie from the tail of the list (just like the first argument to +-- 'foldr'). Thus, @fromList = fromListWith const@. fromListWith :: (a -> a -> a) -> [(ByteString,a)] -> Trie a {-# INLINE fromListWith #-} fromListWith f = foldr (uncurry $ alterBy g) empty @@ -83,46 +103,116 @@ g _ v Nothing = Just v g _ v (Just w) = Just (f v w) + +-- | A variant of 'fromListWith' which applies the combining +-- function strictly. This function is a good consumer for list +-- fusion. If you need list fusion and are running into stack +-- overflow problems with 'fromListWith', then this function may +-- solve the problem. +fromListWith' :: (a -> a -> a) -> [(ByteString,a)] -> Trie a +{-# INLINE fromListWith' #-} +fromListWith' f = foldr (uncurry $ alterBy g') empty + where + g' _ v Nothing = Just v + g' _ v (Just w) = Just $! f v w + + +-- | A left-fold variant of 'fromListWith'. Note that the arguments +-- to the combining function are swapped: the first is the value +-- in the trie which has been accumulated from the initial part of +-- the list; the second argument is the ``new'' value from the +-- remaining tail of the list (just like the first argument to +-- 'foldl'). Thus, @fromListL = fromListWithL const@. +fromListWithL :: (a -> a -> a) -> [(ByteString,a)] -> Trie a +{-# INLINE fromListWithL #-} +fromListWithL f = foldl' (flip . uncurry $ alterBy flipG) empty + where + flipG _ v Nothing = Just v + flipG _ v (Just w) = Just (f w v) + + +-- | A variant of 'fromListWithL' which applies the combining +-- function strictly. +fromListWithL' :: (a -> a -> a) -> [(ByteString,a)] -> Trie a +{-# INLINE fromListWithL' #-} +fromListWithL' f = foldl' (flip . uncurry $ alterBy flipG') empty + where + flipG' _ v Nothing = Just v + flipG' _ v (Just w) = Just $! f w v + ---------------------------------------------------------------- -- | Lookup a key, returning a default value if it's not found. lookupWithDefault :: a -> ByteString -> Trie a -> a -lookupWithDefault x = lookupBy_ (\mv _ -> case mv of - Nothing -> x - Just v -> v) x (const x) +lookupWithDefault def = lookupBy_ f def (const def) + where + f Nothing _ = def + f (Just v) _ = v ---------------------------------------------------------------- -- | Insert a new key, retaining old value on conflict. insertIfAbsent :: ByteString -> a -> Trie a -> Trie a -insertIfAbsent = alterBy $ \_ x mv -> case mv of - Nothing -> Just x - Just _ -> mv +insertIfAbsent = + alterBy $ \_ x mv -> + case mv of + Nothing -> Just x + Just _ -> mv -- | Insert a new key, with a function to resolve conflicts. insertWith :: (a -> a -> a) -> ByteString -> a -> Trie a -> Trie a -insertWith f = alterBy $ \_ x mv -> case mv of - Nothing -> Just x - Just v -> Just (f x v) +insertWith f = + alterBy $ \_ x mv -> + case mv of + Nothing -> Just x + Just v -> Just (f x v) +-- | A variant of 'insertWith' which applies the combining function +-- strictly. +insertWith' :: (a -> a -> a) -> ByteString -> a -> Trie a -> Trie a +insertWith' f = + alterBy $ \_ x mv -> + case mv of + Nothing -> Just x + Just v -> Just $! f x v + +-- | A variant of 'insertWith' which also provides the key to the +-- combining function. insertWithKey :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> Trie a -insertWithKey f = alterBy $ \k x mv -> case mv of - Nothing -> Just x - Just v -> Just (f k x v) +insertWithKey f = + alterBy $ \k x mv -> + case mv of + Nothing -> Just x + Just v -> Just (f k x v) +-- | A variant of 'insertWithKey' which applies the combining +-- function strictly. +insertWithKey' :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> Trie a +insertWithKey' f = + alterBy $ \k x mv -> + case mv of + Nothing -> Just x + Just v -> Just $! f k x v + {- This is a tricky one... insertLookupWithKey :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> (Maybe a, Trie a) -} +---------------------------------------------------------------- -- | Apply a function to change the value at a key. -adjustWithKey :: (ByteString -> a -> a) -> ByteString -> Trie a -> Trie a -adjustWithKey f q = alterBy (\k _ -> liftM (f k)) q undefined +adjustWithKey :: (ByteString -> a -> a) -> ByteString -> Trie a -> Trie a +adjustWithKey f q = + adjustBy (\k _ -> f k) q (impossible "Convenience.adjustWithKey") +-- TODO: benchmark vs the definition with alterBy/liftM -- | Apply a function to the value at a key, possibly removing it. update :: (a -> Maybe a) -> ByteString -> Trie a -> Trie a -update f q = alterBy (\_ _ mx -> mx >>= f) q undefined +update f q = + alterBy (\_ _ mx -> mx >>= f) q (impossible "Convenience.update") +-- | A variant of 'update' which also provides the key to the function. updateWithKey :: (ByteString -> a -> Maybe a) -> ByteString -> Trie a -> Trie a -updateWithKey f q = alterBy (\k _ mx -> mx >>= f k) q undefined +updateWithKey f q = + alterBy (\k _ mx -> mx >>= f k) q (impossible "Convenience.updateWithKey") {- updateLookupWithKey :: (ByteString -> a -> Maybe a) -> ByteString -> Trie a -> (Maybe a, Trie a) @@ -131,13 +221,20 @@ ---------------------------------------------------------------- --- | Combine two tries. If they define the same key, it is removed. +-- | Combine two tries, a la symmetric difference. If they define +-- the same key, it is removed; otherwise it is retained with the +-- value it has in whichever trie. disunion :: Trie a -> Trie a -> Trie a disunion = mergeBy (\_ _ -> Nothing) -- | Combine two tries, using a function to resolve conflicts. unionWith :: (a -> a -> a) -> Trie a -> Trie a -> Trie a unionWith f = mergeBy (\x y -> Just (f x y)) + +-- | A variant of 'unionWith' which applies the combining function +-- strictly. +unionWith' :: (a -> a -> a) -> Trie a -> Trie a -> Trie a +unionWith' f = mergeBy (\x y -> Just $! f x y) {- TODO: (efficiently) difference, intersection diff --git a/src/Data/Trie/Errors.hs b/src/Data/Trie/Errors.hs new file mode 100644 --- /dev/null +++ b/src/Data/Trie/Errors.hs @@ -0,0 +1,30 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} + +---------------------------------------------------------------- +-- ~ 2011.02.12 +-- | +-- Module : Data.Trie.Errors +-- Copyright : Copyright (c) 2008--2011 wren ng thornton +-- License : BSD3 +-- Maintainer : wren@community.haskell.org +-- Stability : experimental +-- Portability : portable +-- +-- Internal convenience functions for giving error messages. +---------------------------------------------------------------- + +module Data.Trie.Errors + ( impossible + ) where + +---------------------------------------------------------------- +---------------------------------------------------------------- + +-- | The impossible happened. Use this instead of 'undefined' just in case. +impossible :: String -> a +{-# NOINLINE impossible #-} +impossible fn = + error $ "Data.Trie." ++ fn ++ ": the impossible happened. This is a bug, please report it to the maintainer." + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/src/Data/Trie/Internal.hs b/src/Data/Trie/Internal.hs --- a/src/Data/Trie/Internal.hs +++ b/src/Data/Trie/Internal.hs @@ -1,13 +1,15 @@ -{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +-- To make GHC stop warning about the Prelude +{-# OPTIONS_GHC -Wall -fwarn-tabs -fno-warn-unused-imports #-} +{-# LANGUAGE NoImplicitPrelude #-} -- For list fusion on toListBy, and for applicative hiding {-# LANGUAGE CPP #-} ---------------------------------------------------------------- --- ~ 2009.01.20 +-- ~ 2010.08.15 -- | -- Module : Data.Trie.Internal --- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- Copyright : Copyright (c) 2008--2011 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org -- Stability : provisional @@ -38,21 +40,26 @@ , lookupBy_, submap -- * Single-value modification - , alterBy, adjustBy + , alterBy, alterBy_, adjustBy -- * Combining tries , mergeBy -- * Mapping functions - , mapBy, filterMap + , mapBy + , filterMap + , contextualMap + , contextualMap' + , contextualFilterMap + , contextualMapBy -- * Priority-queue functions , minAssoc, maxAssoc , updateMinViewBy, updateMaxViewBy ) where -import Prelude hiding (null, lookup) -import qualified Prelude +import Prelude hiding (null, lookup) +import qualified Prelude (null, lookup) import qualified Data.ByteString as S import Data.Trie.ByteStringInternal @@ -204,16 +211,9 @@ -- TODO: consider an instance more like the new one for Data.Map. Better? 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 + put Empty = do 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 @@ -227,10 +227,12 @@ ---------------------------------------------------------------} instance Functor Trie where - fmap _ Empty = Empty - fmap f (Arc k Nothing t) = Arc k Nothing (fmap f t) - fmap f (Arc k (Just v) t) = Arc k (Just (f v)) (fmap f t) - fmap f (Branch p m l r) = Branch p m (fmap f l) (fmap f r) + fmap f = go + where + go Empty = Empty + go (Arc k Nothing t) = Arc k Nothing (go t) + go (Arc k (Just v) t) = Arc k (Just (f v)) (go t) + go (Branch p m l r) = Branch p m (go l) (go r) #ifdef APPLICATIVE_IN_BASE @@ -238,10 +240,12 @@ -- If our definition of foldr is so much faster than the Endo -- default, then maybe we should remove this and use the default -- foldMap based on foldr - foldMap _ Empty = mempty - foldMap f (Arc _ Nothing t) = foldMap f t - foldMap f (Arc _ (Just v) t) = f v `mappend` foldMap f t - foldMap f (Branch _ _ l r) = foldMap f l `mappend` foldMap f r + foldMap f = go + where + go Empty = mempty + go (Arc _ Nothing t) = go t + go (Arc _ (Just v) t) = f v `mappend` go t + go (Branch _ _ l r) = go l `mappend` go r {- This definition is much faster, but it's also wrong -- (or at least different than foldrWithKey) @@ -264,10 +268,12 @@ -- TODO: newtype Assoc = A Trie ; instance Foldable Assoc instance Traversable Trie where - traverse _ Empty = pure Empty - traverse f (Arc k Nothing t) = Arc k Nothing <$> traverse f t - traverse f (Arc k (Just v) t) = Arc k . Just <$> f v <*> traverse f t - traverse f (Branch p m l r) = Branch p m <$> traverse f l <*> traverse f r + traverse f = go + where + go Empty = pure Empty + go (Arc k Nothing t) = Arc k Nothing <$> go t + go (Arc k (Just v) t) = Arc k . Just <$> f v <*> go t + go (Branch p m l r) = Branch p m <$> go l <*> go r instance Applicative Trie where pure = return @@ -331,11 +337,14 @@ -- | 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) +filterMap f = go + where + go Empty = empty + go (Arc k Nothing t) = arc k Nothing (go t) + go (Arc k (Just v) t) = arc k (f v) (go t) + go (Branch p m l r) = branch p m (go l) (go r) + -- | Generic version of 'fmap'. This function is notably more -- expensive than 'fmap' or 'filterMap' because we have to reconstruct -- the keys. @@ -348,12 +357,56 @@ go q (Branch p m l r) = branch p m (go q l) (go q r) +-- | A variant of 'fmap' which provides access to the subtrie rooted +-- at each value. +contextualMap :: (a -> Trie a -> b) -> Trie a -> Trie b +contextualMap f = go + where + go Empty = Empty + go (Arc k Nothing t) = Arc k Nothing (go t) + go (Arc k (Just v) t) = Arc k (Just (f v t)) (go t) + go (Branch p m l r) = Branch p m (go l) (go r) + + +-- | A variant of 'contextualMap' which applies the function strictly. +contextualMap' :: (a -> Trie a -> b) -> Trie a -> Trie b +contextualMap' f = go + where + go Empty = Empty + go (Arc k Nothing t) = Arc k Nothing (go t) + go (Arc k (Just v) t) = Arc k (Just $! f v t) (go t) + go (Branch p m l r) = Branch p m (go l) (go r) + + +-- | A contextual variant of 'filterMap'. +contextualFilterMap :: (a -> Trie a -> Maybe b) -> Trie a -> Trie b +contextualFilterMap f = go + where + go Empty = empty + go (Arc k Nothing t) = arc k Nothing (go t) + go (Arc k (Just v) t) = arc k (f v t) (go t) + go (Branch p m l r) = branch p m (go l) (go r) + + +-- | A contextual variant of 'mapBy'. Again note that this is +-- expensive since we must reconstruct the keys. +contextualMapBy :: (ByteString -> a -> Trie a -> Maybe b) -> Trie a -> Trie b +contextualMapBy f = go S.empty + where + go _ Empty = empty + go q (Arc k Nothing t) = arc k Nothing (go (S.append q k) t) + go q (Arc k (Just v) t) = let q' = S.append q k + in arc k (f q' v t) (go q' t) + go q (Branch p m l r) = branch p m (go q l) (go q r) + + {--------------------------------------------------------------- -- Smart constructors and helper functions for building tries ---------------------------------------------------------------} -- | Smart constructor to prune @Empty@ from @Branch@es. branch :: Prefix -> Mask -> Trie a -> Trie a -> Trie a +{-# INLINE branch #-} branch _ _ Empty r = r branch _ _ l Empty = l branch p m l r = Branch p m l r @@ -362,6 +415,7 @@ -- | Smart constructor to prune @Arc@s that lead nowhere. -- N.B if mv=Just then doesn't check whether t=epsilon. It's up to callers to ensure that invariant isn't broken. arc :: ByteString -> Maybe a -> Trie a -> Trie a +{-# INLINE arc #-} arc k mv@(Just _) t = Arc k mv t arc _ Nothing Empty = Empty arc k Nothing t@(Branch _ _ _ _) | S.null k = t @@ -375,6 +429,7 @@ -- -- N.B. /do not/ use if prefixes could match entirely! branchMerge :: Prefix -> Trie a -> Prefix -> Trie a -> Trie a +{-# INLINE branchMerge #-} branchMerge _ Empty _ t2 = t2 branchMerge _ t1 _ Empty = t1 branchMerge p1 t1 p2 t2 @@ -390,6 +445,7 @@ -- we can see 4/8/?*Word8 at a time instead of just one. -- But that makes maintaining invariants ...difficult :( getPrefix :: Trie a -> Prefix +{-# INLINE getPrefix #-} getPrefix (Branch p _ _ _) = p getPrefix (Arc k _ _) | S.null k = 0 -- for lack of a better value | otherwise = S.head k @@ -400,10 +456,12 @@ -- Error messages ---------------------------------------------------------------} +-- TODO: shouldn't we inline the logic and just NOINLINE the string constant? There are only three use sites, which themselves aren't inlined... errorLogHead :: String -> ByteString -> ByteStringElem {-# NOINLINE errorLogHead #-} -errorLogHead s q | S.null q = error (s ++": found null subquery") - | otherwise = S.head q +errorLogHead fn q + | S.null q = error $ "Data.Trie.Internal." ++ fn ++": found null subquery" + | otherwise = S.head q ---------------------------------------------------------------- @@ -418,18 +476,21 @@ {-# INLINE empty #-} empty = Empty + -- | /O(1)/, Is the trie empty? null :: Trie a -> Bool {-# INLINE null #-} null Empty = True null _ = False + -- | /O(1)/, Construct a singleton trie. singleton :: ByteString -> a -> Trie a {-# INLINE singleton #-} singleton k v = Arc k (Just v) Empty -- For singletons, don't need to verify invariant on arc length >0 + -- | /O(n)/, Get count of elements in trie. size :: Trie a -> Int {-# INLINE size #-} @@ -471,29 +532,27 @@ where go _ Empty = id go q (Branch _ _ l r) = go q l . go q r - go q (Arc k mv t) = case mv of - Nothing -> rest - Just v -> (fcons k' v) . rest - where - rest = go k' t - k' = S.append q k + go q (Arc k mv t) = + case mv of + Nothing -> rest + Just v -> fcons k' v . rest + where + rest = go k' t + k' = S.append q k + -- cf Data.ByteString.unpack -- -- -- | Convert a trie into a list using a function. Resulting values -- are in key-sorted order. toListBy :: (ByteString -> a -> b) -> Trie a -> [b] - +{-# INLINE toListBy #-} #if !defined(__GLASGOW_HASKELL__) -- TODO: should probably inline foldrWithKey -- TODO: compare performance of that vs both this and the GHC version -{-# INLINE toListBy #-} toListBy f t = foldrWithKey (((:) .) . f) [] t - #else - -{-# INLINE toListBy #-} -- Written with 'build' to enable the build\/foldr fusion rules. toListBy f t = build (toListByFB f t) @@ -504,7 +563,6 @@ toListByFB :: (ByteString -> a -> b) -> Trie a -> (b -> c -> c) -> c -> c {-# INLINE [0] toListByFB #-} toListByFB f t cons nil = foldrWithKey ((cons .) . f) nil t - #endif @@ -602,17 +660,25 @@ -- to resolve conflicts (or non-conflicts). alterBy :: (ByteString -> a -> Maybe a -> Maybe a) -> ByteString -> a -> Trie a -> Trie a -alterBy f_ q_ x_ +alterBy f = alterBy_ (\k v mv t -> (f k v mv, t)) +-- TODO: use GHC's 'inline' function so that this gets specialized away. +-- TODO: benchmark to be sure that this doesn't introduce unforseen performance costs because of the uncurrying etc. + + +-- | A variant of 'alterBy' which also allows modifying the sub-trie. +alterBy_ :: (ByteString -> a -> Maybe a -> Trie a -> (Maybe a, Trie a)) + -> ByteString -> a -> Trie a -> Trie a +alterBy_ f_ q_ x_ | S.null q_ = alterEpsilon | otherwise = go q_ where f = f_ q_ x_ - nothing q = arc q (f Nothing) Empty + nothing q = uncurry (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_ + alterEpsilon t_@Empty = uncurry (arc q_) (f Nothing t_) + alterEpsilon t_@(Branch _ _ _ _) = uncurry (arc q_) (f Nothing t_) + alterEpsilon t_@(Arc k mv t) | S.null k = uncurry (arc q_) (f mv t) + | otherwise = uncurry (arc q_) (f Nothing t_) go q Empty = nothing q @@ -625,38 +691,41 @@ qh = errorLogHead "alterBy" q go q t_@(Arc k mv t) = - let (p,k',q') = breakMaximalPrefix 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) + let (p,k',q') = breakMaximalPrefix k q in + case (not $ S.null k', S.null q') of + (True, True) -> -- add node to middle of arc + uncurry (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) -> uncurry (arc k) (f mv t) + (False, False) -> arc k mv (go q' t) --- | ... +-- | Alter the value associated with a given key. If the key is not +-- present, then the trie is returned unaltered. See 'alterBy' if +-- you are interested in inserting new keys or deleting old keys. +-- Because this function does not need to worry about changing the +-- trie structure, it is somewhat faster than 'alterBy'. adjustBy :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> Trie a adjustBy f_ q_ x_ - | S.null q_ = \t_ -> case t_ of - (Arc k (Just v) t) - | S.null k -> Arc k (Just (f v)) t - _ -> t_ + | S.null q_ = adjustEpsilon | otherwise = go q_ where f = f_ q_ x_ + adjustEpsilon (Arc k (Just v) t) | S.null k = Arc k (Just (f v)) t + adjustEpsilon t_ = t_ + go _ Empty = Empty go q t@(Branch p m l r) @@ -667,12 +736,12 @@ qh = errorLogHead "adjustBy" q go q t_@(Arc k mv t) = - let (_,k',q') = breakMaximalPrefix k q - in case (not $ S.null k', S.null q') of - (True, True) -> t_ -- don't break arc inline - (True, False) -> t_ -- don't break arc branching - (False, True) -> Arc k (liftM f mv) t - (False, False) -> Arc k mv (go q' t) + let (_,k',q') = breakMaximalPrefix k q in + case (not $ S.null k', S.null q') of + (True, True) -> t_ -- don't break arc inline + (True, False) -> t_ -- don't break arc branching + (False, True) -> Arc k (liftM f mv) t + (False, False) -> Arc k mv (go q' t) {--------------------------------------------------------------- @@ -729,6 +798,8 @@ | zero p0 m1 = branch p1 m1 (go t0 l1) r1 | otherwise = branch p1 m1 l1 (go t0 r1) + -- We combine these branches of 'go' in order to clarify where the definitions of 'p0', 'p1', 'm'', 'p'' are relevant. However, this may introduce inefficiency in the pattern matching automaton... + -- TODO: check. And get rid of 'go'' if it does. go t0_ t1_ = go' t0_ t1_ where p0 = getPrefix t0_ @@ -739,8 +810,8 @@ go' (Arc k0 mv0 t0) (Arc k1 mv1 t1) | m' == 0 = - let (pre,k0',k1') = breakMaximalPrefix k0 k1 - in if S.null pre + let (pre,k0',k1') = breakMaximalPrefix 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') @@ -765,6 +836,7 @@ go' _ _ | zero p0 m' = Branch p' m' t0_ t1_ go' _ _ = Branch p' m' t1_ t0_ + mergeMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a {-# INLINE mergeMaybe #-} mergeMaybe _ Nothing Nothing = Nothing @@ -806,7 +878,7 @@ updateMinViewBy f = go S.empty where go _ Empty = Nothing - go q (Arc k (Just v) t) = let q' = (S.append q k) + go q (Arc k (Just v) t) = let q' = S.append q k in Just (q',v, arc k (f q' v) t) go q (Arc k Nothing t) = mapView (arc k Nothing) (go (S.append q k) t) go q (Branch p m l r) = mapView (\l' -> branch p m l' r) (go q l) @@ -817,7 +889,7 @@ updateMaxViewBy f = go S.empty where go _ Empty = Nothing - go q (Arc k (Just v) Empty) = let q' = (S.append q k) + go q (Arc k (Just v) Empty) = let q' = S.append q k in Just (q',v, arc k (f q' v) Empty) go q (Arc k mv t) = mapView (arc k mv) (go (S.append q k) t) go q (Branch p m l r) = mapView (branch p m l) (go q r)