diff --git a/bytestring-trie.cabal b/bytestring-trie.cabal --- a/bytestring-trie.cabal +++ b/bytestring-trie.cabal @@ -1,13 +1,13 @@ ---------------------------------------------------------------- --- wren ng thornton ~ 2009.01.04 +-- wren ng thornton ~ 2010.06.10 ---------------------------------------------------------------- Name: bytestring-trie -Version: 0.1.4 +Version: 0.2.2 Cabal-Version: >= 1.2 Build-Type: Simple -Stability: beta -Copyright: Copyright (c) 2008--2009 wren ng thornton +Stability: provisional +Copyright: Copyright (c) 2008--2010 wren ng thornton License: BSD3 License-File: LICENSE Author: wren ng thornton @@ -25,6 +25,26 @@ and other merging operations, but they're also quick for lookups and insertions. +Flag base4 + Description: base-4.0 deprecated Prelude which is imported qualified + Default: True + +Flag useCinternal + 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 + +Flag bytestringInBase + 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 @@ -32,7 +52,37 @@ , Data.Trie.Convenience Other-Modules: Data.Trie.BitTwiddle , Data.Trie.ByteStringInternal - Build-Depends: base, bytestring, binary + -- I think this is all that needs doing to get rid of the warnings? + if flag(base4) + Build-Depends: base >= 4 && < 5 + else + Build-Depends: base < 4 + + 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 + + if flag(applicativeInBase) + Build-Depends: base >= 2.0 + Cpp-Options: -DAPPLICATIVE_IN_BASE + -- BUG (Cabal 1.2 + Haddock): enable for Haddock, disable + -- for Hackage. Fixed in Cabal 1.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 + CC-Options: -O3 + Cpp-Options: -D__USE_C_INTERNAL__ +-- Also need to add stuff to run Configure.hs, FWIW ---------------------------------------------------------------- ----------------------------------------------------------- fin. diff --git a/src/Data/Trie.hs b/src/Data/Trie.hs --- a/src/Data/Trie.hs +++ b/src/Data/Trie.hs @@ -1,17 +1,16 @@ {-# OPTIONS_GHC -Wall -fwarn-tabs #-} ---------------------------------------------------------------- --- ~ 2009.01.05 +-- ~ 2009.01.20 -- | -- Module : Data.Trie -- Copyright : Copyright (c) 2008--2009 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org --- Stability : beta +-- Stability : experimental -- Portability : portable -- -- An efficient implementation of finite 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 @@ -19,23 +18,27 @@ -- -- * 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. +-- +-- 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". ---------------------------------------------------------------- module Data.Trie ( - -- * Data types - Trie(), KeyString, KeyElem + -- * Data type + Trie() -- * Basic functions , empty, null, singleton, size -- * Conversion functions - , fromList, toListBy, toList, keys + , fromList, toListBy, toList, keys, elems -- * Query functions , lookupBy, lookup, member, submap @@ -54,6 +57,7 @@ import qualified Prelude import Data.Trie.Internal +import Data.ByteString (ByteString) import Data.Maybe (isJust) import Control.Monad (liftM) @@ -67,38 +71,47 @@ -- | Convert association list into a trie. On key conflict, values -- earlier in the list shadow later ones. +fromList :: [(ByteString,a)] -> Trie a {-# INLINE fromList #-} -fromList :: [(KeyString,a)] -> Trie a fromList = foldr (uncurry insert) empty -- | Convert trie into association list. Keys will be in sorted order. -toList :: Trie a -> [(KeyString,a)] +toList :: Trie a -> [(ByteString,a)] +{-# INLINE toList #-} toList = toListBy (,) +-- FIX? should 'keys' and 'elems' move to Data.Trie.Convenience instead? + -- | Return all keys in the trie, in sorted order. -keys :: Trie a -> [KeyString] +keys :: Trie a -> [ByteString] +{-# INLINE keys #-} keys = toListBy const +-- | Return all values in the trie, in sorted order according to the keys. +elems :: Trie a -> [a] +{-# INLINE elems #-} +elems = toListBy (flip const) + {--------------------------------------------------------------- -- Query functions (just recurse) ---------------------------------------------------------------} -- | Generic function to find a value (if it exists) and the subtrie -- rooted at the prefix. +lookupBy :: (Maybe a -> Trie a -> b) -> ByteString -> Trie a -> b {-# 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. +lookup :: ByteString -> Trie a -> Maybe a {-# INLINE lookup #-} -lookup :: KeyString -> Trie a -> Maybe a lookup = lookupBy_ const Nothing (const Nothing) -- TODO? move to "Data.Trie.Conventience"? -- | Does a string have a value in the trie? +member :: ByteString -> Trie a -> Bool {-# INLINE member #-} -member :: KeyString -> Trie a -> Bool member q = isJust . lookup q @@ -108,18 +121,19 @@ -- | Insert a new key. If the key is already present, overrides the -- old value +insert :: ByteString -> a -> Trie a -> Trie a {-# INLINE insert #-} -insert :: KeyString -> a -> Trie a -> Trie a insert = alterBy (\_ x _ -> Just x) -- | Apply a function to the value at a key. +adjust :: (a -> a) -> ByteString -> Trie a -> Trie a {-# INLINE adjust #-} -adjust :: (a -> a) -> KeyString -> Trie a -> Trie a adjust f q = alterBy (\_ _ -> liftM f) q undefined +-- TODO: use adjustBy, and benchmark differences -- | Remove the value stored at a key. +delete :: ByteString -> Trie a -> Trie a {-# INLINE delete #-} -delete :: KeyString -> Trie a -> Trie a delete q = alterBy (\_ _ _ -> Nothing) q undefined @@ -129,14 +143,14 @@ -- | Combine two tries, resolving conflicts by choosing the value -- from the left trie. -{-# INLINE unionL #-} unionL :: Trie a -> Trie a -> Trie a +{-# INLINE unionL #-} 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 +{-# INLINE unionR #-} unionR = mergeBy (\_ y -> Just y) ---------------------------------------------------------------- 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 @@ -1,6 +1,8 @@ {-# OPTIONS_GHC -Wall -fwarn-tabs -fno-warn-name-shadowing #-} -{-# OPTIONS_GHC -cpp -fglasgow-exts #-} +-- The MagicHash is for unboxed primitives (-fglasgow-exts also works) +{-# LANGUAGE CPP, MagicHash #-} + ---------------------------------------------------------------- -- ~ 2009.01.05 -- | @@ -8,7 +10,7 @@ -- Copyright : Copyright (c) Daan Leijen 2002 -- License : BSD3 -- Maintainer : libraries@haskell.org, wren@community.haskell.org --- Stability : provisional +-- Stability : stable -- Portability : portable (with CPP) -- -- Functions to treat 'Word' as a bit-vector for big-endian patricia @@ -25,9 +27,10 @@ , mask, shorter, branchMask ) where -import Data.Bits import Data.Trie.ByteStringInternal (ByteStringElem) +import Data.Bits + #if __GLASGOW_HASKELL__ >= 503 import GHC.Exts ( Word(..), Int(..), shiftRL# ) #elif __GLASGOW_HASKELL__ @@ -38,23 +41,20 @@ ---------------------------------------------------------------- --- 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 -{-# INLINE elemToNat #-} elemToNat :: KeyElem -> Word +{-# INLINE elemToNat #-} elemToNat i = fromIntegral i -{-# INLINE natToElem #-} natToElem :: Word -> KeyElem +{-# INLINE natToElem #-} natToElem w = fromIntegral w -{-# INLINE shiftRL #-} shiftRL :: Word -> Int -> Word +{-# INLINE shiftRL #-} #if __GLASGOW_HASKELL__ -- GHC: use unboxing to get @shiftRL@ inlined. shiftRL (W# x) (I# i) = W# (shiftRL# x i) @@ -68,18 +68,18 @@ ---------------------------------------------------------------} -- | Is the value under the mask zero? -{-# INLINE zero #-} zero :: KeyElem -> Mask -> Bool +{-# INLINE zero #-} 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.) -{-# INLINE nomatch #-} nomatch :: KeyElem -> Prefix -> Mask -> Bool +{-# INLINE nomatch #-} nomatch i p m = mask i m /= p -{-# INLINE mask #-} mask :: KeyElem -> Mask -> Prefix +{-# INLINE mask #-} mask i m = maskW (elemToNat i) (elemToNat m) @@ -89,19 +89,19 @@ -- | Get mask by setting all bits higher than the smallest bit in -- @m@. Then apply that mask to @i@. -{-# INLINE maskW #-} maskW :: Word -> Word -> Prefix +{-# INLINE maskW #-} maskW i m = natToElem (i .&. (complement (m-1) `xor` m)) -- | Determine whether the first mask denotes a shorter prefix than -- the second. -{-# INLINE shorter #-} shorter :: Mask -> Mask -> Bool +{-# INLINE shorter #-} shorter m1 m2 = elemToNat m1 > elemToNat m2 -- | Determine first differing bit of two prefixes. -{-# INLINE branchMask #-} branchMask :: Prefix -> Prefix -> Mask +{-# INLINE branchMask #-} branchMask p1 p2 = natToElem (highestBitMask (elemToNat p1 `xor` elemToNat p2)) @@ -150,8 +150,8 @@ into highly efficient machine code. The algorithm is derived from Jorg Arndt's FXT library. ---------------------------------------------------------------} -{-# INLINE highestBitMask #-} highestBitMask :: Word -> Word +{-# INLINE highestBitMask #-} highestBitMask x = case (x .|. shiftRL x 1) of x -> case (x .|. shiftRL x 2) of 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 @@ -1,110 +1,66 @@ +{- +The C algorithm does not appear to give notable performance +improvements, at least when building tries based on /usr/dict on +little-endian 32-bit machines. The implementation also appears +somewhat buggy (cf test/TrieFile.hs) and using the FFI complicates +distribution. + +{-# LANGUAGE CPP, ForeignFunctionInterface #-} +{-# CFILES ByteStringInternal/indexOfDifference.c #-} +-} + {-# OPTIONS_GHC -Wall -fwarn-tabs #-} ---------------------------------------------------------------- --- ~ 2008.12.24 +-- ~ 2009.02.06 -- | -- Module : Data.Trie.ByteStringInternal -- Copyright : Copyright (c) 2008--2009 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org --- Stability : beta +-- Stability : experimental -- Portability : portable -- --- Helper functions on 'ByteString's for "Data.Trie". +-- Helper functions on 'ByteString's for "Data.Trie.Internal". ---------------------------------------------------------------- module Data.Trie.ByteStringInternal ( ByteString, ByteStringElem - {-, unsafeWordHead-} - , splitMaximalPrefix + , breakMaximalPrefix ) 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) +import Foreign.Storable (Storable(..)) -import Foreign.ForeignPtr (ForeignPtr, withForeignPtr) -import Foreign.Ptr (Ptr, plusPtr{-, castPtr-}) -import Foreign.Storable (Storable(..)) -import System.IO.Unsafe (unsafePerformIO) {- -import Foreign.Marshal.Alloc (alloca) -import Data.Bits +#ifdef __USE_C_INTERNAL__ +import Foreign.C.Types (CInt) +import Control.Monad (liftM) +#endif -} ----------------------------------------------------------------- -type ByteStringElem = Word8 -- Associated type of ByteString - -{- -- Not used at present --- | The size of 'Word' in bytes. -{-# INLINE sizeOfWord #-} -sizeOfWord :: Int -sizeOfWord = sizeOf (undefined :: Word) - ---------------------------------------------------------------- --- | Return the first natural 'Word' worth of string, padding by --- zeros as necessary. The position of elements within the word --- varies by architecture, hence this function is quasi-unsafe to --- use for trieing on the bit-vector representation. A safer version --- may be forthcoming, though trieing on multiple bytes at once --- appears impractical at the moment. -unsafeWordHead :: ByteString -> Word -unsafeWordHead (PS s o l) = inlinePerformIO $ - withForeignPtr s $ \p -> - liftM (maskInitialBytes l .&.) - (peek (p `plusPtr` o :: Ptr Word)) - -maskInitialBytes :: Int -> Word -maskInitialBytes byteCount - | isLittleEndian = case effectiveByteCount of - 0 -> 0x0000000000000000 - 1 -> 0x00000000000000FF - 2 -> 0x000000000000FFFF - 3 -> 0x0000000000FFFFFF - 4 -> 0x00000000FFFFFFFF - 5 -> 0x000000FFFFFFFFFF - 6 -> 0x0000FFFFFFFFFFFF - 7 -> 0x00FFFFFFFFFFFFFF - _ -> 0xFFFFFFFFFFFFFFFF - | otherwise = case effectiveByteCount of - 0 -> 0x0000000000000000 - 1 -> 0xFF00000000000000 - 2 -> 0xFFFF000000000000 - 3 -> 0xFFFFFF0000000000 - 4 -> 0xFFFFFFFF00000000 - 5 -> 0xFFFFFFFFFF000000 - 6 -> 0xFFFFFFFFFFFF0000 - 7 -> 0xFFFFFFFFFFFFFF00 - _ -> 0xFFFFFFFFFFFFFFFF - where - effectiveByteCount = 0 `max` (byteCount `min` sizeOfWord) +-- | Associated type of 'ByteString' +type ByteStringElem = Word8 --- 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.") - --- End unused code -} - ---------------------------------------------------------------- -- | Returns the longest shared prefix and the two remaining suffixes -- for a pair of strings. -splitMaximalPrefix :: ByteString -> ByteString +-- +-- > s == (\(pre,s',z') -> pre `append` s') (breakMaximalPrefix s z) +-- > z == (\(pre,s',z') -> pre `append` z') (breakMaximalPrefix s z) + +breakMaximalPrefix :: ByteString -> ByteString -> (ByteString, ByteString, ByteString) -splitMaximalPrefix +breakMaximalPrefix str1@(PS s1 off1 len1) str2@(PS s2 off2 len2) | len1 == 0 = (S.empty, S.empty, str2) @@ -125,44 +81,42 @@ 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 ptrElemOff #-} +ptrElemOff p i = + p `plusPtr` (i * sizeOf (undefined `asTypeOf` inlinePerformIO (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 newPS #-} +newPS s o l = + if l <= 0 then S.empty else PS s o l -{-# INLINE (!$) #-} +-- | fix associativity bug (!$) :: (a -> b) -> a -> b -(!$) = ($!) -- fix associativity bug +{-# INLINE (!$) #-} +(!$) = ($!) +---------------------------------------------------------------- -- | Calculates the first index where values differ. --- --- BUG: There has to be a smarter algorithm for this. In particular, --- it'd be nice if there was a function like C's @memcmp@, but which --- returns the index rather than the 'Ordering'. If it were agnostic --- ot architecture, that'd be even nicer. --- --- TODO: rather than revert to @goByte@, we should use xor- and --- mask-munging so @goWord@ can reuse the information it already --- has to discover the byte of difference. The trick is making that --- faster than the current version. And ensuring alignment. -{-# INLINE indexOfDifference #-} + indexOfDifference :: Ptr ByteStringElem -> Ptr ByteStringElem -> Int -> IO Int +{- +#ifdef __USE_C_INTERNAL__ + +indexOfDifference p q i = + liftM fromIntegral $! c_indexOfDifference p q (fromIntegral i) + +-- This could probably be not IO, but the wrapper requires that anyways... +foreign import ccall unsafe "ByteStringInternal/indexOfDifference.h indexOfDifference" + c_indexOfDifference :: Ptr ByteStringElem -> Ptr ByteStringElem -> CInt -> IO CInt + +#else +-} + +-- Use the naive algorithm which doesn't depend on architecture details indexOfDifference p1 p2 limit = goByte 0 where - {- BUG: using this assumes ByteStrings are Word-aligned - goWord n = if n + sizeOfWord >= limit - then goByte n - else do w1 <- peek (p1 `plusPtr` n :: Ptr Word) - w2 <- peek (p2 `plusPtr` n :: Ptr Word) - if w1 == w2 - then goWord $! n + sizeOfWord - else goByte n - -} - goByte n = if n >= limit then return limit else do c1 <- peekElemOff p1 n @@ -170,6 +124,9 @@ if c1 == c2 then goByte $! n+1 else return n +{- +#endif +-} ---------------------------------------------------------------- ----------------------------------------------------------- fin. diff --git a/src/Data/Trie/ByteStringInternal/indexOfDifference.c b/src/Data/Trie/ByteStringInternal/indexOfDifference.c new file mode 100644 --- /dev/null +++ b/src/Data/Trie/ByteStringInternal/indexOfDifference.c @@ -0,0 +1,175 @@ +/* ------------------------------------------------------------- +-- ~ 2009.01.07 +-- | +-- Module : Data.Trie.ByteStringInternal.indexOfDifference +-- Copyright : Copyright (c) 2008--2009 wren ng thornton +-- License : BSD3 +-- Maintainer : wren@community.haskell.org +-- Stability : beta +-- Portability : portable +-- +-- More efficient implementation (in theory) than testing characters +-- byte by byte. However, the gains seem to be minimal and ---more +-- importantly--- there seem to be obscure bugs that only show up +-- in extensive testing (e.g. running TrieFile on /usr/dict). Maybe +-- worth pursuing further, especially if SSE's 126-bit registers +-- or GCC's vectors can be leveraged in a portable way. +------------------------------------------------------------- */ + +/* Defines certain architecture characteristics. + * Created by Configure.hs */ +#include "indexOfDifference.h" + +typedef unsigned char Word8; +typedef unsigned short Word16; +typedef unsigned long Word32; +typedef unsigned long long Word64; + +/* Hopefully this makes Nat the most optimal size, or the same as int */ +#ifdef __hDataTrie_Nat64__ + typedef Word64 Nat; +# define MIN_NAT ((Nat) 0x8000000000000000ull) +#elif defined(__hDataTrie_Nat32__) + typedef Word32 Nat; +# define MIN_NAT ((Nat) 0x80000000) +#else + /* No definition. Unknown architecture. + * Maybe it'd be worth supporting 16-bit as well...? + * or maybe fail back to 8-bit? */ +#endif + +/* cf also */ +#define NAT_MISALIGNMENT(p) (((int)(p)) % sizeof(Nat)) + + +#define READ_WORD8(p) (*((Word8*) (p))) +#define READ_NAT(p) (*((Nat*) (p))) + + +/* ---------------------------------------------------------- */ +/* Compare up to the first @limit@ bytes of @p1@ against @p2@ + * and return the first index where they differ. */ + +/* TODO: Consider replacing loops by Duff's Device, or similar + */ + +int indexOfDifference(const void* p1, const void* p2, const int limit) { + /* @i@ measures how many bytes are shared, + * Thus it's the 0-based index of difference. */ + int i = 0; + if (limit <= 0) return 0; + + + /* Munge until Nat-aligned (They seem to always be). + * Should fail back to this naive version on unknown arch */ + { + int x1 = NAT_MISALIGNMENT(p1); + int x2 = NAT_MISALIGNMENT(p2); + if (x1 != x2) { + /* FIX: what if they're misaligned differently? + * For now we'll use Word8 all the way */ + x1 = limit; + } + while (i < x1) { + if (READ_WORD8(p1+i) == READ_WORD8(p2+i)) { + i += sizeof(Word8); + } else { + return i; + } + } + if (x1 == limit) { + return limit; + } else { + /* Fall through */ + } + } + + + /* Check one Nat at a time until we find a mismatch */ + Nat diff; + do { + diff = READ_NAT(p1+i) ^ READ_NAT(p2+i); + + /* If the diff is valid and zero, then increment the loop */ + if (i + sizeof(Nat) <= limit && diff == 0) { + i += sizeof(Nat); + } else { + break; + } + } while (i < limit); + + + /* Trim incomplete Nat at the end of the strings */ + if (i + sizeof(Nat) > limit) { + #ifdef __hDataTrie_isLittleEndian__ + diff &= ((1 << ((limit-i) * 8*sizeof(Word8))) - 1); + #else + diff &= MIN_NAT >> ((limit-i) * 8*sizeof(Word8) - 1); + #endif + + if (0 == diff) return limit; + } + + + /* Found a difference. Do binary search to identify first + * byte in Nat which doesn't match. */ + Word32 w32; + #ifdef __hDataTrie_Nat64__ + # ifdef __hDataTrie_isLittleEndian__ + const Word32 first32 = (Word32) (diff & 0x00000000FFFFFFFFull); + # else + const Word32 first32 = (Word32)((diff & 0xFFFFFFFF00000000ull) >> 32); + # endif + if (0 == first32) { + i += 4; + # ifdef __hDataTrie_isLittleEndian__ + w32 = (Word32)((diff & 0xFFFFFFFF00000000ull) >> 32); + # else + w32 = (Word32) (diff & 0x00000000FFFFFFFFull); + # endif + } else { + w32 = first32; + } + #elif defined(__hDataTrie_Nat32__) + w32 = diff; + #else + /* WTF? */ + #endif + + + Word16 w16; + #if defined(__hDataTrie_Nat32__) || defined(__hDataTrie_Nat64__) + # ifdef __hDataTrie_isLittleEndian__ + const Word16 first16 = (Word16) (w32 & 0x0000FFFF); + # else + const Word16 first16 = (Word16)((w32 & 0xFFFF0000) >> 16); + # endif + if (0 == first16) { + i += 2; + # ifdef __hDataTrie_isLittleEndian__ + w16 = (Word16)((w32 & 0xFFFF0000) >> 16); + # else + w16 = (Word16) (w32 & 0x0000FFFF); + # endif + } else { + w16 = first16; + } + #else + /* WTF? */ + #endif + + + Word8 w8; + #ifdef __hDataTrie_isLittleEndian__ + const Word8 first8 = (Word8) (w16 & 0x00FF); + #else + const Word8 first8 = (Word8)((w16 & 0xFF00) >> 8); + #endif + if (0 == first8) { + i += 1; + } + + return i; +} +/* ---------------------------------------------------------- */ +/* ----------------------------------------------------- fin. */ 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,76 +1,91 @@ {-# OPTIONS_GHC -Wall -fwarn-tabs #-} ---------------------------------------------------------------- --- ~ 2009.01.05 +-- ~ 2009.01.20 -- | -- Module : Data.Trie.Convenience -- Copyright : Copyright (c) 2008--2009 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org --- Stability : beta +-- Stability : provisional -- Portability : portable -- --- Additional convenience versions of the generic functions. +-- Additional convenience functions. In order to keep "Data.Trie" +-- concise, non-essential and uncommonly used functions have been +-- moved here. Most of these functions simplify the generic functions +-- from "Data.Trie", following after the interface for "Data.Map" +-- and "Data.IntMap". ---------------------------------------------------------------- module Data.Trie.Convenience ( + -- * Conversion functions + -- $fromList + fromListL, fromListR, fromListS, fromListWith + -- * 'lookupBy' variants - lookupWithDefault + , lookupWithDefault -- * 'alterBy' variants , insertIfAbsent, insertWith, insertWithKey , adjustWithKey , update, updateWithKey - -- ** Conversion functions - -- $fromList - , fromListL, fromListR, fromListS - -- * 'mergeBy' variants , disunion, unionWith ) where import Data.Trie import Data.Trie.Internal (lookupBy_) +import Data.ByteString (ByteString) import Data.List (foldl', sortBy) +import Data.Ord (comparing) import Control.Monad (liftM) ---------------------------------------------------------------- ---------------------------------------------------------------- - -- $fromList --- Just like 'fromList' both of these functions convert an association +-- Just like 'fromList' all of these functions convert an association -- list into a trie, with earlier values shadowing later ones when -- keys conflict. Depending on the order of keys in the list, there --- can be as much as 5x speed difference between the two. Yet, --- performance is about the same when matching best-case to best-case --- and worst-case to worst-case (which is which is swapped when --- reversing the list or changing which function is used). +-- can be as much as 5x speed difference between the left and right +-- variants. Yet, performance is about the same when matching +-- best-case to best-case and worst-case to worst-case (which is +-- which is swapped when reversing the list or changing which +-- function is used). +fromListL :: [(ByteString,a)] -> Trie a {-# INLINE fromListL #-} -fromListL :: [(KeyString,a)] -> Trie a 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'. +fromListR :: [(ByteString,a)] -> Trie a {-# INLINE fromListR #-} -fromListR :: [(KeyString,a)] -> Trie a fromListR = fromList +-- TODO: compare performance against a fromListL definition, adjusting the sort +-- -- | This version 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. +fromListS :: [(ByteString,a)] -> Trie a {-# INLINE fromListS #-} -fromListS :: [(KeyString,a)] -> Trie a -fromListS = fromListR . sortBy (\(k,_) (q,_) -> k `compare` q) +fromListS = fromListR . sortBy (comparing fst) +-- | A variant of 'fromListR' that takes a function for combining values on conflict. +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) ---------------------------------------------------------------- -- | Lookup a key, returning a default value if it's not found. -lookupWithDefault :: a -> KeyString -> Trie a -> a +lookupWithDefault :: a -> ByteString -> Trie a -> a lookupWithDefault x = lookupBy_ (\mv _ -> case mv of Nothing -> x Just v -> v) x (const x) @@ -78,39 +93,39 @@ ---------------------------------------------------------------- -- | Insert a new key, retaining old value on conflict. -insertIfAbsent :: KeyString -> a -> Trie a -> Trie a +insertIfAbsent :: ByteString -> 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 :: (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) -insertWithKey :: (KeyString -> a -> a -> a) -> KeyString -> a -> Trie a -> Trie a +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 :: (KeyString -> a -> a -> a) -> KeyString -> a -> Trie a -> (Maybe a, Trie a) +insertLookupWithKey :: (ByteString -> a -> a -> a) -> ByteString -> 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 :: (ByteString -> a -> a) -> ByteString -> 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 :: (a -> Maybe a) -> ByteString -> Trie a -> Trie a update f q = alterBy (\_ _ mx -> mx >>= f) q undefined -updateWithKey :: (KeyString -> a -> Maybe a) -> KeyString -> Trie a -> Trie a +updateWithKey :: (ByteString -> a -> Maybe a) -> ByteString -> 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) +updateLookupWithKey :: (ByteString -> a -> Maybe a) -> ByteString -> Trie a -> (Maybe a, Trie a) -- Also tricky -} 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,42 +1,54 @@ {-# OPTIONS_GHC -Wall -fwarn-tabs #-} +-- For list fusion on toListBy, and for applicative hiding +{-# LANGUAGE CPP #-} + ---------------------------------------------------------------- --- ~ 2009.01.10 +-- ~ 2009.01.20 -- | -- Module : Data.Trie.Internal -- Copyright : Copyright (c) 2008--2009 wren ng thornton -- License : BSD3 -- Maintainer : wren@community.haskell.org --- Stability : beta --- Portability : portable +-- Stability : provisional +-- Portability : portable (with CPP) -- -- Internal definition of the 'Trie' data type and generic functions -- for manipulating them. Almost everything here is re-exported --- from "Data.Trie". +-- from "Data.Trie", which is the preferred API for users. This +-- module is for developers who need deeper (and potentially fragile) +-- access to the abstract type. ---------------------------------------------------------------- module Data.Trie.Internal ( -- * Data types - Trie(), KeyString, KeyElem, showTrie + Trie(), showTrie + -- * Functions for 'ByteString's + , breakMaximalPrefix + -- * Basic functions , empty, null, singleton, size - -- * Conversion functions - , toListBy + -- * Conversion and folding functions + , foldrWithKey, toListBy -- * Query functions , lookupBy_, submap -- * Single-value modification - , alterBy + , alterBy, adjustBy -- * Combining tries , mergeBy -- * Mapping functions , mapBy, filterMap + + -- * Priority-queue functions + , minAssoc, maxAssoc + , updateMinViewBy, updateMaxViewBy ) where import Prelude hiding (null, lookup) @@ -46,31 +58,42 @@ import Data.Trie.ByteStringInternal import Data.Trie.BitTwiddle -import Control.Monad (liftM3, liftM4) -import Control.Applicative (Applicative(..), (<$>)) +import Data.Binary + import Data.Monoid (Monoid(..)) -import Data.Foldable (Foldable(foldMap)) +import Control.Monad (liftM, liftM3, liftM4) +#ifdef APPLICATIVE_IN_BASE +import Control.Monad (ap) +import Control.Applicative (Applicative(..), (<$>)) +import Data.Foldable (Foldable(..)) import Data.Traversable (Traversable(traverse)) -import Data.Binary +#endif + +#ifdef __GLASGOW_HASKELL__ +import GHC.Exts (build) +#endif ---------------------------------------------------------------- ---------------------------------------------------------------- {--------------------------------------------------------------- --- ByteString Big-endian Patricia Trie +-- ByteString Big-endian Patricia Trie datatype ---------------------------------------------------------------} -type KeyString = ByteString -type KeyElem = ByteStringElem +{- +In our idealized representation, we use a (directed) discrete graph +to represent our finite state machine. To organize the set of +outgoing arcs from a given Node we have ArcSet be a big-endian +patricia tree like Data.IntMap. In order to simplify things we then +go through a series of derivations. -{- 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 Arc a = Arc ByteString (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] + | Start ByteString (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 @@ -82,43 +105,44 @@ data Node a = Accept a (ArcSet a) | Reject (Branch a) data ArcSet a = None - | Arc KeyString (Node a) + | Arc ByteString (Node a) | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) data Trie a = Empty - | Start KeyString (Node a) + | Start ByteString (Node a) -- Squash Node together: -- (most likely good) data Node a = Node (Maybe a) (ArcSet a) data ArcSet a = None - | Arc KeyString (Node a) + | Arc ByteString (Node a) | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) data Trie a = Empty - | Start KeyString (Node a) + | Start ByteString (Node a) -- Squash Empty/None and Arc/Start together: --- (Complicates invariants about non-empty strings and Node's recursion) +-- (This complicates invariants about non-empty strings and Node's +-- recursion, but those can be circumvented by using smart +-- constructors.) data Node a = Node (Maybe a) (ArcSet a) data Trie a = Empty - | Arc KeyString (Node a) + | Arc ByteString (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 +-- * ByteString 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... +-- [2] Maybe we shouldn't unpack the ByteString. We could specialize +-- or inline the breakMaximalPrefix function to prevent constructing +-- a new ByteString 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@. -- @@ -129,7 +153,7 @@ -- be curious to know. data Trie a = Empty - | Arc {-# UNPACK #-} !KeyString + | Arc {-# UNPACK #-} !ByteString !(Maybe a) !(Trie a) | Branch {-# UNPACK #-} !Prefix @@ -138,9 +162,24 @@ !(Trie a) deriving Eq -- Prefix/Mask should be deterministic regardless of insertion order - -- TODO: verify this is so. + -- TODO: prove this is so. +-- TODO? add Ord instance like Data.Map? + +{--------------------------------------------------------------- +-- Trie instances: serialization et cetera +---------------------------------------------------------------} + +-- This instance does not unveil the innards of our abstract type. +-- It doesn't emit truly proper Haskell code though, since ByteStrings +-- are printed as (ASCII) Strings, but that's not our fault. (Also +-- 'fromList' is in "Data.Trie" instead of here.) +instance (Show a) => Show (Trie a) where + showsPrec p t = showParen (p > 10) + $ ("Data.Trie.fromList "++) . shows (toListBy (,) t) + + -- | Visualization fuction for debugging. showTrie :: (Show a) => Trie a -> String showTrie t = shows' id t "" @@ -161,11 +200,10 @@ in s' . shows' (ss . (spaces s' ++)) t' -{--------------------------------------------------------------- --- Trie instances ----------------------------------------------------------------} +-- TODO?? a Read instance? hrm... should I? -instance Binary a => Binary (Trie a) where +-- 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 @@ -184,6 +222,10 @@ _ -> liftM4 Branch get get get get +{--------------------------------------------------------------- +-- Trie instances: Abstract Nonsense +---------------------------------------------------------------} + instance Functor Trie where fmap _ Empty = Empty fmap f (Arc k Nothing t) = Arc k Nothing (fmap f t) @@ -191,6 +233,47 @@ fmap f (Branch p m l r) = Branch p m (fmap f l) (fmap f r) +#ifdef APPLICATIVE_IN_BASE +instance Foldable Trie where + -- 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 + + {- This definition is much faster, but it's also wrong + -- (or at least different than foldrWithKey) + foldr f = \z t -> go t id z + where + go Empty k x = k x + go (Branch _ _ l r) k x = go r (go l k) x + go (Arc _ Nothing t) k x = go t k x + go (Arc _ (Just v) t) k x = go t k (f v x) + + foldl f = \z t -> go t id z + where + go Empty k x = k x + go (Branch _ _ l r) k x = go l (go r k) x + go (Arc _ Nothing t) k x = go t k x + go (Arc _ (Just v) t) k x = go t k (f x v) + -} + +-- TODO: newtype Keys = K Trie ; instance Foldable Keys +-- 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 + +instance Applicative Trie where + pure = return + (<*>) = ap +#endif + -- 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 @@ -212,31 +295,38 @@ unionL = mergeBy (\x _ -> Just x) -instance Monoid a => Monoid (Trie a) where +-- This instance is more sensible than Data.IntMap and Data.Map's +instance (Monoid a) => Monoid (Trie a) where mempty = empty 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 - 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 +-- Since the Monoid instance isn't natural in @a@, I can't think +-- of any other sensible instance for MonadPlus. It's as specious +-- as Maybe, IO, and STM's instances though. +-- +-- MonadPlus laws: +-- 1. forms a monoid +-- 2. mzero >>= f === mzero +-- 3. m >> mzero === mzero +-- 4. mplus m n >>= k === mplus (m >>= k) (n >>= k) +-- 4' mplus (return a) n === return a +{- +-- Follows #1, #1, and #3. But it does something like 4' instead +-- of actually doing #4 (since we'd merge the trees generated by +-- @k@ for conflicting values) +-- +-- TODO: cf Control.Applicative.Alternative (base-4, but not Hugs). +-- But (<*>) gets odd when the function is not 'pure'... maybe +-- helpful though. +instance MonadPlus Trie where + mzero = empty + mplus = unionL where unionL = mergeBy (\x _ -> Just x) +-} {--------------------------------------------------------------- --- Mapping functions +-- Extra mapping functions ---------------------------------------------------------------} -- | Apply a function to all values, potentially removing them. @@ -249,7 +339,7 @@ -- | 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 :: (ByteString -> a -> Maybe b) -> Trie a -> Trie b mapBy f = go S.empty where go _ Empty = empty @@ -271,7 +361,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 :: KeyString -> Maybe a -> Trie a -> Trie a +arc :: ByteString -> 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 @@ -310,7 +400,8 @@ -- Error messages ---------------------------------------------------------------} -errorLogHead :: String -> KeyString -> KeyElem +errorLogHead :: String -> ByteString -> ByteStringElem +{-# NOINLINE errorLogHead #-} errorLogHead s q | S.null q = error (s ++": found null subquery") | otherwise = S.head q @@ -322,30 +413,29 @@ -- Basic functions ---------------------------------------------------------------} --- | /O(1)/, The empty trie. -{-# INLINE empty #-} +-- | /O(1)/, Construct the empty trie. empty :: Trie a +{-# INLINE empty #-} empty = Empty -- | /O(1)/, Is the trie empty? -{-# INLINE null #-} null :: Trie a -> Bool +{-# INLINE null #-} null Empty = True null _ = False --- | /O(1)/, A singleton trie. +-- | /O(1)/, Construct a singleton trie. +singleton :: ByteString -> a -> Trie a {-# 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 +{-# INLINE size #-} size t = size' t id 0 --- | /O(n)/, Internal CPS accumulator function for calculating --- 'size'. +-- | /O(n)/, CPS accumulator helper 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 @@ -362,26 +452,62 @@ -- TODO: rewrite list-catenation to be lazier (real CPS instead of -- function building? is the function building really better than -- (++) anyways?) +-- N.B. If our manual definition of foldr/foldl (using function +-- application) is so much faster than the default Endo definition +-- (using function composition), then we should make this use +-- application instead too. -- -- TODO: the @q@ accumulator should be lazy ByteString and only --- forced by @f@. It's already non-strict, but we should ensure +-- forced by @fcons@. 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. -toListBy :: (KeyString -> a -> b) -> Trie a -> [b] -toListBy f = \t -> go S.empty t [] +-- BUG: not safe for deep strict @fcons@, only for WHNF-strict like (:) +-- Where to put the strictness to amortize it? +-- +-- | Convert a trie into a list (in key-sorted order) using a +-- function, folding the list as we go. +foldrWithKey :: (ByteString -> a -> b -> b) -> b -> Trie a -> b +foldrWithKey fcons nil = \t -> go S.empty t nil 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 -> (f k' v :) . 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] +#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) + +-- TODO: should probably have a specialized version for strictness, +-- and a rule to rewrite generic lazy version into it. As per +-- Data.ByteString.unpack and the comments there about strictness +-- and fusion. +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 + + {--------------------------------------------------------------- -- Query functions (just recurse) ---------------------------------------------------------------} @@ -395,7 +521,7 @@ -- This function is intended for internal use. For the public-facing -- version, see @lookupBy@ in "Data.Trie". lookupBy_ :: (Maybe a -> Trie a -> b) -> b -> (Trie a -> b) - -> KeyString -> Trie a -> b + -> ByteString -> Trie a -> b lookupBy_ f z a = lookupBy_' where -- | Deal with epsilon query (when there is no epsilon value) @@ -406,7 +532,7 @@ go _ Empty = z go q (Arc k mv t) = - let (_,k',q') = splitMaximalPrefix k q + let (_,k',q') = breakMaximalPrefix k q in case (not $ S.null k', S.null q') of (True, True) -> a (Arc k' mv t) (True, False) -> z @@ -436,8 +562,8 @@ -- of (>>=) for epsilon`elem`t) -- -- | Return the subtrie containing all keys beginning with a prefix. +submap :: ByteString -> Trie a -> Trie a {-# 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_@ @@ -447,14 +573,17 @@ submap' mx t = Arc q mx t errorInvariantBroken :: String -> String -> a +{-# NOINLINE errorInvariantBroken #-} errorInvariantBroken s e = error (s ++ ": Invariant was broken" ++ e') where e' = if Prelude.null e then e else ", found: " ++ e errorArcAfterNothing :: String -> a +{-# NOINLINE errorArcAfterNothing #-} errorArcAfterNothing s = errorInvariantBroken s "Arc after Nothing" errorEmptyAfterNothing :: String -> a +{-# NOINLINE errorEmptyAfterNothing #-} errorEmptyAfterNothing s = errorInvariantBroken s "Empty after Nothing" -- -} @@ -471,8 +600,8 @@ -- -- | 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 :: (ByteString -> a -> Maybe a -> Maybe a) + -> ByteString -> a -> Trie a -> Trie a alterBy f_ q_ x_ | S.null q_ = alterEpsilon | otherwise = go q_ @@ -496,7 +625,7 @@ qh = errorLogHead "alterBy" q go q t_@(Arc k mv t) = - let (p,k',q') = splitMaximalPrefix k q + 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) @@ -516,6 +645,36 @@ (False, False) -> arc k mv (go q' t) +-- | ... +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_ + | otherwise = go q_ + where + f = f_ q_ x_ + + go _ Empty = Empty + + go q t@(Branch p m l r) + | nomatch qh p m = t + | zero qh m = Branch p m (go q l) r + | otherwise = Branch p m l (go q r) + where + 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) + + {--------------------------------------------------------------- -- Trie-combining functions ---------------------------------------------------------------} @@ -580,7 +739,7 @@ go' (Arc k0 mv0 t0) (Arc k1 mv1 t1) | m' == 0 = - let (pre,k0',k1') = splitMaximalPrefix k0 k1 + 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 #-} @@ -588,7 +747,7 @@ 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,True) -> arcMerge mv1 (Arc k0' mv0 t0) t1 (False,False) -> arcMerge Nothing (Arc k0' mv0 t0) (Arc k1' mv1 t1) go' (Arc _ _ _) @@ -599,19 +758,69 @@ 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) + | zero p1 m0 = branch p0 m0 (go l t1_) r + | otherwise = branch p0 m0 l (go r t1_) -- 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 +{-# INLINE mergeMaybe #-} mergeMaybe _ Nothing Nothing = Nothing mergeMaybe _ Nothing mv1@(Just _) = mv1 mergeMaybe _ mv0@(Just _) Nothing = mv0 mergeMaybe f (Just v0) (Just v1) = f v0 v1 + + +{--------------------------------------------------------------- +-- Priority-queue functions +---------------------------------------------------------------} + +minAssoc :: Trie a -> Maybe (ByteString, a) +minAssoc = go S.empty + where + go _ Empty = Nothing + go q (Arc k (Just v) _) = Just (S.append q k,v) + go q (Arc k Nothing t) = go (S.append q k) t + go q (Branch _ _ l _) = go q l + + +maxAssoc :: Trie a -> Maybe (ByteString, a) +maxAssoc = go S.empty + where + go _ Empty = Nothing + go q (Arc k (Just v) Empty) = Just (S.append q k,v) + go q (Arc k _ t) = go (S.append q k) t + go q (Branch _ _ _ r) = go q r + + +mapView :: (Trie a -> Trie a) + -> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a) +mapView _ Nothing = Nothing +mapView f (Just (k,v,t)) = Just (k,v, f t) + + +updateMinViewBy :: (ByteString -> a -> Maybe a) + -> Trie a -> Maybe (ByteString, a, Trie a) +updateMinViewBy f = go S.empty + where + go _ Empty = Nothing + 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) + + +updateMaxViewBy :: (ByteString -> a -> Maybe a) + -> Trie a -> Maybe (ByteString, a, Trie a) +updateMaxViewBy f = go S.empty + where + go _ Empty = Nothing + 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) ---------------------------------------------------------------- ----------------------------------------------------------- fin.