radix-tree (empty) → 0.1
raw patch · 7 files changed
+924/−0 lines, 7 filesdep +HUnitdep +QuickCheckdep +basesetup-changed
Dependencies added: HUnit, QuickCheck, base, bytestring, containers, deepseq, gauge, hashtables, primitive, radix-tree, tasty, tasty-hunit, tasty-quickcheck, text, unordered-containers
Files
- LICENSE +27/−0
- Setup.hs +2/−0
- bench/RadixTreeBench.hs +146/−0
- radix-tree.cabal +143/−0
- src/Data/RadixTree.hs +33/−0
- src/Data/RadixTree/Internal.hs +454/−0
- test/TestMain.hs +119/−0
+ LICENSE view
@@ -0,0 +1,27 @@+Copyright (c) 2018 Sergey Vinokurov++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions+are met:+1. Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.+2. 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.+3. Neither the name of the author nor the names of his contributors+ may be used to endorse or promote products derived from this software+ without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE REGENTS 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 AUTHORS 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.
+ Setup.hs view
@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain
+ bench/RadixTreeBench.hs view
@@ -0,0 +1,146 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}++module Main (main) where++import Control.Arrow+import Control.DeepSeq+import Control.Exception++import Data.Foldable++import qualified Data.ByteString.Short as BSS+import qualified Data.HashMap.Strict as HM+import qualified Data.Map.Strict as M+import qualified Data.Text.Encoding as TE+import qualified Data.Text.Lazy as TL+import qualified Data.Text.Lazy.IO as TLIO++import qualified Data.HashTable.IO as HT++import Gauge++import qualified Data.RadixTree.Internal as RT++main :: IO ()+main = do+ let config = defaultConfig+ { resamples = 10000+ , displayMode = Condensed+ , rerunsLimit = 1+ }++ contents <- TLIO.readFile "/tmp/tags-ebac8dcc87fd1f1b1e7016d6585549309e3c5016-haskell-mode"+ let tags :: [TL.Text]+ tags = filter (not . TL.null) $ map (head . TL.splitOn "\t") $ drop 1 $ TL.lines contents++ decodeBS = TE.encodeUtf8 . TL.toStrict+ decode = BSS.toShort . decodeBS++ tags' :: [BSS.ShortByteString]+ tags' = map decode tags++ tags'' :: [(BSS.ShortByteString, ())]+ tags'' = map (id &&& const ()) tags'++ tagsRev'' :: [(BSS.ShortByteString, ())]+ tagsRev'' = map ((BSS.pack . reverse . BSS.unpack) &&& const ()) tags'++ -- tagsBS :: [(BS.ByteString, ())]+ -- tagsBS = map (decodeBS &&& const ()) tags++ queriesPresent :: [BSS.ShortByteString]+ queriesPresent = tags' ++ map (BSS.pack . reverse . BSS.unpack) tags'++ queriesMissing :: [BSS.ShortByteString]+ queriesMissing = map (BSS.pack . reverse . BSS.unpack) tags'++ queriesBoth :: [BSS.ShortByteString]+ queriesBoth = tags' ++ map (BSS.pack . reverse . BSS.unpack) tags'++ evaluate $ rnf tags'+ evaluate $ rnf tags''+ evaluate $ rnf tagsRev''+ evaluate $ rnf queriesPresent+ evaluate $ rnf queriesMissing+ evaluate $ rnf queriesBoth++ let radixTree = RT.fromList tags''+ radixTreeRev = RT.fromList tagsRev''+ treeMap = M.fromList tags''+ treeMapRev = M.fromList tagsRev''+ hashMap = HM.fromList tags''+ hashMapRev = HM.fromList tagsRev''++ evaluate $ rnf radixTree+ evaluate $ rnf radixTreeRev+ evaluate $ rnf treeMap+ evaluate $ rnf treeMapRev+ evaluate $ rnf hashMap+ evaluate $ rnf hashMapRev++ (basic :: HT.BasicHashTable BSS.ShortByteString ()) <- HT.new+ -- (linear :: HT.LinearHashTable BSS.ShortByteString ()) <- HT.new+ (cuckoo :: HT.CuckooHashTable BSS.ShortByteString ()) <- HT.new+ for_ tags'' $ \(k, v) -> do+ HT.insert basic k v+ -- HT.insert linear k v+ HT.insert cuckoo k v++ defaultMainWith config+ [ bgroup "creation"+ [ bench "Data.RadixTree" $ nf RT.fromList tags''+ , bench "Data.Map" $ nf M.fromList tags''+ , bench "Data.HashMap" $ nf HM.fromList tags''+ , bench "BasicHashTable" $ nfIO $ do+ (ht :: HT.BasicHashTable BSS.ShortByteString ()) <- HT.new+ for_ tags'' $ \(k, v) -> HT.insert ht k v+ -- , bench "LinearHashTable" $ nfIO $ do+ -- (ht :: HT.LinearHashTable BSS.ShortByteString ()) <- HT.new+ -- for_ tags'' $ \(k, v) -> HT.insert ht k v+ , bench "CuckooHashTable" $ nfIO $ do+ (ht :: HT.CuckooHashTable BSS.ShortByteString ()) <- HT.new+ for_ tags'' $ \(k, v) -> HT.insert ht k v+ ]+ , bgroup "lookup"+ [ bgroup "present"+ [ bench "Data.RadixTree" $ nf (map (`RT.lookup` radixTree)) queriesPresent+ , bench "Data.Map" $ nf (map (`M.lookup` treeMap)) queriesPresent+ , bench "Data.HashMap" $ nf (map (`HM.lookup` hashMap)) queriesPresent+ , bench "BasicHashTable" $ nfIO $ traverse (HT.lookup basic) queriesPresent+ -- , bench "LinearHashTable" $ nfIO $ traverse (HT.lookup linear) queriesPresent+ , bench "CuckooHashTable" $ nfIO $ traverse (HT.lookup cuckoo) queriesPresent+ ]+ , bgroup "missing"+ [ bench "Data.RadixTree" $ nf (map (`RT.lookup` radixTree)) queriesMissing+ , bench "Data.Map" $ nf (map (`M.lookup` treeMap)) queriesMissing+ , bench "Data.HashMap" $ nf (map (`HM.lookup` hashMap)) queriesMissing+ , bench "BasicHashTable" $ nfIO $ traverse (HT.lookup basic) queriesMissing+ -- , bench "LinearHashTable" $ nfIO $ traverse (HT.lookup linear) queriesMissing+ , bench "CuckooHashTable" $ nfIO $ traverse (HT.lookup cuckoo) queriesMissing+ ]+ , bgroup "both"+ [ bench "Data.RadixTree" $ nf (map (`RT.lookup` radixTree)) queriesBoth+ , bench "Data.Map" $ nf (map (`M.lookup` treeMap)) queriesBoth+ , bench "Data.HashMap" $ nf (map (`HM.lookup` hashMap)) queriesBoth+ , bench "BasicHashTable" $ nfIO $ traverse (HT.lookup basic) queriesBoth+ -- , bench "LinearHashTable" $ nfIO $ traverse (HT.lookup linear) queriesBoth+ , bench "CuckooHashTable" $ nfIO $ traverse (HT.lookup cuckoo) queriesBoth+ ]+ ]+ , bgroup "keys"+ [ bench "Data.RadixTree" $ nf RT.keys radixTree+ , bench "Data.Map" $ nf M.keys treeMap+ , bench "Data.HashMap" $ nf HM.keys hashMap+ ]+ , bgroup "toList"+ [ bench "Data.RadixTree" $ nf RT.toList radixTree+ , bench "Data.Map" $ nf M.toList treeMap+ , bench "Data.HashMap" $ nf HM.toList hashMap+ ]+ , bgroup "union"+ [ bench "Data.RadixTree" $ nf (uncurry RT.union) (radixTree, radixTreeRev)+ , bench "Data.Map" $ nf (uncurry M.union) (treeMap, treeMapRev)+ , bench "Data.HashMap" $ nf (uncurry HM.union) (hashMap, hashMapRev)+ ]+ ]
+ radix-tree.cabal view
@@ -0,0 +1,143 @@+name:+ radix-tree+version:+ 0.1+category:+ Data Structures+synopsis:+ Radix tree data structive over short byte-strings+description:+ This module provides a memory-efficient map from+ Data.ByteString.Short keys to arbitrary values implemented as a radix+ tree datastructure. Memory efficiency is achieved by sharing common+ prefixes of all keys.+license:+ BSD3+license-file:+ LICENSE+author:+ Sergey Vinokurov+maintainer:+ Sergey Vinokurov <serg.foo@gmail.com>+copyright:+ (c) 2018 Sergey Vinokurov+tested-with:+ GHC == 8.0.2, GHC == 8.2.2, GHC == 8.4.3++cabal-version:+ 2.0+build-type:+ Simple++homepage: https://github.com/sergv/radix-tree++source-repository head+ type: git+ location: https://github.com/sergv/radix-tree.git++library+ exposed-modules:+ Data.RadixTree+ Data.RadixTree.Internal+ hs-source-dirs:+ src+ build-depends:+ base >= 4.9 && < 5,+ bytestring,+ containers,+ deepseq,+ primitive+ default-language:+ Haskell2010+ ghc-options:+ -Wall+ -fwarn-name-shadowing+ -fno-warn-type-defaults+ if impl(ghc >= 8.0)+ ghc-options:+ -Wcompat+ -Whi-shadowing+ -Widentities+ -Wincomplete-record-updates+ -Wincomplete-uni-patterns+ -Wmissing-exported-signatures+ if impl(ghc >= 8.2)+ ghc-options:+ -Wcpp-undef+ -Wmissing-home-modules+ -Wunbanged-strict-patterns++test-suite radix-tree-test+ type:+ exitcode-stdio-1.0+ main-is:+ test/TestMain.hs+ build-depends:+ HUnit,+ QuickCheck,+ base >= 4.9 && < 5,+ bytestring,+ containers,+ tasty,+ tasty-hunit,+ tasty-quickcheck,+ radix-tree+ default-language:+ Haskell2010+ ghc-options:+ -rtsopts+ -Wall+ -fwarn-name-shadowing+ -fno-warn-type-defaults+ if impl(ghc >= 8.0)+ ghc-options:+ -Wall-missed-specialisations+ -Wcompat+ -Whi-shadowing+ -Widentities+ -Wincomplete-record-updates+ -Wincomplete-uni-patterns+ -Wmissing-exported-signatures+ if impl(ghc >= 8.2)+ ghc-options:+ -Wcpp-undef+ -Wmissing-home-modules+ -Wunbanged-strict-patterns++benchmark radix-tree-bench+ type:+ exitcode-stdio-1.0+ main-is:+ bench/RadixTreeBench.hs+ hs-source-dirs:+ . bench+ build-depends:+ base >= 4.9 && < 5,+ bytestring,+ containers,+ deepseq,+ gauge >= 0.2.3,+ hashtables,+ radix-tree,+ text,+ unordered-containers+ default-language:+ Haskell2010+ ghc-options:+ -rtsopts+ -Wall+ -fwarn-name-shadowing+ -fno-warn-type-defaults+ if impl(ghc >= 8.0)+ ghc-options:+ -Wcompat+ -Whi-shadowing+ -Widentities+ -Wincomplete-record-updates+ -Wincomplete-uni-patterns+ -Wmissing-exported-signatures+ if impl(ghc >= 8.2)+ ghc-options:+ -Wcpp-undef+ -Wmissing-home-modules+ -Wunbanged-strict-patterns
+ src/Data/RadixTree.hs view
@@ -0,0 +1,33 @@+----------------------------------------------------------------------------+-- |+-- Module : Data.RadixTree+-- Copyright : (c) Sergey Vinokurov 2018+-- License : BSD3-style (see LICENSE)+-- Maintainer : serg.foo@gmail.com+--+-- This is an implementation of the radix tree datastructure. Interface+-- is designed to be compatible with what 'Data.Map' provides.+----------------------------------------------------------------------------++module Data.RadixTree+ ( RadixTree+ , empty+ , null+ , size+ , insert+ , insertWith+ , lookup+ , fromList+ , toList+ , toAscList+ , keys+ , keysSet+ , elems+ , mapMaybe+ , union+ , unionWith+ ) where++import Prelude hiding (lookup, null)++import Data.RadixTree.Internal
+ src/Data/RadixTree/Internal.hs view
@@ -0,0 +1,454 @@+----------------------------------------------------------------------------+-- |+-- Module : Data.RadixTree.Internal+-- Copyright : (c) Sergey Vinokurov 2018+-- License : BSD3-style (see LICENSE)+-- Maintainer : serg.foo@gmail.com+--+-- This is an internal module that exposes innards of the 'RadixTree'+-- data structure. This API may change in any new release, even in a+-- patch release - depend on it at your own risk.+----------------------------------------------------------------------------++{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DeriveFoldable #-}+{-# LANGUAGE DeriveFunctor #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE DeriveTraversable #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE ScopedTypeVariables #-}++{-# OPTIONS_HADDOCK not-home #-}++module Data.RadixTree.Internal+ ( RadixTree(..)+ , empty+ , null+ , size+ , insert+ , insertWith+ , lookup+ , fromList+ , toList+ , toAscList+ , keys+ , keysSet+ , elems+ , mapMaybe+ , union+ , unionWith+ ) where++import Prelude hiding (lookup, null)++import Control.Arrow (first)+import Control.DeepSeq+import Control.Monad.ST+import Control.Monad.ST.Unsafe++import Data.ByteString.Short (ShortByteString)+import qualified Data.ByteString.Short as BSS+import qualified Data.ByteString.Short.Internal as BSSI+import qualified Data.Foldable as Foldable+import Data.IntMap (IntMap)+import qualified Data.IntMap.Strict as IM+import qualified Data.List as L+import Data.Maybe (fromMaybe)+import Data.Primitive.ByteArray+import Data.Semigroup as Semigroup+import Data.Set (Set)+import qualified Data.Set as S+import Data.Word+import GHC.Generics (Generic)++-- | A tree data structure that efficiently indexes values by string keys.+--+-- This type can be more memory-efficient than 'Data.Map' because it combines+-- common prefixes of all keys. Specific savings will vary depending on+-- concrete data set.+data RadixTree a+ = RadixNode+ !(Maybe a)+ !(IntMap (RadixTree a)) -- ^ Either has 0 or 2 or more children, never 1.+ | RadixStr+ !(Maybe a)+ {-# UNPACK #-} !ShortByteString -- ^ Non-empty+ !(RadixTree a)+ deriving (Show, Functor, Foldable, Traversable, Generic)++instance NFData a => NFData (RadixTree a)++-- | Radix tree with no elements.+empty :: RadixTree a+empty = RadixNode Nothing IM.empty++{-# INLINE interleaveST #-}+interleaveST :: ST s a -> ST s a+interleaveST =+#if MIN_VERSION_base(4, 10, 0)+ unsafeDupableInterleaveST+#else+ unsafeInterleaveST+#endif++splitShortByteString :: Int -> ShortByteString -> (ShortByteString, ShortByteString, Word8, ShortByteString)+splitShortByteString n (BSSI.SBS source) = runST $ do+ prefix <- newByteArray prefixSize+ copyByteArray prefix 0 source' 0 prefixSize+ ByteArray prefix# <- unsafeFreezeByteArray prefix+ midSuffix <- interleaveST $ do+ midSuffix <- newByteArray midSuffixSize+ copyByteArray midSuffix 0 source' n midSuffixSize+ unsafeFreezeByteArray midSuffix+ suffix <- interleaveST $ do+ suffix <- newByteArray suffixSize+ copyByteArray suffix 0 source' (n + 1) suffixSize+ unsafeFreezeByteArray suffix+ pure (BSSI.SBS prefix#, byteArrayToBSS midSuffix, indexByteArray source' n, byteArrayToBSS suffix)+ where+ source' = ByteArray source+ prefixSize = n+ midSuffixSize = sizeofByteArray source' - prefixSize+ suffixSize = midSuffixSize - 1++{-# INLINE byteArrayToBSS #-}+byteArrayToBSS :: ByteArray -> BSS.ShortByteString+byteArrayToBSS (ByteArray xs) = BSSI.SBS xs++dropShortByteString :: Int -> ShortByteString -> ShortByteString+dropShortByteString 0 src = src+dropShortByteString !n (BSSI.SBS source) = runST $ do+ dest <- newByteArray sz+ copyByteArray dest 0 source' n sz+ byteArrayToBSS <$> unsafeFreezeByteArray dest+ where+ source' = ByteArray source+ !sz = sizeofByteArray source' - n++singletonShortByteString :: Word8 -> ShortByteString+singletonShortByteString !c = runST $ do+ dest <- newByteArray 1+ writeByteArray dest 0 c+ byteArrayToBSS <$> unsafeFreezeByteArray dest++{-# INLINE unsafeHeadeShortByteString #-}+unsafeHeadeShortByteString :: ShortByteString -> Word8+unsafeHeadeShortByteString = (`BSSI.unsafeIndex` 0)++data Mismatch+ = IsPrefix+ | CommonPrefixThenMismatch+ !ShortByteString -- ^ Prefix of node contents common with the key+ ShortByteString -- ^ Suffix with the first mismatching byte+ Word8 -- ^ First byte of the suffix that caused mismatch+ ShortByteString -- ^ Rest of node contents, suffix+ deriving (Show, Generic)++analyseMismatch+ :: ShortByteString -- ^ Key+ -> Int -- ^ Key offset+ -> ShortByteString -- ^ Node contents+ -> Mismatch+analyseMismatch (BSSI.SBS key) !keyOffset nodeContentsBS@(BSSI.SBS nodeContents) =+ case findMismatch 0 of+ Nothing -> IsPrefix+ Just mismatchIdx ->+ case splitShortByteString mismatchIdx nodeContentsBS of+ (prefix, midSuffix, mid, suffix) -> CommonPrefixThenMismatch prefix midSuffix mid suffix+ where+ keySize = sizeofByteArray key'+ keyLeft = keySize - keyOffset+ contentsSize = sizeofByteArray nodeContents'++ key' = ByteArray key+ nodeContents' = ByteArray nodeContents++ limit :: Int+ limit = min keyLeft contentsSize++ findMismatch :: Int -> Maybe Int+ findMismatch !i+ | i == limit+ = if i == contentsSize+ then Nothing+ else Just i -- Key ended in the middle of node's packed key.+ | (indexByteArray key' (keyOffset + i) :: Word8) == indexByteArray nodeContents' i+ = findMismatch $ i + 1+ | otherwise+ = Just i++mkRadixNodeFuse :: Maybe a -> IntMap (RadixTree a) -> Maybe (RadixTree a)+mkRadixNodeFuse val children =+ case val of+ Nothing | IM.null children+ -> Nothing+ val' | [(c, child)] <- IM.toList children+ -> Just $ RadixStr val' (singletonShortByteString $ fromIntegral c) child+ _ -> Just $ RadixNode val children++-- Precondition: input string is non-empty+mkRadixStrFuse :: Maybe a -> ShortByteString -> RadixTree a -> Maybe (RadixTree a)+mkRadixStrFuse val str rest =+ case (val, rest) of+ (val', RadixStr Nothing str' rest') ->+ Just $ RadixStr val' (str Semigroup.<> str') rest'+ (Nothing, node)+ | null node -> Nothing+ (val', rest') ->+ Just $ RadixStr val' str rest'++mkRadixStr :: ShortByteString -> RadixTree a -> RadixTree a+mkRadixStr str rest+ | BSS.null str = rest+ | otherwise = RadixStr Nothing str rest++-- TODO: prove following function correct.++-- | Check whether radix tree is empty+null :: RadixTree a -> Bool+null = \case+ RadixNode Nothing children -> IM.null children+ RadixStr Nothing _ rest -> null rest+ _ -> False++-- | O(n) Get number of elements in a radix tree.+size :: RadixTree a -> Int+size = length++-- | Add new element to a radix tree.+insert :: forall a. ShortByteString -> a -> RadixTree a -> RadixTree a+insert = insertWith const++-- | Add new element to a radix tree. If an element was already present for+-- the given key, use supplied funciton @f@ to produce a new value. The+-- function will be called like this @f newValue oldValue@.+insertWith :: forall a. (a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a+insertWith = insert'++{-# INLINE insert' #-}+insert' :: forall a. (a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a+insert' f key value = go 0+ where+ len = BSS.length key++ readKey :: Int -> Int+ readKey = fromIntegral . BSSI.unsafeIndex key++ go :: Int -> RadixTree a -> RadixTree a+ go i+ | i < len+ = \case+ RadixNode oldValue children+ | IM.null children ->+ RadixStr oldValue (dropShortByteString i key) $ RadixNode (Just value) IM.empty+ | otherwise ->+ RadixNode oldValue $+ IM.alter (Just . maybe optNode (go i')) c children+ where+ c :: Int+ c = readKey i+ i' = i + 1+ optNode =+ mkRadixStr (dropShortByteString i' key) $ RadixNode (Just value) IM.empty+ RadixStr oldValue packedKey rest ->+ case analyseMismatch key i packedKey of+ IsPrefix ->+ RadixStr oldValue packedKey $ go (i + BSS.length packedKey) rest+ CommonPrefixThenMismatch prefix midSuffix mid suffix ->+ (if BSS.null prefix then id else RadixStr oldValue prefix) $+ if isKeyEnded+ then+ RadixStr (Just value) midSuffix rest+ else+ RadixNode (if BSS.null prefix then oldValue else Nothing) $+ IM.fromList+ [ ( mid'+ , mkRadixStr suffix rest+ )+ , ( readKey i'+ , mkRadixStr (dropShortByteString (i' + 1) key) $ RadixNode (Just value) IM.empty+ )+ ]+ where+ i' = i + BSS.length prefix+ isKeyEnded = i' >= len+ mid' = fromIntegral mid+ | otherwise+ = \case+ RadixNode oldValue children ->+ RadixNode (Just (maybe value (f value) oldValue)) children+ RadixStr oldValue key' rest ->+ RadixStr (Just (maybe value (f value) oldValue)) key' rest++canStripPrefixFromShortByteString+ :: Int -> ShortByteString -> ShortByteString -> Bool+canStripPrefixFromShortByteString bigStart (BSSI.SBS small) (BSSI.SBS big)+ | bigStart + smallSize > bigSize = False+ | otherwise = findMismatch 0+ where+ small' = ByteArray small+ big' = ByteArray big++ smallSize = sizeofByteArray small'+ bigSize = sizeofByteArray big'++ findMismatch :: Int -> Bool+ findMismatch !i+ | i == smallSize+ = True+ | (indexByteArray small' i :: Word8) == indexByteArray big' (bigStart + i)+ = findMismatch $ i + 1+ | otherwise+ = False++-- | O(length(key)) Try to find a value associated with the given key.+lookup :: forall a. ShortByteString -> RadixTree a -> Maybe a+lookup key = go 0+ where+ len = BSS.length key++ readKey :: Int -> Int+ readKey = fromIntegral . BSSI.unsafeIndex key++ go :: Int -> RadixTree a -> Maybe a+ go !n tree+ | n == len+ = case tree of+ RadixNode val _ -> val+ RadixStr val _ _ -> val+ | otherwise+ = case tree of+ RadixNode _ children ->+ IM.lookup (readKey n) children >>= go (n + 1)+ RadixStr _ packedKey rest+ | canStripPrefixFromShortByteString n packedKey key+ -> go (n + BSS.length packedKey) rest+ | otherwise+ -> Nothing++-- | Construct a radix tree from list of key-value pairs. If some key+-- appears twice in the input list, later occurrences will override+-- earlier ones.+fromList :: [(ShortByteString, a)] -> RadixTree a+fromList =+ L.foldl' (\acc (k, v) -> insert' const k v acc) empty++-- | O(n) Convert a radix tree to a list of key-value pairs.+toList :: RadixTree a -> [(ShortByteString, a)]+toList = toAscList++-- | O(n) Convert a radix tree to an ascending list of key-value pairs.+toAscList :: forall a. RadixTree a -> [(ShortByteString, a)]+toAscList = map (first BSS.pack) . go+ where+ go :: RadixTree a -> [([Word8], a)]+ go = \case+ RadixNode val children ->+ maybe id (\val' ys -> ([], val') : ys) val $+ IM.foldMapWithKey (\c child -> map (first (fromIntegral c :)) $ go child) children+ RadixStr val packedKey rest ->+ maybe id (\val' ys -> ([], val') : ys) val $+ map (first (BSS.unpack packedKey ++)) $+ go rest++-- | O(n) Get all keys stored in a radix tree.+keys :: RadixTree a -> [ShortByteString]+keys = map BSS.pack . go+ where+ go :: RadixTree a -> [[Word8]]+ go = \case+ RadixNode val children ->+ maybe id (\_ ys -> [] : ys) val $+ IM.foldMapWithKey (\c child -> map (fromIntegral c :) $ go child) children+ RadixStr val packedKey rest ->+ maybe id (\_ ys -> [] : ys) val $+ map (BSS.unpack packedKey <>) $+ go rest++-- | O(n) Get set of all keys stored in a radix tree.+keysSet :: RadixTree a -> Set ShortByteString+keysSet = S.fromDistinctAscList . keys++-- | O(n) Get all values stored in a radix tree.+elems :: RadixTree a -> [a]+elems = Foldable.toList++-- | O(n) Map a function that can remove some existing elements over a+-- radix tree.+mapMaybe :: forall a b. (a -> Maybe b) -> RadixTree a -> RadixTree b+mapMaybe f = fromMaybe empty . go+ where+ go :: RadixTree a -> Maybe (RadixTree b)+ go = \case+ RadixNode val children ->+ mkRadixNodeFuse (f =<< val) $ IM.mapMaybe go children+ RadixStr val str rest ->+ mkRadixStrFuse (f =<< val) str $ fromMaybe empty $ go rest++-- | O(n + m) Combine two radix trees trees. If a key is present in both+-- trees then the value from left one will be retained.+union :: RadixTree a -> RadixTree a -> RadixTree a+union = unionWith const++-- | O(n + m) Combine two trees using supplied function to resolve+-- values that have the same key in both trees.+unionWith :: forall a. (a -> a -> a) -> RadixTree a -> RadixTree a -> RadixTree a+unionWith f = go+ where+ combineVals :: Maybe a -> Maybe a -> Maybe a+ combineVals x y = case (x, y) of+ (Nothing, Nothing) -> Nothing+ (Nothing, y'@Just{}) -> y'+ (x'@Just{}, Nothing) -> x'+ (Just x', Just y') -> Just $ f x' y'++ go :: RadixTree a -> RadixTree a -> RadixTree a+ go x y = case (x, y) of+ (RadixNode val children, RadixNode val' children') ->+ RadixNode (combineVals val val') (IM.unionWith go children children')+ (RadixNode val children, RadixStr val' str' rest') ->+ RadixNode (combineVals val val') $+ (\g -> IM.alter g h children) $ \child ->+ Just $!+ let rest'' = mkRadixStr (dropShortByteString 1 str') rest' in+ case child of+ Nothing -> rest''+ Just child' -> go child' rest''+ where+ h = fromIntegral $ unsafeHeadeShortByteString str'+ (RadixStr val str rest, RadixNode val' children') ->+ RadixNode (combineVals val val') $+ (\g -> IM.alter g h children') $ \child ->+ Just $!+ let rest' = mkRadixStr (dropShortByteString 1 str) rest in+ case child of+ Nothing -> rest'+ Just child' -> go rest' child'+ where+ h = fromIntegral $ unsafeHeadeShortByteString str+ (RadixStr val str rest, RadixStr val' str' rest') ->+ case analyseMismatch str 0 str' of+ -- str' is a prefix of str+ IsPrefix ->+ RadixStr (combineVals val val') str' $+ go (mkRadixStr (dropShortByteString (BSS.length str') str) rest) rest'+ -- str' = prefix + firstMismatchStr' + suffixStr'+ -- = prefix + midSuffixStr'+ CommonPrefixThenMismatch prefix midSuffixStr' firstMismatchStr' suffixStr' ->+ (if BSS.null prefix then id else RadixStr (combineVals val val') prefix) $+ if BSS.length prefix == BSS.length str+ then+ go rest $ RadixStr+ (if BSS.null prefix then combineVals val val' else Nothing)+ midSuffixStr'+ rest'+ else RadixNode (if BSS.null prefix then combineVals val val' else Nothing) $ IM.fromList+ [ ( fromIntegral firstMismatchStr'+ , mkRadixStr suffixStr' rest'+ )+ , ( fromIntegral $ BSSI.unsafeIndex str $ BSS.length prefix+ , mkRadixStr (dropShortByteString (BSSI.length prefix + 1) str) rest+ )+ ]
+ test/TestMain.hs view
@@ -0,0 +1,119 @@+----------------------------------------------------------------------------+-- |+-- Module : TestMain+-- Copyright : (c) Sergey Vinokurov 2018+-- License : BSD3-style (see LICENSE)+-- Maintainer : serg.foo@gmail.com+----------------------------------------------------------------------------++{-# LANGUAGE ScopedTypeVariables #-}++{-# OPTIONS_GHC -Wno-orphans #-}++module Main (main) where++import Data.ByteString.Short (ShortByteString)+import qualified Data.ByteString.Short as BSS++import Data.Char+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as M+import Data.RadixTree (RadixTree)+import qualified Data.RadixTree as RT+import Data.Word++import Test.QuickCheck+import Test.QuickCheck.Poly+import Test.Tasty+import Test.Tasty.QuickCheck as QC++newtype AsciiChar = AsciiChar { unAsciiChar :: Char }++instance Arbitrary AsciiChar where+ arbitrary = AsciiChar <$> choose ('a', 'z')+ shrink (AsciiChar 'a') = []+ shrink (AsciiChar c) = [AsciiChar c' | c' <- ['a'..pred c]]++mkAsciiChar :: Word8 -> AsciiChar+mkAsciiChar = AsciiChar . chr. fromIntegral++asciiByte :: AsciiChar -> Word8+asciiByte = fromIntegral . ord . unAsciiChar++instance Arbitrary ShortByteString where+ arbitrary =+ BSS.pack . map asciiByte <$> listOf arbitrary+ shrink =+ map (BSS.pack . map asciiByte) . shrink . map mkAsciiChar . BSS.unpack++instance Arbitrary a => Arbitrary (RadixTree a) where+ arbitrary = RT.fromList <$> arbitrary+ shrink = map RT.fromList . shrink . RT.toAscList++main :: IO ()+main = defaultMain tests++tests :: TestTree+tests = testGroup "Tests" [properties]++properties :: TestTree+properties = testGroup "Properties" [qcProps]++qcProps :: TestTree+qcProps = adjustOption (\(QuickCheckTests n) -> QuickCheckTests (max 10000 n)) $ testGroup "radix tree"+ [ QC.testProperty "∀ t: RT.lookup k (RT.insert k v t) == v" $+ \(t :: RadixTree A) (k :: ShortByteString) (v :: A) ->+ RT.lookup k (RT.insert k v t) == Just v+ , QC.testProperty "∀ t: RT.lookup k (RT.insert k v2 (RT.insert k v1 t)) == v2" $+ \(t :: RadixTree A) (k :: ShortByteString) (v1 :: A) (v2 :: A) ->+ RT.lookup k (RT.insert k v2 (RT.insert k v1 t)) == Just v2++ , QC.testProperty "∀ xs: RT.fromList xs == M.fromList xs" $+ \(xs :: [(ShortByteString, A)]) ->+ RT.toAscList (RT.fromList xs) == M.toAscList (M.fromList xs)++ , QC.testProperty "∀ xs: RT.size (RT.fromList xs) == M.size (M.fromList xs)" $+ \(xs :: [(ShortByteString, A)]) ->+ RT.size (RT.fromList xs) == M.size (M.fromList xs)++ , QC.testProperty "∀ f: RT.mapMaybe f == M.mapMaybe f" $+ \(f :: Fun A (Maybe B)) ->+ RT.mapMaybe (applyFun f) ==== M.mapMaybe (applyFun f)++ , QC.testProperty "∀ k v t: RT.insert k v t == M.insert k v t" $+ \(k :: ShortByteString) (v :: A) ->+ RT.insert k v ==== M.insert k v++ , QC.testProperty "∀ f xs ys: RT.mergeWith f xs ys == M.mergeWith f xs ys" $+ \(f :: Fun (A, A) A) ->+ RT.unionWith (curry (applyFun f)) ===== M.unionWith (curry (applyFun f))+ ]++(====)+ :: Eq b+ => (RadixTree a -> RadixTree b)+ -> (Map ShortByteString a -> Map ShortByteString b)+ -> [(ShortByteString, a)]+ -> Bool+(====) f g xs =+ RT.toAscList (f (RT.fromList xs)) == M.toAscList (g (M.fromList xs))++(=====)+ :: Eq a+ => (RadixTree a -> RadixTree a -> RadixTree a)+ -> (Map ShortByteString a -> Map ShortByteString a -> Map ShortByteString a)+ -> [(ShortByteString, a)]+ -> [(ShortByteString, a)]+ -> Bool+(=====) f g xs ys =+ RT.toAscList (f (RT.fromList xs) (RT.fromList ys)) == M.toAscList (g (M.fromList xs) (M.fromList ys))++-- unitTests :: TestTree+-- unitTests = testGroup "Unit tests"+-- [ testCase "List comparison (different length)" $+-- [1, 2, 3] `compare` [1,2] @?= GT+--+-- -- the following test does not hold+-- , testCase "List comparison (same length)" $+-- [1, 2, 3] `compare` [1,2,2] @?= LT+-- ]