bytetrie (empty) → 0.1.0.0
raw patch · 7 files changed
+1008/−0 lines, 7 filesdep +basedep +bifunctorsdep +byte-containerssetup-changed
Dependencies added: base, bifunctors, byte-containers, bytebuild, byteslice, bytetrie, primitive, primitive-unlifted, quickcheck-classes, tasty, tasty-hunit, tasty-quickcheck, unpacked-maybe
Files
- CHANGELOG.md +5/−0
- LICENSE +30/−0
- Setup.hs +3/−0
- bytetrie.cabal +59/−0
- common/TestData.hs +232/−0
- src/Data/Trie/Word8.hs +534/−0
- test/Main.hs +145/−0
+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Revision history for bytetrie++## 0.1.0.0 -- YYYY-mm-dd++* First version. Released on an unsuspecting world.
+ LICENSE view
@@ -0,0 +1,30 @@+Copyright (c) 2020, Eric Demko++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 Eric Demko 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.
+ Setup.hs view
@@ -0,0 +1,3 @@+import Distribution.Simple++main = defaultMain
+ bytetrie.cabal view
@@ -0,0 +1,59 @@+cabal-version: 3.0+name: bytetrie+version: 0.1.0.0+author: Eric Demko+maintainer: amartin@layer3com.com+copyright: 2020 Layer 3 Communications+category: Data+build-type: Simple+synopsis: Tries with Bytes as keys.+description:+ Tries with 'Bytes' (equiv. 'ByteArray') as keys. This implementation is+ optimized for performing queries rather than updating the structure+ repeatedly.++homepage: https://github.com/byteverse/bytetrie+bug-reports: https://github.com/byteverse/bytetrie/issues+license: BSD-3-Clause+license-file: LICENSE+extra-doc-files: CHANGELOG.md++library+ exposed-modules: Data.Trie.Word8+ build-depends:+ , base >=4.12 && <5+ , bifunctors >=5.6.1 && <5.7+ , byte-containers >=0.1 && <0.2.0+ , bytebuild >=0.3.16 && <0.4+ , byteslice >=0.2.13 && <0.3+ , primitive >=0.7 && <0.10+ , unpacked-maybe >=0.1.0 && <0.2++ hs-source-dirs: src+ default-language: Haskell2010+ ghc-options: -Wall -O2++test-suite test+ main-is: Main.hs+ type: exitcode-stdio-1.0+ build-depends:+ , base >=4.12 && <5+ , byte-containers >=0.1 && <0.2.0+ , byteslice >=0.2.13 && <0.3+ , bytetrie+ , primitive >=0.7 && <0.10+ , primitive-unlifted >=0.1.2 && <0.2+ , quickcheck-classes >=0.6.4+ , tasty+ , tasty-hunit >=0.10+ , tasty-quickcheck+ , unpacked-maybe >=0.1.0 && <0.2++ other-modules: TestData+ hs-source-dirs: test common+ default-language: Haskell2010+ ghc-options: -Wall -O2++source-repository head+ type: git+ location: git://github.com/byteverse/bytetrie.git
+ common/TestData.hs view
@@ -0,0 +1,232 @@+module TestData+ ( replacements+ , bigstring+ ) where++import Data.Bifunctor (bimap)+import Data.Bytes (Bytes)+import Data.Char (ord)+import Data.Trie.Word8 (Trie)+import Data.Word (Word8)++import qualified Data.Bytes.Text.Latin1 as Latin1+import qualified Data.Trie.Word8 as Trie++_c2w :: Char -> Word8+_c2w = fromIntegral . ord++replacements :: Trie Bytes+replacements =+ Trie.fromList $+ map+ (bimap Latin1.fromString Latin1.fromString)+ [ ("Francisco", "Frank")+ , ("Bernardo", "Bob")+ , ("Marcellus", "Mark")+ , ("Horatio", "Henry")+ , ("Ghost", "Spock")+ , ("What", "wat")+ , ("Why", "ytho")+ ]++bigstring :: Bytes+bigstring =+ Latin1.fromString . unlines $+ [ "Enter two Sentinels-[first,] Francisco, [who paces up and down at his post; then] Bernardo, [who approaches him]."+ , "Bernardo. Who's there?"+ , "Francisco. Nay, answer me. Stand and unfold yourself."+ , "Bernardo. Long live the King!"+ , "Francisco. Bernardo?"+ , "Bernardo. He."+ , "Francisco. You come most carefully upon your hour."+ , "Bernardo. 'Tis now struck twelve. Get thee to bed, Francisco."+ , "Francisco. For this relief much thanks. 'Tis bitter cold,"+ , "And I am sick at heart."+ , "Bernardo. Have you had quiet guard?"+ , "If you do meet Horatio and Marcellus,"+ , "The rivals of my watch, bid them make haste."+ , "Enter Horatio and Marcellus."+ , "Francisco. I think I hear them. Stand, ho! Who is there?"+ , "Horatio. Friends to this ground."+ , "Marcellus. And liegemen to the Dane."+ , "Francisco. Give you good night."+ , "Marcellus. O, farewell, honest soldier."+ , "Who hath reliev'd you?"+ , "Francisco. Bernardo hath my place."+ , "Give you good night. Exit."+ , "Marcellus. Holla, Bernardo!"+ , "Bernardo. Say-"+ , "What, is Horatio there?"+ , "Horatio. A piece of him."+ , "Bernardo. Welcome, Horatio. Welcome, good Marcellus."+ , "Marcellus. What, has this thing appear'd again to-night?"+ , "Bernardo. I have seen nothing."+ , "Marcellus. Horatio says 'tis but our fantasy,"+ , "And will not let belief take hold of him"+ , "Touching this dreaded sight, twice seen of us."+ , "Therefore I have entreated him along,"+ , "With us to watch the minutes of this night,"+ , "That, if again this apparition come,"+ , "He may approve our eyes and speak to it."+ , "Horatio. Tush, tush, 'twill not appear."+ , "Bernardo. Sit down awhile,"+ , "And let us once again assail your ears,"+ , "That are so fortified against our story,"+ , "What we two nights have seen."+ , "Horatio. Well, sit we down,"+ , "And let us hear Bernardo speak of this."+ , "Bernardo. Last night of all,"+ , "When yond same star that's westward from the pole"+ , "Had made his course t' illume that part of heaven"+ , "Where now it burns, Marcellus and myself,"+ , "The bell then beating one-"+ , "Enter Ghost."+ , "Marcellus. Peace! break thee off! Look where it comes again!"+ , "Bernardo. In the same figure, like the King that's dead."+ , "Marcellus. Thou art a scholar; speak to it, Horatio."+ , "Bernardo. Looks it not like the King? Mark it, Horatio."+ , "Horatio. Most like. It harrows me with fear and wonder."+ , "Bernardo. It would be spoke to."+ , "Marcellus. Question it, Horatio."+ , "Horatio. What art thou that usurp'st this time of night"+ , "Together with that fair and warlike form"+ , "In which the majesty of buried Denmark"+ , "Did sometimes march? By heaven I charge thee speak!"+ , "Marcellus. It is offended."+ , "Bernardo. See, it stalks away!"+ , "Horatio. Stay! Speak, speak! I charge thee speak!"+ , "Exit Ghost."+ , "Marcellus. 'Tis gone and will not answer."+ , "Bernardo. How now, Horatio? You tremble and look pale."+ , "Is not this something more than fantasy?"+ , "What think you on't?"+ , "Horatio. Before my God, I might not this believe"+ , "Without the sensible and true avouch"+ , "Of mine own eyes."+ , "Marcellus. Is it not like the King?"+ , "Horatio. As thou art to thyself."+ , "Such was the very armour he had on"+ , "When he th' ambitious Norway combated."+ , "So frown'd he once when, in an angry parle,"+ , "He smote the sledded Polacks on the ice."+ , "'Tis strange."+ , "Marcellus. Thus twice before, and jump at this dead hour,"+ , "With martial stalk hath he gone by our watch."+ , "Horatio. In what particular thought to work I know not;"+ , "But, in the gross and scope of my opinion,"+ , "This bodes some strange eruption to our state."+ , "Marcellus. Good now, sit down, and tell me he that knows,"+ , "Why this same strict and most observant watch"+ , "So nightly toils the subject of the land,"+ , "And why such daily cast of brazen cannon"+ , "And foreign mart for implements of war;"+ , "Why such impress of shipwrights, whose sore task"+ , "Does not divide the Sunday from the week."+ , "What might be toward, that this sweaty haste"+ , "Doth make the night joint-labourer with the day?"+ , "Who is't that can inform me?"+ , "Horatio. That can I."+ , "At least, the whisper goes so. Our last king,"+ , "Whose image even but now appear'd to us,"+ , "Was, as you know, by Fortinbras of Norway,"+ , "Thereto prick'd on by a most emulate pride,"+ , "Dar'd to the combat; in which our valiant Hamlet"+ , "(For so this side of our known world esteem'd him)"+ , "Did slay this Fortinbras; who, by a seal'd compact,"+ , "Well ratified by law and heraldry,"+ , "Did forfeit, with his life, all those his lands"+ , "Which he stood seiz'd of, to the conqueror;"+ , "Against the which a moiety competent"+ , "Was gaged by our king; which had return'd"+ , "To the inheritance of Fortinbras,"+ , "Had he been vanquisher, as, by the same cov'nant"+ , "And carriage of the article design'd,"+ , "His fell to Hamlet. Now, sir, young Fortinbras,"+ , "Of unimproved mettle hot and full,"+ , "Hath in the skirts of Norway, here and there,"+ , "Shark'd up a list of lawless resolutes,"+ , "For food and diet, to some enterprise"+ , "That hath a stomach in't; which is no other,"+ , "As it doth well appear unto our state,"+ , "But to recover of us, by strong hand"+ , "And terms compulsatory, those foresaid lands"+ , "So by his father lost; and this, I take it,"+ , "Is the main motive of our preparations,"+ , "The source of this our watch, and the chief head"+ , "Of this post-haste and romage in the land."+ , "Bernardo. I think it be no other but e'en so."+ , "Well may it sort that this portentous figure"+ , "Comes armed through our watch, so like the King"+ , "That was and is the question of these wars."+ , "Horatio. A mote it is to trouble the mind's eye."+ , "In the most high and palmy state of Rome,"+ , "A little ere the mightiest Julius fell,"+ , "The graves stood tenantless, and the sheeted dead"+ , "Did squeak and gibber in the Roman streets;"+ , "As stars with trains of fire, and dews of blood,"+ , "Disasters in the sun; and the moist star"+ , "Upon whose influence Neptune's empire stands"+ , "Was sick almost to doomsday with eclipse."+ , "And even the like precurse of fierce events,"+ , "As harbingers preceding still the fates"+ , "And prologue to the omen coming on,"+ , "Have heaven and earth together demonstrated"+ , "Unto our climature and countrymen."+ , "[Enter Ghost again.]"+ , "But soft! behold! Lo, where it comes again!"+ , "I'll cross it, though it blast me.- Stay illusion!"+ , "[Spreads his arms.]"+ , "If thou hast any sound, or use of voice,"+ , "Speak to me."+ , "If there be any good thing to be done,"+ , "That may to thee do ease, and, grace to me,"+ , "Speak to me."+ , "If thou art privy to thy country's fate,"+ , "Which happily foreknowing may avoid,"+ , "O, speak!"+ , "Or if thou hast uphoarded in thy life"+ , "Extorted treasure in the womb of earth"+ , "(For which, they say, you spirits oft walk in death),"+ , "[The cock crows.]"+ , "Speak of it! Stay, and speak!- Stop it, Marcellus!"+ , "Marcellus. Shall I strike at it with my partisan?"+ , "Horatio. Do, if it will not stand."+ , "Bernardo. 'Tis here!"+ , "Horatio. 'Tis here!"+ , "Marcellus. 'Tis gone!"+ , "[Exit Ghost.]"+ , "We do it wrong, being so majestical,"+ , "To offer it the show of violence;"+ , "For it is as the air, invulnerable,"+ , "And our vain blows malicious mockery."+ , "Bernardo. It was about to speak, when the cock crew."+ , "Horatio. And then it started, like a guilty thing"+ , "Upon a fearful summons. I have heard"+ , "The cock, that is the trumpet to the morn,"+ , "Doth with his lofty and shrill-sounding throat"+ , "Awake the god of day; and at his warning,"+ , "Whether in sea or fire, in earth or air,"+ , "Th' extravagant and erring spirit hies"+ , "To his confine; and of the truth herein"+ , "This present object made probation."+ , "Marcellus. It faded on the crowing of the cock."+ , "Some say that ever, 'gainst that season comes"+ , "Wherein our Saviour's birth is celebrated,"+ , "The bird of dawning singeth all night long;"+ , "And then, they say, no spirit dare stir abroad,"+ , "The nights are wholesome, then no planets strike,"+ , "No fairy takes, nor witch hath power to charm,"+ , "So hallow'd and so gracious is the time."+ , "Horatio. So have I heard and do in part believe it."+ , "But look, the morn, in russet mantle clad,"+ , "Walks o'er the dew of yon high eastward hill."+ , "Break we our watch up; and by my advice"+ , "Let us impart what we have seen to-night"+ , "Unto young Hamlet; for, upon my life,"+ , "This spirit, dumb to us, will speak to him."+ , "Do you consent we shall acquaint him with it,"+ , "As needful in our loves, fitting our duty?"+ , "Let's do't, I pray; and I this morning know"+ , "Where we shall find him most conveniently."+ , "Exeunt."+ ]
+ src/Data/Trie/Word8.hs view
@@ -0,0 +1,534 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DeriveFunctor #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE ViewPatterns #-}++{- | Tries with 'Bytes' (equiv. 'ByteArray') as keys.+This implementation is optimized for performing queries rather+than updating the structure repeatedly.+-}+module Data.Trie.Word8+ ( -- * Trie Type+ Trie+ , valid++ -- * Query++ -- ** Lookup+ , lookup+ , lookupTrie+ , lookupPrefixes++ -- ** Search+ , multiFindReplace+ , search+ , replace+ , stripPrefix+ , stripPrefixWithKey++ -- ** Size+ , null+ , size++ -- * Construction+ , empty+ , singleton++ -- ** Conversion+ , fromList+ , toList+ , foldl'+ , traverse_++ -- ** Insertion+ , insert+ , insertWith++ -- ** Deletion+ , delete++ -- ** Combine+ , union+ , unionWith+ , append+ , prepend+ ) where++import Prelude hiding (lookup, null)++import Control.Applicative ((<|>))+import Data.Bifunctor (first)+import Data.Bytes (Bytes, fromByteArray, toByteArray)+import Data.Bytes.Chunks (Chunks)+import Data.Map.Word8 (Map)+import Data.Maybe (isNothing)+import Data.Monoid (Any (Any), getAny)+import Data.Primitive.ByteArray (ByteArray, indexByteArray)+import Data.Word (Word8)++import qualified Data.Bytes as Bytes+import qualified Data.Bytes.Builder as Build+import qualified Data.Foldable as Foldable+import qualified Data.Map.Word8 as Map+import qualified Data.Maybe.Unpacked as U++{- | Tries implemented using a 256-entry bitmap as given in+"Data.Map.Word8".+This means that each branch point can be navigated with only+some bit manipulations and adding an offset.+On sparse data, this should save a lot of space relative to holding+a 256-entry pointer array.++This data type has 'Tip', 'Run', and 'Branch' nodes.+Branches always have at least two children,+and Runs always have at least one byte.+Leaves are 'Tip's.+Once the invariants are met (see below),+there is exactly one 'Trie' representation for each trie.++In each constructor, the @U.Maybe a@ is a possible entry;+it comes before any child bytes.++INVARIANT: The Run constructor never has a linear child.+ Linear nodes are those with no value and exactly one child,+ which in this implementation is only valueless runs.+INVARIANT: The Run constructor never has zero bytes.+INVARIANT: The Branch constructor has at least two children.+INVARIANT: No child of a node has size zero. That includes:+ The next node after a run is never null.+ No child of a branch is ever null.+-}+data Trie a+ = Tip {-# UNPACK #-} !(U.Maybe a)+ | -- ByteArray uses more copying on modification,+ -- but the data structures are smaller than with Bytes, making lookup faster+ UnsafeRun {-# UNPACK #-} !(U.Maybe a) {-# UNPACK #-} !ByteArray !(Trie a)+ | UnsafeBranch {-# UNPACK #-} !(U.Maybe a) !(Map (Trie a))+ deriving (Eq, Functor)++instance (Semigroup a) => Semigroup (Trie a) where (<>) = append+instance (Semigroup a) => Monoid (Trie a) where mempty = empty+instance (Show a) => Show (Trie a) where show = show . toList++{-# COMPLETE Tip, Run, Branch #-}++pattern Run :: U.Maybe a -> ByteArray -> Trie a -> Trie a+pattern Run v run next <- UnsafeRun v run next+ where+ Run v run next+ | null next = Tip v+ | Bytes.null (fromByteArray run) = next -- WARNING: throws away `v`, value/non-value from `next`+ | Just (run', next') <- fromLinear next =+ UnsafeRun v (run <> run') next'+ | otherwise = UnsafeRun v run next++pattern Branch :: U.Maybe a -> Map (Trie a) -> Trie a+pattern Branch v children <- UnsafeBranch v children+ where+ Branch v (removeEmptyChildren -> children)+ | Map.null children = Tip v+ | Just (c, child) <- fromSingletonMap children =+ Run v (toByteArray $ Bytes.singleton c) child+ | otherwise = UnsafeBranch v children+removeEmptyChildren :: Map (Trie a) -> Map (Trie a)+removeEmptyChildren = Map.foldrWithKeys f Map.empty+ where+ f k v xs = if null v then xs else Map.insert k v xs++-- Get nodes with no value, and exactly one possible next byte.+-- I.e. it never returns an empty bytes in the tuple.+fromLinear :: Trie a -> Maybe (ByteArray, Trie a)+fromLinear (Run U.Nothing run next) = Just (run, next)+fromLinear _ = Nothing++valid :: Trie a -> Bool+valid (Tip _) = True+valid (Run _ run next) =+ not (Bytes.null (fromByteArray run))+ && isNothing (fromLinear next)+ && not (null next)+ && valid next+valid (Branch _ children) =+ Map.size children > 1+ && Map.foldrWithKeys nonNullChild True children+ where+ nonNullChild _ child !acc = acc && not (null child)++------------ Find/Replace ------------++{- | The raison-d'etre of this library: repeatedly search in a byte string+for the longest of multiple patterns and make replacements.+-}+multiFindReplace ::+ (Semigroup b) =>+ -- | construct a portion of the result from unmatched bytes+ (Bytes -> b) ->+ -- | construct a replacement from the found value+ (a -> b) ->+ -- | the dictionary of all replacements+ Trie a ->+ -- | input to be edited+ Bytes ->+ -- | result of replacement+ b+{-# INLINE multiFindReplace #-}+multiFindReplace fromNoMatch fromMatch = \t ->+ let needles = delete mempty t+ -- `into` counts up until the first index where a replacement is found+ go !into rawInp =+ let inp = Bytes.unsafeDrop into rawInp+ unMatched = Bytes.unsafeTake into rawInp+ in if+ | Bytes.null inp -> fromNoMatch unMatched+ | Just (val, rest) <- stripPrefix needles inp ->+ fromNoMatch unMatched <> fromMatch val <> go 0 rest+ | otherwise -> go (into + 1) rawInp+ in go 0++replace :: Trie Bytes -> Bytes -> Chunks+replace t inp = Build.run 4080 $ go t inp+ where+ go = multiFindReplace Build.bytes Build.bytes++search :: Trie a -> Bytes -> Bool+search t inp = getAny $ go t inp+ where+ go = multiFindReplace (const mempty) (const $ Any True)++------------ Construction ------------++-- | The empty trie.+empty :: Trie a+empty = Tip U.Nothing++-- | A trie with a single element.+singleton :: Bytes -> a -> Trie a+singleton k v = prepend k $ Tip (U.Just v)++{- | Prepend every key in the 'Trie' with the given 'Bytes'.++This should be used internally instead of the Run ctor,+thereby ensuring the run length >= 2 invariant is maintained.+It is exported anyway because someone may find it useful.+-}+prepend :: Bytes -> Trie a -> Trie a+prepend bytes next = Run U.Nothing (toByteArray bytes) next++{- | Insert a new key/value into the trie.+If the key is already present in the trie, the associated value is+replaced with the new one.+'insert' is equivalent to 'insertWith' 'const'.+-}+insert :: Bytes -> a -> Trie a -> Trie a+insert = insertWith const++{- | Insert with a function, combining new value and old value.+@'insertWith' f key value trie@ will insert the pair @(key, value)@+into @trie@ if @key@ does not exist in the trie.+If the key does exist, the function will insert the pair+@(key, f new_value old_value)@.+-}+insertWith :: (a -> a -> a) -> Bytes -> a -> Trie a -> Trie a+insertWith f k v = unionWith f (singleton k v)++delete :: Bytes -> Trie a -> Trie a+delete k0 trie = go k0 trie+ where+ -- `go` is not always tail-recursive.+ -- Instead, each node with exactly one child must be checked after the+ -- deletion to ensure that child is non-empty.+ -- TODO+ -- However, as soon as it is known that the size must be greater than one,+ -- we can throw away all queued normalizations so far.+ -- Therefore, we maintain a delimited continuation as an accumulator,+ -- but I'm not yet sure how to manually store it.+ -- go :: Bytes -> Trie a+ go key node@(Tip v)+ | Bytes.null key+ , U.Just _ <- v -- NOTE this is redundant now, but when I use cps, it won't be+ =+ empty+ | otherwise = node+ go key node@(Run v (fromByteArray -> run) next)+ -- found key, therefore delete+ | Bytes.null key+ , U.Just _ <- v -- NOTE this is redundant now, but when I use cps, it won't be+ =+ prepend run next+ -- carry on searching for the key+ | Just key' <- Bytes.stripPrefix run key =+ Run v (toByteArray run) (go key' next)+ -- key not present+ | otherwise = node+ go key node@(Branch v children)+ -- found key, therefore delete+ | Bytes.null key+ , U.Just _ <- v =+ UnsafeBranch U.Nothing children+ -- carry on searching for the key+ | Just (c, key') <- Bytes.uncons key+ , Just child <- Map.lookup c children =+ Branch v (Map.insert c (go key' child) children)+ -- key not present+ | otherwise = node++{- | Union of the two tries, but where a key appears in both,+the associated values are combined with '(<>)' to produce the new value,+i.e. @append == unionWith (<>)@.+-}+append :: (Semigroup a) => Trie a -> Trie a -> Trie a+append = unionWith (<>)++{- | The left-biased union of the two tries.+It prefers the first when duplicate keys are encountered,+i.e. @union == unionWith const@.+-}+union :: Trie a -> Trie a -> Trie a+union = unionWith const++-- | Union with a combining function.+unionWith :: (a -> a -> a) -> Trie a -> Trie a -> Trie a+unionWith f trieA trieB = case (trieA, trieB) of+ (Tip a, Tip b) -> Tip (a `mergeValue` b)+ (Tip a, Run b run next) -> UnsafeRun (a `mergeValue` b) run next+ (Run a run next, Tip b) -> UnsafeRun (a `mergeValue` b) run next+ (Tip a, Branch b children) -> UnsafeBranch (a `mergeValue` b) children+ (Branch a children, Tip b) -> UnsafeBranch (a `mergeValue` b) children+ -- all non-Tip cases+ (Branch a children, Branch b children') ->+ UnsafeBranch (a `mergeValue` b) (mergeChildren children children')+ (Branch a children, r@(Run _ _ _)) ->+ UnsafeBranch (a `mergeValue` b) (mergeChildren children children')+ where+ (b, c, child') = unsafeUnconsRun r+ children' = Map.singleton c child'+ (r@(Run _ _ _), Branch b children') ->+ UnsafeBranch (a `mergeValue` b) (mergeChildren children children')+ where+ (a, c, child') = unsafeUnconsRun r+ children = Map.singleton c child'+ (Run a (fromByteArray -> run) next, Run b (fromByteArray -> run') next') ->+ if Bytes.null common+ then+ let mkChild bytes trie = case Bytes.uncons bytes of+ Prelude.Just (c, k) -> Map.singleton c (prepend k trie)+ Prelude.Nothing -> error "invariant violation: empty run bytes"+ child = mkChild run next+ child' = mkChild run' next'+ in UnsafeBranch (a `mergeValue` b) $ child `Map.union` child'+ else+ let child = prepend (uncommon run) next+ child' = prepend (uncommon run') next'+ in Run (a `mergeValue` b) (toByteArray common) $ unionWith f child child'+ where+ common = Bytes.longestCommonPrefix run run'+ uncommon bytes = Bytes.unsafeDrop (Bytes.length common) bytes+ where+ mergeChildren left right = Map.unionWith (unionWith f) left right+ mergeValue U.Nothing U.Nothing = U.Nothing+ mergeValue (U.Just x) (U.Just y) = U.Just (f x y)+ mergeValue x y = x <|> y++------------ Conversion ------------++{- | Build a trie from a list of key/value pairs.+If more than one value for the same key appears, the last value for that+key is retained.+-}+fromList :: [(Bytes, a)] -> Trie a+fromList kvs = Foldable.foldl' (\xs (k, v) -> insert k v xs) empty kvs++{- | Convert the trie to a list of key/value pairs.+The resulting list has its keys sorted in ascending order.+-}+toList :: Trie a -> [(Bytes, a)]+toList = \case+ Tip valO -> fromValue valO+ Run valO run next -> fromValue valO ++ prependList (fromByteArray run) (toList next)+ Branch valO children -> fromValue valO ++ Map.foldrWithKeys f [] children+ where+ f k v acc = prependList (Bytes.singleton k) (toList v) ++ acc+ where+ fromValue valO = (mempty,) <$> U.maybeToList valO+ prependList run list = first (run <>) <$> list++foldl' :: (b -> a -> b) -> b -> Trie a -> b+{-# INLINE foldl' #-}+foldl' f !b0 t0 = go b0 t0+ where+ go !b t = case t of+ Tip valO -> case valO of+ U.Just x -> f b x+ _ -> b+ Run valO _ next ->+ let b' = case valO of+ U.Just x -> f b x+ _ -> b+ in go b' next+ Branch valO children ->+ let b' = case valO of+ U.Just x -> f b x+ _ -> b+ in Map.foldl' go b' children++traverse_ :: (Applicative m) => (a -> m b) -> Trie a -> m ()+{-# INLINE traverse_ #-}+traverse_ f t0 = go t0+ where+ go t = case t of+ Tip valO -> case valO of+ U.Just x -> f x *> pure ()+ _ -> pure ()+ Run valO _ next -> case valO of+ U.Just x -> f x *> go next+ _ -> go next+ Branch valO children -> case valO of+ U.Just x -> f x *> Map.traverse_ go children+ _ -> Map.traverse_ go children++------------ Query ------------++-- | Lookup the value at the 'Bytes' key in the trie.+lookup :: Bytes -> Trie a -> Maybe a+lookup k (Tip v)+ | Bytes.null k = U.toBaseMaybe v+ | otherwise = Nothing+lookup k (Run v (fromByteArray -> run) next)+ | Bytes.null k = U.toBaseMaybe v+ | run `Bytes.isPrefixOf` k =+ let k' = Bytes.unsafeDrop (Bytes.length run) k+ in lookup k' next+ | otherwise = Prelude.Nothing+lookup k (Branch valO children) = case Bytes.uncons k of+ Prelude.Nothing -> U.toBaseMaybe valO+ Prelude.Just (c, k') -> lookup k' =<< Map.lookup c children++{- | Lookup the value at the 'Bytes' key in the trie. Returns the value+of the exact match and the values for any keys that are prefixes of+the search key. The shortest prefix is first. The exact match (if there+is one) is last.+-}+lookupPrefixes :: Bytes -> Trie a -> [a]+{-# INLINE lookupPrefixes #-}+lookupPrefixes = lookupPrefixesGo []++lookupPrefixesGo :: [a] -> Bytes -> Trie a -> [a]+lookupPrefixesGo !acc !_ (Tip v)+ | U.Just x <- v = x : acc+ | otherwise = acc+lookupPrefixesGo !acc !k (UnsafeRun v (fromByteArray -> run) next) =+ let acc' = case v of+ U.Just x -> x : acc+ _ -> acc+ in if+ | Bytes.null k -> acc'+ | run `Bytes.isPrefixOf` k ->+ let k' = Bytes.unsafeDrop (Bytes.length run) k+ in lookupPrefixesGo acc' k' next+ | otherwise -> acc'+lookupPrefixesGo !acc !k (UnsafeBranch valO children) =+ let acc' = case valO of+ U.Just x -> x : acc+ _ -> acc+ in case Bytes.uncons k of+ Prelude.Nothing -> acc'+ Prelude.Just (c, k') -> case Map.lookup c children of+ Nothing -> acc'+ Just child -> lookupPrefixesGo acc' k' child++{- | Lookup the trie at the 'Bytes' key in the trie. Returns the subtrie+at this position.++>>> (k1 <> k2 == k) ==> (lookup k v t == lookup k2 (lookupTrie k1 t))+-}+lookupTrie :: Bytes -> Trie a -> Trie a+lookupTrie !k trie+ | Bytes.null k = trie+ | otherwise = case trie of+ Tip {} -> empty+ Run _ (fromByteArray -> run) next+ | run `Bytes.isPrefixOf` k ->+ let k' = Bytes.unsafeDrop (Bytes.length run) k+ in lookupTrie k' next+ | k `Bytes.isPrefixOf` run ->+ let run' = Bytes.unsafeDrop (Bytes.length k) run+ in Run U.Nothing (Bytes.toByteArrayClone run') next+ | otherwise -> empty+ Branch _ children ->+ let !k' = Bytes.unsafeDrop 1 k+ !c = Bytes.unsafeIndex k 0+ in case Map.lookup c children of+ Nothing -> empty+ Just child -> lookupTrie k' child++{- | Find the longest prefix of the input 'Bytes' which has a value in the trie.+Returns the associated value and the remainder of the input after the prefix.+-}+stripPrefix :: Trie a -> Bytes -> Maybe (a, Bytes)+stripPrefix trie inp = first snd <$> stripPrefixWithKey trie inp++{- | Find the longest prefix of the input 'Bytes' which has a value in the trie.+Returns the prefix and associated value found as a key/value tuple,+and also the remainder of the input after the prefix.+-}+stripPrefixWithKey :: forall a. Trie a -> Bytes -> Maybe ((Bytes, a), Bytes)+stripPrefixWithKey trie0 rawInp = go 0 Nothing trie0+ where+ go :: Int -> Maybe (Bytes, a) -> Trie a -> Maybe ((Bytes, a), Bytes)+ go !into !prior node =+ let inp = Bytes.unsafeDrop into rawInp+ candidate = (Bytes.unsafeTake into rawInp,) <$> topValue node+ found = candidate <|> prior+ in if+ | Run _ (fromByteArray -> run) next <- node+ , run `Bytes.isPrefixOf` inp ->+ go (into + Bytes.length run) found next+ | Branch _ children <- node+ , Just (c, _) <- Bytes.uncons inp+ , Just next <- Map.lookup c children ->+ go (into + 1) found next+ | otherwise -> mkReturn <$> found+ mkReturn (prefix, v) =+ let post = Bytes.unsafeDrop (Bytes.length prefix) rawInp+ in ((prefix, v), post)++null :: Trie a -> Bool+null (Tip U.Nothing) = True+null _ = False++size :: Trie a -> Int+size node = here + under+ where+ here = maybe 0 (const 1) (topValue node)+ under = case node of+ Tip _ -> 0+ Run _ _ next -> size next+ Branch _ children -> Map.foldrWithKeys (\_ child !acc -> acc + size child) 0 children++------ Helpers ------++topValue :: Trie a -> Maybe a+topValue = \case+ Tip v -> U.toBaseMaybe v+ Run v _ _ -> U.toBaseMaybe v+ Branch v _ -> U.toBaseMaybe v++unsafeUnconsRun :: Trie a -> (U.Maybe a, Word8, Trie a)+unsafeUnconsRun (Run v0 bs next) = (v0, c, run')+ where+ c = indexByteArray bs 0+ bs' = Bytes.unsafeDrop 1 (fromByteArray bs)+ run' = prepend bs' next+unsafeUnconsRun (Tip _) = error "unsafeUnconsRun on Tip trie"+unsafeUnconsRun (Branch _ _) = error "unsafeUnconsRun on Branch trie"++-- TODO is this really a decent way to do this?+fromSingletonMap :: Map a -> Maybe (Word8, a)+fromSingletonMap mp = case Map.toList mp of+ [(c, v)] -> Just (c, v)+ _ -> Nothing
+ test/Main.hs view
@@ -0,0 +1,145 @@+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# OPTIONS_GHC -fno-warn-orphans #-}++import Control.Monad (forM_)+import Data.Bifunctor (first)+import Data.Bytes.Types (Bytes (Bytes))+import Data.List (isInfixOf, sort)+import Data.Monoid (Sum)+import Data.Proxy (Proxy (..))+import Data.Trie.Word8 (Trie)+import Data.Word (Word8)+import Test.Tasty (TestTree, defaultMain, testGroup)+import Test.Tasty.HUnit (assertBool, testCase)+import Test.Tasty.QuickCheck (Arbitrary, Discard (Discard), property, testProperty, (===))++import qualified Data.Bytes as Bytes+import qualified Data.Bytes.Text.Latin1 as Latin1+import qualified Data.Trie.Word8 as Trie+import qualified GHC.Exts as Exts+import qualified Test.QuickCheck.Classes as QCC+import qualified Test.Tasty.QuickCheck as TQC+import qualified TestData++main :: IO ()+main = defaultMain tests++tests :: TestTree+tests =+ testGroup+ "bytetrie"+ [ testGroup+ "validity"+ [ testProperty "fromList validity" $ \alist ->+ let a = Trie.fromList alist :: Trie Int+ in Trie.valid a+ , testProperty "unionWith maintains invariants" $ \xsList ysList ->+ let xs = Trie.fromList xsList :: Trie Int+ ys = Trie.fromList ysList+ in Trie.valid (Trie.unionWith (+) xs ys)+ , testProperty "delete maintains invariants" $ \alist k v ->+ let a = Trie.fromList alist :: Trie Int+ a' = Trie.insert k v a+ in Trie.valid (Trie.delete k a) && Trie.valid (Trie.delete k a')+ ]+ , testGroup+ "lookup"+ [ testProperty "lookup existing key" $+ \(xs :: [(Bytes, Int)]) (k, v) ys ->+ let alist = xs ++ [(k, v)] ++ ys+ in if+ | Nothing <- lookup k ys ->+ Trie.lookup k (Trie.fromList alist) === lookup k (reverse alist)+ | otherwise -> property Discard+ , testProperty "lookup missing key" $+ \(alist :: [(Bytes, Int)]) (k :: Bytes) ->+ if+ | Just _ <- lookup k alist -> property Discard+ | otherwise -> Trie.lookup k (Trie.fromList alist) === Nothing+ ]+ , testGroup+ "unionWith"+ [ lawsToTest (QCC.semigroupLaws (Proxy :: Proxy (Trie [Integer])))+ , lawsToTest (QCC.monoidLaws (Proxy :: Proxy (Trie [Integer])))+ , lawsToTest (QCC.commutativeMonoidLaws (Proxy :: Proxy (Trie (Sum Integer))))+ ]+ , lawsToTest (QCC.functorLaws (Proxy :: Proxy Trie))+ , testGroup+ "lookupTrie"+ [ testProperty "alpha" $+ \(x :: Bytes) (y :: Bytes) (val :: Int) (xs :: [(Bytes, Int)]) ->+ let base = Trie.insert (x <> y) val (Trie.fromList xs)+ in Trie.lookup (x <> y) base === Trie.lookup y (Trie.lookupTrie x base)+ , testProperty "beta" $+ \(x :: Bool) (y :: Bool) (val :: Int) (xs :: [([Bool], Int)]) ->+ let xs' = map (first (foldMap boolToBytes)) xs+ base = Trie.insert (boolToBytes x <> boolToBytes y) val (Trie.fromList xs')+ in Trie.lookup (boolToBytes x <> boolToBytes y) base === Trie.lookup (boolToBytes y) (Trie.lookupTrie (boolToBytes x) base)+ ]+ , testGroup+ "stripPrefix"+ [ testProperty "finds longest prefix" $+ \(a :: Bytes) (b :: Bytes) (c :: Bytes) (d :: Bytes) ->+ if+ | Just (wc, _) <- Bytes.uncons c+ , Just (wd, _) <- Bytes.uncons d+ , wc /= wd ->+ let t :: Trie Int = Trie.fromList [(a, 1), (a <> b, 2), (a <> b <> d, 3)]+ expected = Just ((a <> b, 2), c)+ found = Trie.stripPrefixWithKey t (a <> b <> c)+ in found === expected+ | otherwise -> property Discard+ , testProperty "finds nothing" $+ \(a :: Bytes) (b :: Bytes) ->+ if+ | Just (wa, _) <- Bytes.uncons a+ , Just (wb, _) <- Bytes.uncons b+ , wa /= wb ->+ let t :: Trie Int = Trie.fromList [(a, 1)]+ in Trie.stripPrefixWithKey t b === Nothing+ | otherwise -> property Discard+ ]+ , testGroup+ "delete"+ [ testProperty "maintains invariant" $+ \(xs :: [(Bytes, Int)]) (k, v) ys ->+ if+ | Nothing <- lookup k xs+ , Nothing <- lookup k ys ->+ let t = Trie.fromList $ xs ++ [(k, v)] ++ ys+ t' = Trie.fromList $ xs ++ ys+ in Trie.delete k t === t'+ | otherwise -> property Discard+ ]+ , testProperty "toList is sorted" $ \alist ->+ let a = Trie.fromList alist :: Trie ()+ out = Trie.toList a+ in out == sort out+ , testCase "sed works" $ do+ let sedList = Trie.multiFindReplace Latin1.toString Latin1.toString+ outp = sedList TestData.replacements TestData.bigstring+ replaced = ["Francisco", "Bernardo", "Marcellus", "Horatio", "Ghost", "What", "Why"]+ forM_ replaced $ \word -> do+ let assertName = "replaced " ++ word+ assertBool assertName $ not (word `isInfixOf` outp)+ ]++instance (Arbitrary a) => Arbitrary (Trie a) where+ arbitrary = Trie.fromList <$> TQC.arbitrary+instance Arbitrary Bytes where+ arbitrary = do+ xs :: [Word8] <- TQC.arbitrary+ front <- TQC.choose (0, 2)+ back <- TQC.choose (0, 2)+ let frontPad = replicate front (254 :: Word8)+ let backPad = replicate back (254 :: Word8)+ let raw = Exts.fromList (frontPad ++ xs ++ backPad)+ pure (Bytes raw front (length xs))++lawsToTest :: QCC.Laws -> TestTree+lawsToTest (QCC.Laws name pairs) = testGroup name (map (uncurry TQC.testProperty) pairs)++boolToBytes :: Bool -> Bytes+boolToBytes True = Bytes.singleton 1+boolToBytes False = Bytes.singleton 0