diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for bytetrie + +## 0.1.0.0 -- YYYY-mm-dd + +* First version. Released on an unsuspecting world. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,3 @@ +import Distribution.Simple + +main = defaultMain diff --git a/bytetrie.cabal b/bytetrie.cabal new file mode 100644 --- /dev/null +++ b/bytetrie.cabal @@ -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 diff --git a/common/TestData.hs b/common/TestData.hs new file mode 100644 --- /dev/null +++ b/common/TestData.hs @@ -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." + ] diff --git a/src/Data/Trie/Word8.hs b/src/Data/Trie/Word8.hs new file mode 100644 --- /dev/null +++ b/src/Data/Trie/Word8.hs @@ -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 diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -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