diff --git a/Data/Trie.hs b/Data/Trie.hs new file mode 100644 --- /dev/null +++ b/Data/Trie.hs @@ -0,0 +1,586 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} + +---------------------------------------------------------------- +-- ~ 2008.12.19 +-- | +-- Module : Data.Trie +-- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- License : BSD3 +-- Maintainer : wren@community.haskell.org +-- Stability : beta +-- Portability : portable +-- +-- An efficient implementation of maps from strings to values. +-- +-- The implementation is based on /big-endian patricia trees/, like +-- "Data.IntMap". We first trie on the elements of "Data.ByteString" +-- and then trie on the big-endian bit representation of those +-- 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, +-- +-- +-- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve +-- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4), +-- October 1968, pages 514-534. +---------------------------------------------------------------- + +module Data.Trie + ( + -- * Data types + Trie(), KeyString, KeyElem, showTrie + + -- * Basic functions + , empty, null, singleton, size, toList, fromList + + -- * Query functions + , lookupBy_, lookupBy, lookup, member, submap + + -- * Single-value modification + , alterBy, insert, adjust, delete + + -- * Combining tries + , mergeBy, unionL, unionR + + -- * Mapping functions + , mapBy, filterMap + ) where + +import Prelude hiding (null, lookup) +import qualified Prelude + +import qualified Data.ByteString as S +import Data.Trie.ByteStringInternal +import Data.Trie.BitTwiddle + +import Data.Maybe (isJust) +import Control.Monad (liftM) +import Control.Applicative (Applicative(..), (<$>)) +import Data.Monoid (Monoid(..)) +import Data.Foldable (Foldable(foldMap)) +import Data.Traversable (Traversable(traverse)) +---------------------------------------------------------------- +---------------------------------------------------------------- + + +{--------------------------------------------------------------- +-- ByteString Big-endian Patricia Trie +---------------------------------------------------------------} +type KeyString = ByteString +type KeyElem = ByteStringElem + +{- Idealized: +data Node a = Accept a (ArcSet a) + | Reject (Branch a) -- Invariant: Must be Branch +data Arc a = Arc KeyString (Node a) -- Invariant: never empty string +data ArcSet a = None + | One {KeyElem} (Arc a) + | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) +data Trie a = Empty + | Start KeyString (Node a) -- Maybe empty string [1] + +[1] If we maintain the invariants on how Nodes recurse, then we +can't simply have Start(Node a) because we may have a shared prefix +where the prefix itself is not Accept'ed. + + +-- Squash Arc into One: +-- (pure good) +data Node a = Accept a (ArcSet a) + | Reject (Branch a) +data ArcSet a = None + | Arc KeyString (Node a) + | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) +data Trie a = Empty + | Start KeyString (Node a) + + +-- Squash Node together: +-- (most likely good) +data Node a = Node (Maybe a) (ArcSet a) +data ArcSet a = None + | Arc KeyString (Node a) + | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) +data Trie a = Empty + | Start KeyString (Node a) + + +-- Squash Empty/None and Arc/Start together: +-- (Complicates invariants about non-empty strings and Node's recursion) +data Node a = Node (Maybe a) (ArcSet a) +data Trie a = Empty + | Arc KeyString (Node a) + | Branch {Prefix} {Mask} (Trie a) (Trie a) + + +-- Squash Node into Arc: +-- (By this point, pure good) +-- Unseen invariants: +-- * KeyString non-empty, unless Arc is absolute root of tree +-- * If (Maybe a) is Nothing, then (Trie a) is Branch +-- * With views, we could re-expand Arc into accepting and +-- nonaccepting variants +-- +-- [2] Maybe we shouldn't unpack the KeyString. We could specialize +-- or inline the splitMaximalPrefix function to prevent constructing +-- a new KeyString from the parts... +-} + +-- | A map from 'ByteString's to @a@. For all the generic functions, +-- note that tries are strict in the @Maybe@ but not in @a@. +-- +-- The 'Monad' instance is strange. If a key @k1@ is a prefix of +-- other keys, then results from binding the value at @k1@ will +-- override values from longer keys when they collide. If this is +-- useful for anything, or if there's a more sensible instance, I'd +-- be curious to know. + +data Trie a = Empty + | Arc {-# UNPACK #-} !KeyString + !(Maybe a) + !(Trie a) + | Branch {-# UNPACK #-} !Prefix + {-# UNPACK #-} !Mask + !(Trie a) + !(Trie a) + +-- | Visualization fuction for debugging. +showTrie :: (Show a) => Trie a -> String +showTrie t = shows' id t "" + where + spaces f = map (const ' ') (f "") + + shows' _ Empty = (".\n"++) + shows' ss (Branch p m l r) = + let s' = ("--"++) . shows p . (","++) . shows m . ("-+"++) + ss' = ss . (tail (spaces s') ++) + in s' . shows' (ss' . ("|"++)) l + . ss' . ("|\n"++) + . ss' . ("`"++) . shows' (ss' . (" "++)) r + shows' ss (Arc k mv t') = + let s' = ("--"++) . shows k + . maybe id (\v -> ("-("++) . shows v . (")"++)) mv + . ("--"++) + in s' . shows' (ss . (spaces s' ++)) t' + + +{--------------------------------------------------------------- +-- Trie instances +---------------------------------------------------------------} + +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) + +-- Does this even make sense? It's not nondeterminism like lists +-- and sets. If no keys were prefixes of other keys it'd make sense +-- as a decision-tree; but since keys /can/ prefix, tries formed +-- from shorter keys can shadow the results from longer keys due +-- to the 'unionL'. It does seem to follow the laws though... What +-- computation could this possibly represent? +-- +-- 1. return x >>= f == f x +-- 2. m >>= return == m +-- 3. (m >>= f) >>= g == m >>= (\x -> f x >>= g) +instance Monad Trie where + return x = singleton S.empty x + + (>>=) Empty _ = empty + (>>=) (Arc k Nothing t) f = arc k Nothing (t >>= f) + (>>=) (Arc k (Just v) t) f = arc k Nothing (f v `unionL` (t >>= f)) + (>>=) (Branch p m l r) f = branch p m (l >>= f) (r >>= f) + + +instance Monoid a => Monoid (Trie a) where + mempty = empty + mappend = mergeBy $ \x y -> Just (x `mappend` y) + + +instance Foldable Trie where + 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 + + +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 + +{--------------------------------------------------------------- +-- 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 +branch _ _ Empty r = r +branch _ _ l Empty = l +branch p m l r = Branch p m l r + + +-- | Smart constructor to prune @Arc@s that lead nowhere. +arc :: KeyString -> Maybe a -> Trie a -> Trie a +arc k mv@(Just _) t = Arc k mv t +arc _ Nothing Empty = Empty +arc k Nothing t@(Branch _ _ _ _) | S.null k = t + | otherwise = Arc k Nothing t +arc k Nothing (Arc k' mv' t') = Arc (S.append k k') mv' t' + + +-- | Smart constructor to join two tries into a @Branch@ with maximal +-- prefix sharing. Requires knowing the prefixes, but can combine +-- either @Branch@es or @Arc@s. +-- +-- N.B. /do not/ use if prefixes could match entirely! +branchMerge :: Prefix -> Trie a -> Prefix -> Trie a -> Trie a +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 + where + m = branchMask p1 p2 + p = mask p1 m + + +-- It would be better if Arc used +-- Data.ByteString.TrieInternal.wordHead somehow, that way +-- we can see 4/8/?*Word8 at a time instead of just one. +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" + + +---------------------------------------------------------------- +---------------------------------------------------------------- + + +{--------------------------------------------------------------- +-- Error messages +---------------------------------------------------------------} + +-- | Once correctness is proven, these error messages could be +-- preprocessed away in order to give minor optimizations +errorInvariantBroken :: String -> String -> a +errorInvariantBroken s e = error (s ++ ": Invariant was broken" ++ e') + where + e' = if Prelude.null e then e else ", found: " ++ e + +errorArcAfterNothing :: String -> a +errorArcAfterNothing s = errorInvariantBroken s "Arc after Nothing" + +errorEmptyAfterNothing :: String -> a +errorEmptyAfterNothing s = errorInvariantBroken s "Empty after Nothing" + + +errorLogHead :: String -> KeyString -> KeyElem +errorLogHead s q | S.null q = error (s ++": found null subquery") + | otherwise = S.head q + +{--------------------------------------------------------------- +-- Basic functions +---------------------------------------------------------------} + +-- | /O(1)/, The empty trie. +{-# INLINE empty #-} +empty :: Trie a +empty = Empty + +-- | /O(1)/, Is the trie empty? +{-# INLINE null #-} +null :: Trie a -> Bool +null Empty = True +null _ = False + +-- | /O(1)/, A singleton trie. +{-# INLINE singleton #-} +singleton :: KeyString -> a -> Trie a +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. +{-# INLINE size #-} +size :: Trie a -> Int +size t = size' t id 0 + +-- | /O(n)/, Internal CPS accumulator function for calculating +-- 'size'. +size' :: Trie a -> (Int -> Int) -> Int -> Int +size' Empty f n = f n +size' (Branch _ _ l r) f n = size' l (size' r f) n +size' (Arc _ Nothing t) f n = size' t f n +size' (Arc _ (Just _) t) f n = size' t f $! n + 1 + +-- BUG: Inefficient for now +-- TODO: rewrite both of these to support list fusion +-- +-- | Convert trie into association list. Keys will be in sorted order. +toList :: Trie a -> [(KeyString,a)] +toList Empty = [] +toList (Branch _ _ l r) = toList l ++ toList r +toList (Arc k mv t) = maybe [] (\v -> [(k,v)]) mv + ++ map (\(q,x) -> (S.append k q, x)) (toList t) + +-- | Convert association list into a trie. On key conflict, values +-- earlier in the list shadow later ones. +fromList :: [(KeyString,a)] -> Trie a +fromList = foldr (uncurry insert) empty + + +{--------------------------------------------------------------- +-- Query functions (just recurse) +---------------------------------------------------------------} + +-- | Generic function to find a value (if it exists) and the subtrie +-- rooted at the prefix. The first function argument is called if and +-- only if a node is exactly reachable by the query; if no node is +-- exactly reachable the default value is used; if the middle of +-- an arc is reached, the second function argument is used. +-- +-- This function is intended for internal use. +lookupBy_ :: (Maybe a -> Trie a -> b) -> b -> (Trie a -> b) + -> KeyString -> Trie a -> b +lookupBy_ f z a = go + where + go _ Empty = z + + go q (Arc k mv t) + | not (S.null k') = a (Arc k' mv t) + | S.null q' = f mv t + | otherwise = go q' t + where + (_,k',q') = splitMaximalPrefix k q + + go q t_@(Branch _ _ _ _) = findArc t_ + where + qh = errorLogHead "lookupBy_" q + + -- | /O(min(m,W))/, where /m/ is number of @Arc@s in this + -- branching, and /W/ is the word size of the Prefix,Mask type. + findArc (Branch p m l r) + | nomatch qh p m = z + | zero qh m = findArc l + | otherwise = findArc r + findArc t@(Arc _ _ _) = go q t + findArc Empty = z + + +-- | Generic function to find a value (if it exists) and the subtrie +-- rooted at the prefix. Intended for public consumption. +{-# INLINE lookupBy #-} +lookupBy :: (Maybe a -> Trie a -> b) -> KeyString -> Trie a -> b +lookupBy f = lookupBy_ f (f Nothing Empty) (f Nothing) + +-- | Return the value associated with a query string if it exists. +{-# INLINE lookup #-} +lookup :: KeyString -> Trie a -> Maybe a +lookup = lookupBy_ const Nothing (const Nothing) + +-- | Return the subtrie containing all keys beginning with a prefix. +submap :: KeyString -> Trie a -> Trie a +submap q = lookupBy_ submap' Empty (arc q Nothing) q + where + submap' Nothing Empty = errorEmptyAfterNothing "submap" + submap' Nothing (Arc _ _ _) = errorArcAfterNothing "submap" + submap' mx t = Arc q mx t + +-- | Does a string have a value in the trie? +{-# INLINE member #-} +member :: KeyString -> Trie a -> Bool +member q = isJust . lookup q + + +{--------------------------------------------------------------- +-- Single-value modification functions (recurse and clone spine) +---------------------------------------------------------------} + +-- TODO: We should CPS on Empty to avoid cloning spine if no change. +-- Difficulties arise with the calls to 'branch' and 'arc'. Will +-- have to create a continuation chain, so no savings on memory +-- allocation; but would have savings on held memory, if they're +-- still holding the old one... +-- +-- | Generic function to alter a trie by one element with a function +-- to resolve conflicts (or non-conflicts). +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_) + | otherwise = go q_ + where + f = f_ q_ x_ + nothing q = arc q (f Nothing) Empty + + go q Empty = nothing q + + go q t@(Branch p m l r) + | nomatch qh p m = branchMerge p t qh (nothing q) + | zero qh m = branch p m (go q l) r + | otherwise = branch p m l (go q r) + where + qh = errorLogHead "alterBy" q + + go q t_@(Arc k mv t) + | not (S.null k') = case nothing q' of + Empty -> t_ + l -> let r = Arc k' mv t + in (if S.null p + then id + else Arc p Nothing) + (branchMerge (getPrefix l) l + (getPrefix r) r) + | S.null q' = arc k (f mv) t + | otherwise = arc k mv (go q' t) + where + (p,k',q') = splitMaximalPrefix k q + + +-- | Insert a new key. If the key is already present, overrides the +-- old value +{-# INLINE insert #-} +insert :: KeyString -> a -> Trie a -> Trie a +insert = alterBy (\_ x _ -> Just x) + +-- | Apply a function to the value at a key. +{-# INLINE adjust #-} +adjust :: (a -> a) -> KeyString -> Trie a -> Trie a +adjust f q = alterBy (\_ _ -> liftM f) q undefined + +-- | Remove the value stored at a key. +{-# INLINE delete #-} +delete :: KeyString -> Trie a -> Trie a +delete q = alterBy (\_ _ _ -> Nothing) q undefined + + +{--------------------------------------------------------------- +-- Trie-combining functions +---------------------------------------------------------------} + +-- TEST CASES: foldr (unionL . uncurry singleton) empty t +-- foldr (uncurry insert) empty t +-- 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 + 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) + + 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 -> 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 + (mv',t',t'') = case (S.null k0', S.null k1') of + (False,False) -> ( Nothing + , Arc k0' mv0 t0 + , Arc k1' mv1 t1 + ) + (False,True) -> ( mv1 + , Arc k0' mv0 t0 + , t1 + ) + (True, False) -> ( mv0 + , t0 + , Arc k1' mv1 t1 + ) + (True, True) -> ( case (mv0,mv1) of + (Nothing,Nothing) -> Nothing + (Nothing,Just _) -> mv1 + (Just _, Nothing) -> mv0 + (Just v0,Just v1) -> f v0 v1 + , t0 + , t1 + ) + in arc pk mv' (mergeBy f t' t'') + + (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) + + -- 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' + + +-- | Combine two tries, resolving conflicts by choosing the value +-- from the left trie. +{-# INLINE unionL #-} +unionL :: Trie a -> Trie a -> Trie a +unionL = mergeBy (\x _ -> Just x) + +-- | Combine two tries, resolving conflicts by choosing the value +-- from the right trie. +{-# INLINE unionR #-} +unionR :: Trie a -> Trie a -> Trie a +unionR = mergeBy (\_ y -> Just y) + + +{--------------------------------------------------------------- +-- 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. diff --git a/Data/Trie/BitTwiddle.hs b/Data/Trie/BitTwiddle.hs new file mode 100644 --- /dev/null +++ b/Data/Trie/BitTwiddle.hs @@ -0,0 +1,156 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs -fno-warn-name-shadowing #-} +{-# OPTIONS_GHC -cpp -fglasgow-exts #-} + +---------------------------------------------------------------- +-- ~ 2008.12.19 +-- | +-- Module : Data.Trie.BitTwiddle +-- Copyright : Copyright (c) Daan Leijen 2002 +-- License : BSD3 +-- Maintainer : libraries@haskell.org, wren@community.haskell.org +-- Stability : provisional +-- Portability : portable (with CPP) +-- +-- Functions to treat 'Word' as a bit-vector for big-endian patricia +-- trees. This code is duplicated from "Data.IntMap". The only +-- differences are that some of the conversion functions are +-- specialized to 'Word8' for bytestrings, instead of being specialized +-- to 'Int'. +---------------------------------------------------------------- + +module Data.Trie.BitTwiddle + ( Prefix, Mask + , elemToNat + , zero, nomatch + , mask, shorter, branchMask + ) where + +import Data.Bits +import Data.Trie.ByteStringInternal (ByteStringElem) + +#if __GLASGOW_HASKELL__ >= 503 +import GHC.Exts ( Word(..), Int(..), shiftRL# ) +#elif __GLASGOW_HASKELL__ +import GlaExts ( Word(..), Int(..), shiftRL# ) +#else +import Data.Word (Word) +#endif + +---------------------------------------------------------------- + +-- TODO: Natural word size, is 4*Word8 on my machine. Which means +-- it'll be more efficient to Branch by the first 4 bytes instead +-- of just one... +type KeyElem = ByteStringElem +type Prefix = KeyElem +type Mask = KeyElem + + +elemToNat :: KeyElem -> Word +elemToNat i = fromIntegral i + +natToElem :: Word -> KeyElem +natToElem w = fromIntegral w + +shiftRL :: Word -> Int -> Word +#if __GLASGOW_HASKELL__ +-- GHC: use unboxing to get @shiftRL@ inlined. +shiftRL (W# x) (I# i) = W# (shiftRL# x i) +#else +shiftRL x i = shiftR x i +#endif + + +{--------------------------------------------------------------- +-- Endian independent bit twiddling (Trie endianness, not architecture) +---------------------------------------------------------------} + +-- | Is the value under the mask zero? +zero :: KeyElem -> Mask -> Bool +zero i m = (elemToNat i) .&. (elemToNat m) == 0 + +-- | Does a value /not/ match some prefix, for all the bits preceding +-- a masking bit? (Hence a subtree matching the value doesn't exist.) +nomatch :: KeyElem -> Prefix -> Mask -> Bool +nomatch i p m = mask i m /= p + +mask :: KeyElem -> Mask -> Prefix +mask i m = maskW (elemToNat i) (elemToNat m) + + +{--------------------------------------------------------------- +-- Big endian operations (Trie endianness, not architecture) +---------------------------------------------------------------} + +-- | Get mask by setting all bits higher than the smallest bit in +-- @m@. Then apply that mask to @i@. +maskW :: Word -> Word -> Prefix +maskW i m = natToElem (i .&. (complement (m-1) `xor` m)) + +-- | Determine whether the first mask denotes a shorter prefix than +-- the second. +shorter :: Mask -> Mask -> Bool +shorter m1 m2 = elemToNat m1 > elemToNat m2 + +-- | Determine first differing bit of two prefixes. +branchMask :: Prefix -> Prefix -> Mask +branchMask p1 p2 + = natToElem (highestBitMask (elemToNat p1 `xor` elemToNat p2)) + +{--------------------------------------------------------------- + Finding the highest bit (mask) in a word [x] can be done efficiently + in three ways: + * convert to a floating point value and the mantissa tells us the + [log2(x)] that corresponds with the highest bit position. The + mantissa is retrieved either via the standard C function [frexp] + or by some bit twiddling on IEEE compatible numbers (float). + Note that one needs to use at least [double] precision for an + accurate mantissa of 32 bit numbers. + * use bit twiddling, a logarithmic sequence of bitwise or's and + shifts (bit). + * use processor specific assembler instruction (asm). + + The most portable way would be [bit], but is it efficient enough? + I have measured the cycle counts of the different methods on an + AMD Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC + instruction: + + highestBitMask: method cycles + -------------- + frexp 200 + float 33 + bit 11 + asm 12 + + highestBit: method cycles + -------------- + frexp 195 + float 33 + bit 11 + asm 11 + + Wow, the bit twiddling is on today's RISC like machines even + faster than a single CISC instruction (BSR)! +---------------------------------------------------------------} + +{--------------------------------------------------------------- + [highestBitMask] returns a word where only the highest bit is + set. It is found by first setting all bits in lower positions + than the highest bit and than taking an exclusive or with the + original value. Allthough the function may look expensive, GHC + compiles this into excellent C code that subsequently compiled + into highly efficient machine code. The algorithm is derived from + Jorg Arndt's FXT library. +---------------------------------------------------------------} +highestBitMask :: Word -> Word +highestBitMask x + = case (x .|. shiftRL x 1) of + x -> case (x .|. shiftRL x 2) of + x -> case (x .|. shiftRL x 4) of + x -> case (x .|. shiftRL x 8) of + x -> case (x .|. shiftRL x 16) of + x -> case (x .|. shiftRL x 32) of -- for 64 bit platforms + x -> (x `xor` shiftRL x 1) + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/Data/Trie/ByteStringInternal.hs b/Data/Trie/ByteStringInternal.hs new file mode 100644 --- /dev/null +++ b/Data/Trie/ByteStringInternal.hs @@ -0,0 +1,142 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} + +---------------------------------------------------------------- +-- ~ 2008.12.19 +-- | +-- Module : Data.Trie.ByteStringInternal +-- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- License : BSD3 +-- Maintainer : wren@community.haskell.org +-- Stability : beta +-- Portability : portable +-- +-- Helper functions on 'ByteString's for "Data.Trie". +---------------------------------------------------------------- + + +module Data.Trie.ByteStringInternal + ( ByteString, ByteStringElem + , wordHead + , splitMaximalPrefix + ) where + +import qualified Data.ByteString as S +import Data.ByteString.Internal (ByteString(..), inlinePerformIO) +import Data.Word + +import Control.Monad + +import Foreign.ForeignPtr (ForeignPtr, withForeignPtr) +import Foreign.Ptr (Ptr, plusPtr, castPtr) +import Foreign.Storable (Storable(..)) +import Foreign.Marshal.Alloc (alloca) +import System.IO.Unsafe (unsafePerformIO) +import Data.Bits +---------------------------------------------------------------- + +type ByteStringElem = Word8 -- Associated type of ByteString + + +---------------------------------------------------------------- +-- | Return the first natural 'Word' worth of string, padding by +-- zeros as necessary. Number of elements per word, and position +-- of elements within the word varies by architecture. +wordHead :: ByteString -> Word +wordHead (PS s o l) = inlinePerformIO $ + withForeignPtr s $ \p -> + liftM (maskBytes l .&.) + (peek (p `plusPtr` o :: Ptr Word)) + +-- The 0x are automatically truncated when too large, don't need +-- to worry about 'max'ing with sizeOf Word. +maskBytes :: Int -> Word +maskBytes i + | isLittleEndian = case 0 `max` i of + 0 -> 0x0000000000000000 + 1 -> 0x00000000000000FF + 2 -> 0x000000000000FFFF + 3 -> 0x0000000000FFFFFF + 4 -> 0x00000000FFFFFFFF + 5 -> 0x000000FFFFFFFFFF + 6 -> 0x0000FFFFFFFFFFFF + 7 -> 0x00FFFFFFFFFFFFFF + _ -> 0xFFFFFFFFFFFFFFFF + | otherwise = case 0 `max` i of + 0 -> 0x0000000000000000 + 1 -> 0xFF00000000000000 + 2 -> 0xFFFF000000000000 + 3 -> 0xFFFFFF0000000000 + 4 -> 0xFFFFFFFF00000000 + 5 -> 0xFFFFFFFFFF000000 + 6 -> 0xFFFFFFFFFFFF0000 + 7 -> 0xFFFFFFFFFFFFFF00 + _ -> 0xFFFFFFFFFFFFFFFF + +-- TODO: How to get this to execute statically?... +-- BUG? is 'alloca' safe in 'inlinePerformIO' (used in 'wordHead')? +{-# NOINLINE isLittleEndian #-} +isLittleEndian :: Bool +isLittleEndian = unsafePerformIO $ alloca $ \p -> do + poke p (0x04030201 :: Word32) + b <- peek (castPtr p :: Ptr Word8) + case b of + 0x01 -> return True + 0x04 -> return False + _ -> error ("non-standard endianness detected! " + ++ "Contact the Data.Trie maintainer.") + +---------------------------------------------------------------- +-- | Returns the longest shared prefix and the two remaining suffixes +-- for a pair of strings. +splitMaximalPrefix :: ByteString -> ByteString + -> (ByteString, ByteString, ByteString) +splitMaximalPrefix + str1@(PS s1 off1 len1) + str2@(PS s2 off2 len2) + | len1 == 0 = (S.empty, S.empty, str2) + | len2 == 0 = (S.empty, str1, S.empty) + | otherwise = inlinePerformIO $ + withForeignPtr s1 $ \p1 -> + withForeignPtr s2 $ \p2 -> do + i <- indexOfDifference + (p1 `ptrElemOff` off1) + (p2 `ptrElemOff` off2) + (min len1 len2) + let pre = if off1 + len1 < off2 + len2 -- share the smaller one + then newPS s1 off1 i + else newPS s2 off2 i + let s1' = newPS s1 (off1 + i) (len1 - i) + let s2' = newPS s2 (off2 + i) (len2 - i) + + return $! (,,) !$ pre !$ s1' !$ s2' + +-- | C-style pointer addition, without the liberal type of 'plusPtr'. +{-# INLINE ptrElemOff #-} +ptrElemOff :: Storable a => Ptr a -> Int -> Ptr a +ptrElemOff p i = p `plusPtr` (i * sizeOf (unsafePerformIO (peek p))) + +{-# INLINE newPS #-} +newPS :: ForeignPtr ByteStringElem -> Int -> Int -> ByteString +newPS s o l = if l <= 0 then S.empty else PS s o l + +{-# INLINE (!$) #-} +(!$) :: (a -> b) -> a -> b +(!$) = ($!) -- fix associativity bug + + +-- Calculates the first index where values differ. +-- BUG: There has to be a smarter algorithm for this +{-# INLINE indexOfDifference #-} +indexOfDifference :: Ptr ByteStringElem -> Ptr ByteStringElem -> Int -> IO Int +indexOfDifference p1 p2 limit = go 0 + where + go n = if n >= limit + then return limit + else do c1 <- peekElemOff p1 n + c2 <- peekElemOff p2 n + if c1 == c2 + then go $! n+1 + else return n + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/Data/Trie/Convenience.hs b/Data/Trie/Convenience.hs new file mode 100644 --- /dev/null +++ b/Data/Trie/Convenience.hs @@ -0,0 +1,96 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} + +---------------------------------------------------------------- +-- ~ 2008.12.19 +-- | +-- Module : Data.Trie.Convenience +-- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- License : BSD3 +-- Maintainer : wren@community.haskell.org +-- Stability : beta +-- Portability : portable +-- +-- Additional convenience versions of the generic functions. +---------------------------------------------------------------- + +module Data.Trie.Convenience + ( + -- * 'lookupBy' variants + lookupWithDefault + + -- * 'alterBy' variants + , insertIfAbsent, insertWith, insertWithKey + , adjustWithKey + , update, updateWithKey + + -- * 'mergeBy' variants + , disunion, unionWith + ) where + +import Data.Trie +import Control.Monad (liftM) + +---------------------------------------------------------------- +---------------------------------------------------------------- + +-- | Lookup a key, returning a default value if it's not found. +lookupWithDefault :: a -> KeyString -> Trie a -> a +lookupWithDefault x = lookupBy_ (\mv _ -> case mv of + Nothing -> x + Just v -> v) x (const x) + +---------------------------------------------------------------- + +-- | Insert a new key, retaining old value on conflict. +insertIfAbsent :: KeyString -> a -> Trie a -> Trie a +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) -> KeyString -> a -> Trie a -> Trie a +insertWith f = alterBy $ \_ x mv -> case mv of + Nothing -> Just x + Just v -> Just (f x v) + +insertWithKey :: (KeyString -> a -> a -> a) -> KeyString -> 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 :: (KeyString -> a -> a -> a) -> KeyString -> a -> Trie a -> (Maybe a, Trie a) +-} + +-- | Apply a function to change the value at a key. +adjustWithKey :: (KeyString -> a -> a) -> KeyString -> Trie a -> Trie a +adjustWithKey f q = alterBy (\k _ -> liftM (f k)) q undefined + +-- | Apply a function to the value at a key, possibly removing it. +update :: (a -> Maybe a) -> KeyString -> Trie a -> Trie a +update f q = alterBy (\_ _ mx -> mx >>= f) q undefined + +updateWithKey :: (KeyString -> a -> Maybe a) -> KeyString -> Trie a -> Trie a +updateWithKey f q = alterBy (\k _ mx -> mx >>= f k) q undefined + +{- +updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> ByteStringTrie a -> (Maybe a, ByteStringTrie a) +-- Also tricky +-} + +---------------------------------------------------------------- + +-- | Combine two tries. If they define the same key, it is removed. +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)) + +{- TODO: (efficiently) +difference, intersection +-} + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,33 @@ +Copyright (c) 2008--2009, wren ng thornton. +ALL RIGHTS RESERVED. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of the copyright holders nor the names of + other contributors may be used to endorse or promote products + derived from this software without specific prior written + permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS +FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN +ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. + diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,7 @@ +#!/usr/bin/env runhaskell + +module Main (main) where +import Distribution.Simple + +main :: IO () +main = defaultMain diff --git a/bytestring-trie.cabal b/bytestring-trie.cabal new file mode 100644 --- /dev/null +++ b/bytestring-trie.cabal @@ -0,0 +1,30 @@ +---------------------------------------------------------------- +-- wren ng thornton ~ 2008.12.19 +---------------------------------------------------------------- + +Name: bytestring-trie +Version: 0.1 +Cabal-Version: >= 1.2 +Build-Type: Simple +Stability: beta +Copyright: Copyright (c) 2008--2009 wren ng thornton +License: BSD3 +License-File: LICENSE +Author: wren ng thornton +Maintainer: wren@community.haskell.org +Homepage: http://code.haskell.org/~wren/ +Category: Data, Data Structures +Synopsis: Efficient map from strings to values. +Description: Efficient map from strings to values. + . + The implementation is based on /big-endian patricia trees/, like "Data.IntMap". We first trie on the elements of "Data.ByteString" and then trie on the big-endian bit representation of those elements. + +Library + Exposed-Modules: Data.Trie + , Data.Trie.Convenience + Other-Modules: Data.Trie.BitTwiddle + , Data.Trie.ByteStringInternal + Build-Depends: base, bytestring + +---------------------------------------------------------------- +----------------------------------------------------------- fin.