diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -30,14 +30,22 @@ data set, the benchmark setup, and a few things to keep in mind when interpreting this graph. +### LLVM + +If you are using LLVM instead of the GHC backend, make sure to compare different versions. +Using LLVM 9 instead of LLVM 12 with GHC 8.10.7 made a significant dent in benchmark times. +For more information, see [this issue][llvm]. + ## Example Check if a string contains one of the needles: ```haskell + +import qualified Data.Text.AhoCorasick.Automaton as Aho import qualified Data.Text.AhoCorasick.Searcher as Searcher -searcher = Searcher.build ["tshirt", "shirts", "shorts"] +searcher = Searcher.build Aho.CaseSensitive ["tshirt", "shirts", "shorts"] Searcher.containsAny searcher "short tshirts" > True @@ -48,7 +56,9 @@ Searcher.containsAny searcher "Short TSHIRTS" > False -Searcher.containsAnyIgnoreCase searcher "Short TSHIRTS" +searcher' = Searcher.build Aho.IgnoreCase ["tshirt", "shirts", "shorts"] + +Searcher.containsAny searcher' "Short TSHIRTS" > True ``` @@ -105,3 +115,4 @@ [text]: https://github.com/haskell/text [hankcs]: https://github.com/hankcs/AhoCorasickDoubleArrayTrie/tree/v1.2.0 [burntsushi]: https://github.com/BurntSushi/aho-corasick/tree/0.6.8 +[llvm]: https://github.com/channable/alfred-margaret/issues/24 diff --git a/alfred-margaret.cabal b/alfred-margaret.cabal --- a/alfred-margaret.cabal +++ b/alfred-margaret.cabal @@ -1,5 +1,5 @@ name: alfred-margaret -version: 1.1.1.0 +version: 1.1.2.0 synopsis: Fast Aho-Corasick string searching description: An efficient implementation of the Aho-Corasick string searching algorithm. @@ -18,6 +18,8 @@ GHC == 8.6.3 -- Stackage LTS 16.18. , GHC == 8.8.4 + -- Stackage LTS 18.27 + , GHC == 8.10.7 source-repository head type: git @@ -46,7 +48,17 @@ , Data.Text.BoyerMoore.Automaton , Data.Text.BoyerMoore.Replacer , Data.Text.BoyerMoore.Searcher + , Data.Text.CaseSensitivity , Data.Text.Utf16 + , Data.Text.Utf8 + , Data.Text.Utf8.AhoCorasick.Automaton + , Data.Text.Utf8.AhoCorasick.Replacer + , Data.Text.Utf8.AhoCorasick.Searcher + , Data.Text.Utf8.AhoCorasick.Splitter + , Data.Text.Utf8.BoyerMoore.Automaton + , Data.Text.Utf8.BoyerMoore.Replacer + , Data.Text.Utf8.BoyerMoore.Searcher + , Data.TypedByteArray build-depends: base >= 4.7 && < 5 , containers >= 0.6 && < 0.7 @@ -56,10 +68,11 @@ , text >= 1.2.3 && < 1.3 , unordered-containers >= 0.2.9 && < 0.3 , vector >= 0.12 && < 0.13 + , bytestring >= 0.10.12 && < 1 ghc-options: -Wall -Wincomplete-record-updates -Wincomplete-uni-patterns -O2 default-language: Haskell2010 if flag(aeson) { - build-depends: aeson >= 1.4.2 && < 1.6 + build-depends: aeson >= 1.4.2 && < 3 cpp-options: -DHAS_AESON } if flag(llvm) { @@ -71,6 +84,9 @@ main-is: Main.hs other-modules: Data.Text.AhoCorasickSpec , Data.Text.BoyerMooreSpec + , Data.Text.Utf8.AhoCorasickSpec + , Data.Text.Utf8.BoyerMooreSpec + , Data.Text.Utf8.Utf8Spec , Data.Text.Orphans hs-source-dirs: tests ghc-options: -Wall -Wincomplete-record-updates -Wno-orphans @@ -80,6 +96,27 @@ , deepseq , hspec , hspec-expectations + , primitive , quickcheck-instances , text default-language: Haskell2010 + +benchmark uvector-vs-tba + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: bench/uvector-vs-tba + ghc-options: -Wall -Wincomplete-record-updates -Wno-orphans + build-depends: base >= 4.7 && < 5 + , alfred-margaret + , vector + , deepseq + , criterion + default-language: Haskell2010 + +executable dump-automaton + main-is: Main.hs + hs-source-dirs: + app/dump-automaton + build-depends: base + , alfred-margaret + default-language: Haskell2010 diff --git a/app/dump-automaton/Main.hs b/app/dump-automaton/Main.hs new file mode 100644 --- /dev/null +++ b/app/dump-automaton/Main.hs @@ -0,0 +1,18 @@ +module Main where + +import Control.Monad (forM) +import qualified Data.Text.Utf8 as Utf8 +import Data.Text.Utf8.AhoCorasick.Automaton (debugBuildDot) +import System.Environment (getArgs) +import System.IO (hPrint, hPutStr, stderr) + +main = do + args <- getArgs + needles <- forM args $ \needle -> do + hPutStr stderr $ needle ++ ": " + let needleBytes = Utf8.unpackUtf8 $ Utf8.pack needle + hPrint stderr needleBytes + pure $ Utf8.pack needle + + let dot = debugBuildDot needles + putStrLn dot diff --git a/bench/uvector-vs-tba/Main.hs b/bench/uvector-vs-tba/Main.hs new file mode 100644 --- /dev/null +++ b/bench/uvector-vs-tba/Main.hs @@ -0,0 +1,76 @@ +{-# LANGUAGE BangPatterns #-} + +-- | Microbenchmark to measure the difference in 'TypedByteArray' and 'Data.Vector.Unboxed.Vector' indexing performance. +-- The latter is a slice into a 'ByteArray' which means that every time you retrieve and element, an offset +-- needs to be added to your index. +-- +-- To reproduce: +-- +-- @ +-- stack bench alfred-margaret:uvector-vs-tba --ba '--output uvector-vs-tba.html' +-- @ +-- +-- You can pass a greater @--time-limit@ (in the single quotes) to increase the number of iterations. +-- +-- NOTE: 'readUVector' and 'genUVector' are marked @NOINLINE@ to prevent GHC optimizing away the indexing addition. +-- 'readTba' and 'genTba' are marked @NOINLINE@ as well for fairness, altough it shouldn't make a difference there. +module Main where + +import Control.Monad.ST (runST) +import Criterion.Main (Benchmark, bench, bgroup, defaultMain, nf) +import Data.Foldable (for_) +import Text.Printf (printf) + +import qualified Data.Vector.Unboxed as UVector +import qualified Data.Vector.Unboxed.Mutable as UMVector + +import qualified Data.TypedByteArray as TBA + +main :: IO () +main = defaultMain + [ bgroup "tba" $ mkReadBenchs readTba genTba [7, 8] + , bgroup "uvector" $ mkReadBenchs readUVector genUVector [7, 8] + ] + +mkReadBenchs + :: (Int -> a -> Int) -- ^ Function that reads from an array @n@ times. + -> (Int -> a) -- ^ Function that constructs an array of length @n@. + -> [Int] -- ^ Which powers of 10 to pass for @n@. + -> [Benchmark] +mkReadBenchs readPattern gen powers = + [ bench (printf "%d reads" n) $ nf (readPattern n) $ gen n + | n <- map (10^) powers + ] + +{-# NOINLINE readTba #-} +readTba :: Int -> TBA.TypedByteArray Int -> Int +readTba !n !arr = go 0 + where + go !i + | i >= n = 42 + | otherwise = go $ TBA.unsafeIndex arr i + +{-# NOINLINE readUVector #-} +readUVector :: Int -> UVector.Vector Int -> Int +readUVector !n !arr = go 0 + where + go !i + | i >= n = 42 + | otherwise = go $ UVector.unsafeIndex arr i + +-- NOTE: We should probably measure pseudo-random access time as well, e.g. by shuffling the generated arrays. + +{-# NOINLINE genTba #-} +genTba :: Int -> TBA.TypedByteArray Int +genTba n = runST $ do + mutArr <- TBA.newTypedByteArray n + for_ [0 .. n-1] $ \i -> TBA.writeTypedByteArray mutArr i $ i + 1 + TBA.unsafeFreezeTypedByteArray mutArr + +-- | Generate an unboxed vector @v@ such that @v[i] == i + 1@. +{-# NOINLINE genUVector #-} +genUVector :: Int -> UVector.Vector Int +genUVector n = runST $ do + mutArr <- UMVector.new n + for_ [0 .. n-1] $ \i -> UMVector.write mutArr i $ i + 1 + UVector.unsafeFreeze mutArr diff --git a/src/Data/Text/AhoCorasick/Automaton.hs b/src/Data/Text/AhoCorasick/Automaton.hs --- a/src/Data/Text/AhoCorasick/Automaton.hs +++ b/src/Data/Text/AhoCorasick/Automaton.hs @@ -7,8 +7,6 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE DerivingStrategies #-} -{-# LANGUAGE DeriveAnyClass #-} {-# LANGUAGE ScopedTypeVariables #-} -- | An efficient implementation of the Aho-Corasick string matching algorithm. @@ -29,49 +27,36 @@ -- automaton as int maps, which are convenient for incremental construction. -- Afterwards we pack the automaton into unboxed vectors. module Data.Text.AhoCorasick.Automaton - ( AcMachine (..) - , build - , runText - , runLower - , debugBuildDot - , CaseSensitivity (..) - , CodeUnitIndex (..) - , Match (..) - , Next (..) - ) - where + ( AcMachine (..) + , CaseSensitivity (..) + , CodeUnitIndex (..) + , Match (..) + , Next (..) + , build + , debugBuildDot + , runLower + , runText + ) where import Prelude hiding (length) import Control.DeepSeq (NFData) -import Data.Bits ((.&.), (.|.), shiftL, shiftR) +import Data.Bits (shiftL, shiftR, (.&.), (.|.)) import Data.Foldable (foldl') -import Data.Hashable (Hashable) import Data.IntMap.Strict (IntMap) import Data.Text.Internal (Text (..)) import Data.Word (Word64) import GHC.Generics (Generic) -#if defined(HAS_AESON) -import Data.Aeson (FromJSON, ToJSON) -#endif - import qualified Data.IntMap.Strict as IntMap import qualified Data.List as List import qualified Data.Vector as Vector -import qualified Data.Vector.Unboxed as UVector +import Data.Text.CaseSensitivity (CaseSensitivity (..)) import Data.Text.Utf16 (CodeUnit, CodeUnitIndex (..), indexTextArray, lowerCodeUnit) +import Data.TypedByteArray (Prim, TypedByteArray) -data CaseSensitivity - = CaseSensitive - | IgnoreCase - deriving stock (Eq, Generic, Show) -#if defined(HAS_AESON) - deriving anyclass (Hashable, NFData, FromJSON, ToJSON) -#else - deriving anyclass (Hashable, NFData) -#endif +import qualified Data.TypedByteArray as TBA -- | A numbered state in the Aho-Corasick automaton. type State = Int @@ -99,7 +84,7 @@ { matchPos :: {-# UNPACK #-} !CodeUnitIndex -- ^ The code unit index past the last code unit of the match. Note that this -- is not a code *point* (Haskell `Char`) index; a code point might be encoded - -- as two code units. + -- as up to four code units. , matchValue :: v -- ^ The payload associated with the matched needle. } deriving (Show, Eq) @@ -109,14 +94,14 @@ { machineValues :: !(Vector.Vector [v]) -- ^ For every state, the values associated with its needles. If the state is -- not a match state, the list is empty. - , machineTransitions :: !(UVector.Vector Transition) + , machineTransitions :: !(TypedByteArray Transition) -- ^ A packed vector of transitions. For every state, there is a slice of this -- vector that starts at the offset given by `machineOffsets`, and ends at the -- first wildcard transition. - , machineOffsets :: !(UVector.Vector Int) + , machineOffsets :: !(TypedByteArray Int) -- ^ For every state, the index into `machineTransitions` where the transition -- list for that state starts. - , machineRootAsciiTransitions :: !(UVector.Vector Transition) + , machineRootAsciiTransitions :: !(TypedByteArray Transition) -- ^ A lookup table for transitions from the root state, an optimization to -- avoid having to walk all transitions, at the cost of using a bit of -- additional memory. @@ -161,11 +146,11 @@ -- the transitions for a specific state, we also produce a vector of start -- indices. All transition lists are terminated by a wildcard transition, so -- there is no need to record the length. -packTransitions :: [[Transition]] -> (UVector.Vector Transition, UVector.Vector Int) +packTransitions :: [[Transition]] -> (TypedByteArray Transition, TypedByteArray Int) packTransitions transitions = let - packed = UVector.fromList $ concat transitions - offsets = UVector.fromList $ scanl (+) 0 $ fmap List.length transitions + packed = TBA.fromList $ concat transitions + offsets = TBA.fromList $ scanl (+) 0 $ fmap List.length transitions in (packed, offsets) @@ -298,8 +283,8 @@ -- | Build a lookup table for the first 128 code units, that can be used for -- O(1) lookup of a transition, rather than doing a linear scan over all -- transitions. The fallback goes back to the initial state, state 0. -buildAsciiTransitionLookupTable :: IntMap State -> UVector.Vector Transition -buildAsciiTransitionLookupTable transitions = UVector.generate asciiCount $ \i -> +buildAsciiTransitionLookupTable :: IntMap State -> TypedByteArray Transition +buildAsciiTransitionLookupTable transitions = TBA.generate asciiCount $ \i -> case IntMap.lookup i transitions of Just state -> newTransition (fromIntegral i) state Nothing -> newWildcardTransition 0 @@ -385,8 +370,9 @@ at :: forall a. Vector.Vector a -> Int -> a at = Vector.unsafeIndex -uAt :: forall a. UVector.Unbox a => UVector.Vector a -> Int -> a -uAt = UVector.unsafeIndex +{-# INLINE uAt #-} +uAt :: Prim a => TypedByteArray a -> Int -> a +uAt = TBA.unsafeIndex -- | Result of handling a match: stepping the automaton can exit early by -- returning a `Done`, or it can continue with a new accumulator with `Step`. diff --git a/src/Data/Text/AhoCorasick/Searcher.hs b/src/Data/Text/AhoCorasick/Searcher.hs --- a/src/Data/Text/AhoCorasick/Searcher.hs +++ b/src/Data/Text/AhoCorasick/Searcher.hs @@ -9,31 +9,33 @@ {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} module Data.Text.AhoCorasick.Searcher - ( Searcher - , build - , buildWithValues - , needles - , numNeedles - , automaton - , caseSensitivity - , containsAny - , setSearcherCaseSensitivity - ) - where + ( Searcher + , automaton + , build + , buildNeedleIdSearcher + , buildWithValues + , caseSensitivity + , containsAll + , containsAny + , needles + , numNeedles + , setSearcherCaseSensitivity + ) where import Control.DeepSeq (NFData) import Data.Hashable (Hashable (hashWithSalt), Hashed, hashed, unhashed) -import Data.Semigroup (Semigroup, (<>)) import Data.Text (Text) import GHC.Generics (Generic) +import qualified Data.IntSet as IS + #if defined(HAS_AESON) import Data.Aeson ((.:), (.=)) import qualified Data.Aeson as AE #endif - import Data.Text.AhoCorasick.Automaton (CaseSensitivity (..)) import qualified Data.Text.AhoCorasick.Automaton as Aho @@ -104,7 +106,7 @@ -- The caller is responsible that the needles are lower case in case the IgnoreCase -- is used for case sensitivity build :: CaseSensitivity -> [Text] -> Searcher () -build case_ = buildWithValues case_ . fmap (\x -> (x, ())) +build case_ = buildWithValues case_ . fmap (, ()) -- | The caller is responsible that the needles are lower case in case the IgnoreCase -- is used for case sensitivity @@ -152,5 +154,28 @@ -- On the first match, return True immediately. f _acc _match = Aho.Done True in case caseSensitivity searcher of - CaseSensitive -> Aho.runText False f (automaton searcher) text - IgnoreCase -> Aho.runLower False f (automaton searcher) text + CaseSensitive -> Aho.runText False f (automaton searcher) text + IgnoreCase -> Aho.runLower False f (automaton searcher) text + +-- | Build a 'Searcher' that returns the needle's index in the needle list when it matches. +buildNeedleIdSearcher :: CaseSensitivity -> [Text] -> Searcher Int +buildNeedleIdSearcher !case_ !ns = + buildWithValues case_ $ zip ns [0..] + +-- | Returns whether the haystack contains all of the needles. +-- This function expects the passed 'Searcher' to be constructed using 'buildNeedleIdAutomaton'. +containsAll :: Searcher Int -> Text -> Bool +containsAll !searcher !haystack = + let + initial = IS.fromDistinctAscList [0..numNeedles searcher - 1] + ac = automaton searcher + + f !acc (Aho.Match _index !needleId) + | IS.null acc' = Aho.Done acc' + | otherwise = Aho.Step acc' + where + !acc' = IS.delete needleId acc + + in IS.null $ case caseSensitivity searcher of + CaseSensitive -> Aho.runText initial f ac haystack + IgnoreCase -> Aho.runLower initial f ac haystack diff --git a/src/Data/Text/BoyerMoore/Automaton.hs b/src/Data/Text/BoyerMoore/Automaton.hs --- a/src/Data/Text/BoyerMoore/Automaton.hs +++ b/src/Data/Text/BoyerMoore/Automaton.hs @@ -6,10 +6,9 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} -{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE DerivingStrategies #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE ScopedTypeVariables #-} -- | An efficient implementation of the Boyer-Moore string search algorithm. @@ -22,19 +21,16 @@ -- The algorithm here can be potentially improved by including the Galil rule -- (https://en.wikipedia.org/wiki/Boyer%E2%80%93Moore_string-search_algorithm#The_Galil_rule) module Data.Text.BoyerMoore.Automaton - ( - Automaton - , CaseSensitivity (..) - , buildAutomaton - , runText - , runLower - , patternLength - , patternText - - , CodeUnitIndex (..) - , Next (..) - ) - where + ( Automaton + , CaseSensitivity (..) + , CodeUnitIndex (..) + , Next (..) + , buildAutomaton + , patternLength + , patternText + , runLower + , runText + ) where import Prelude hiding (length) @@ -43,18 +39,20 @@ import Control.Monad.ST (runST) import Data.Hashable (Hashable (..), Hashed, hashed, unhashed) import Data.Text.Internal (Text (..)) +import Data.TypedByteArray (Prim, TypedByteArray) import GHC.Generics (Generic) #if defined(HAS_AESON) import qualified Data.Aeson as AE #endif import qualified Data.HashMap.Strict as HashMap -import qualified Data.Vector.Unboxed as UVector -import qualified Data.Vector.Unboxed.Mutable as UMVector -import Data.Text.AhoCorasick.Automaton (CaseSensitivity (..), Next (..)) -import Data.Text.Utf16 (CodeUnit, CodeUnitIndex (..), lengthUtf16, lowerCodeUnit, lengthUtf16, unsafeIndexUtf16) +import Data.Text.AhoCorasick.Automaton (Next (..)) +import Data.Text.CaseSensitivity (CaseSensitivity (..)) +import Data.Text.Utf16 (CodeUnit, CodeUnitIndex (..), lengthUtf16, lowerCodeUnit, unsafeIndexUtf16) +import qualified Data.TypedByteArray as TBA + -- | A Boyer-Moore automaton is based on lookup-tables that allow skipping through the haystack. -- This allows for sub-linear matching in some cases, as we do not have to look at every input -- character. @@ -182,7 +180,7 @@ -- | The suffix table tells us for each character of the pattern how many characters we can -- jump ahead if the match fails at that point. -newtype SuffixTable = SuffixTable (UVector.Vector Int) +newtype SuffixTable = SuffixTable (TypedByteArray Int) deriving stock (Generic, Show) deriving anyclass (NFData) @@ -195,7 +193,7 @@ buildSuffixTable pattern = runST $ do let patLen = lengthUtf16 pattern - table <- UMVector.new $ codeUnitIndex patLen + table <- TBA.newTypedByteArray $ codeUnitIndex patLen let -- Case 1: For each position of the pattern we record the shift that would align the pattern so @@ -221,7 +219,7 @@ prefixIndex | isPrefix pattern (p + 1) = p + 1 | otherwise = lastPrefixIndex - UMVector.write table (codeUnitIndex p) (codeUnitIndex $ prefixIndex + patLen - 1 - p) + TBA.writeTypedByteArray table (codeUnitIndex p) (codeUnitIndex $ prefixIndex + patLen - 1 - p) init1 prefixIndex (p - 1) | otherwise = pure () @@ -234,21 +232,21 @@ let suffixLen = suffixLength pattern p when (indexCodePoint pattern (p - suffixLen) /= indexCodePoint pattern (patLen - 1 - suffixLen)) $ - UMVector.write table (codeUnitIndex $ patLen - 1 - suffixLen) (codeUnitIndex $ patLen - 1 - p + suffixLen) + TBA.writeTypedByteArray table (codeUnitIndex $ patLen - 1 - suffixLen) (codeUnitIndex $ patLen - 1 - p + suffixLen) init2 (p + 1) | otherwise = pure () init1 (patLen - 1) (patLen - 1) init2 0 - SuffixTable <$> UVector.unsafeFreeze table + SuffixTable <$> TBA.unsafeFreezeTypedByteArray table -- | The bad char table tells us how far we may skip ahead when encountering a certain character -- in the input string. For example, if there's a character that is not contained in the pattern at -- all, we can skip ahead until after that character. data BadCharTable = BadCharTable - { badCharTableAscii :: {-# UNPACK #-} !(UVector.Vector Int) + { badCharTableAscii :: {-# UNPACK #-} !(TypedByteArray Int) -- ^ The element type should be CodeUnitIndex, but there's no unboxed vector for that type, and -- defining it would be a lot of boilerplate. , badCharTableNonAscii :: !(HashMap.HashMap CodeUnit CodeUnitIndex) @@ -301,7 +299,7 @@ -- Initialize table with the maximum skip distance, which is the length of the pattern. -- This applies to all characters that are not part of the pattern. - asciiTable <- UMVector.replicate asciiCount $ codeUnitIndex patLen + asciiTable <- TBA.replicate asciiCount $ codeUnitIndex patLen let -- Fill the bad character table based on the rightmost occurrence of a character in the pattern. @@ -338,7 +336,7 @@ let patChar = indexCodePoint pattern i if fromIntegral patChar < asciiCount then do - UMVector.write asciiTable (fromIntegral patChar) (codeUnitIndex $ patLen - 1 - i) + TBA.writeTypedByteArray asciiTable (fromIntegral patChar) (codeUnitIndex $ patLen - 1 - i) fillTable (i + 1) nonAsciis else fillTable (i + 1) (HashMap.insert patChar (patLen - 1 - i) nonAsciis) @@ -347,7 +345,7 @@ nonAsciis <- fillTable 0 HashMap.empty - asciiTableFrozen <- UVector.unsafeFreeze asciiTable + asciiTableFrozen <- TBA.unsafeFreezeTypedByteArray asciiTable pure BadCharTable { badCharTableAscii = asciiTableFrozen @@ -359,9 +357,9 @@ -- Helper functions for easily toggling the safety of this module -- | Read from a lookup table at the specified index. -indexTable :: UVector.Unbox a => UVector.Vector a -> Int -> a +indexTable :: Prim a => TypedByteArray a -> Int -> a {-# INLINE indexTable #-} -indexTable = (UVector.!) -- UVector.unsafeIndex +indexTable = TBA.unsafeIndex -- | Read from a lookup table at the specified index. diff --git a/src/Data/Text/BoyerMoore/Replacer.hs b/src/Data/Text/BoyerMoore/Replacer.hs --- a/src/Data/Text/BoyerMoore/Replacer.hs +++ b/src/Data/Text/BoyerMoore/Replacer.hs @@ -14,7 +14,7 @@ where import Data.Text (Text) -import Data.Text.BoyerMoore.Automaton (CodeUnitIndex, Automaton) +import Data.Text.BoyerMoore.Automaton (Automaton, CodeUnitIndex) import qualified Data.Text as Text diff --git a/src/Data/Text/BoyerMoore/Searcher.hs b/src/Data/Text/BoyerMoore/Searcher.hs --- a/src/Data/Text/BoyerMoore/Searcher.hs +++ b/src/Data/Text/BoyerMoore/Searcher.hs @@ -9,18 +9,18 @@ {-# LANGUAGE FlexibleInstances #-} module Data.Text.BoyerMoore.Searcher - ( Searcher - , build - , buildWithValues - , needles - , numNeedles - , automata - , caseSensitivity - , containsAny - , setSearcherCaseSensitivity - ) - where - + ( Searcher + , automata + , build + , buildNeedleIdSearcher + , buildWithValues + , caseSensitivity + , containsAll + , containsAny + , needles + , numNeedles + , setSearcherCaseSensitivity + ) where import Control.DeepSeq (NFData) import Data.Bifunctor (first) @@ -123,3 +123,23 @@ any (\(automaton, ()) -> BoyerMoore.runText False f automaton text) (automata searcher) IgnoreCase -> any (\(automaton, ()) -> BoyerMoore.runLower False f automaton text) (automata searcher) + +-- | Build a 'Searcher' that returns the needle's index in the needle list when it matches. +buildNeedleIdSearcher :: CaseSensitivity -> [Text] -> Searcher Int +buildNeedleIdSearcher !case_ !ns = + buildWithValues case_ $ zip ns [0..] + +-- | Like 'containsAny', but checks whether all needles match instead. +-- Use 'buildNeedleIdSearcher' to get an appropriate 'Searcher'. +{-# NOINLINE containsAll #-} +containsAll :: Searcher Int -> Text -> Bool +containsAll !searcher !text = + let + -- On the first match, return True immediately. + f _acc _match = BoyerMoore.Done True + in + case caseSensitivity searcher of + CaseSensitive -> + all (\(automaton, _) -> BoyerMoore.runText False f automaton text) (automata searcher) + IgnoreCase -> + all (\(automaton, _) -> BoyerMoore.runLower False f automaton text) (automata searcher) diff --git a/src/Data/Text/CaseSensitivity.hs b/src/Data/Text/CaseSensitivity.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/CaseSensitivity.hs @@ -0,0 +1,22 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DerivingStrategies #-} + +module Data.Text.CaseSensitivity where + +import Control.DeepSeq (NFData) +import Data.Hashable (Hashable) +import GHC.Generics (Generic) +#if defined(HAS_AESON) +import Data.Aeson (FromJSON, ToJSON) +#endif +data CaseSensitivity + = CaseSensitive + | IgnoreCase + deriving stock (Eq, Generic, Show) +#if defined(HAS_AESON) + deriving anyclass (Hashable, NFData, FromJSON, ToJSON) +#else + deriving anyclass (Hashable, NFData) +#endif diff --git a/src/Data/Text/Utf16.hs b/src/Data/Text/Utf16.hs --- a/src/Data/Text/Utf16.hs +++ b/src/Data/Text/Utf16.hs @@ -10,23 +10,23 @@ {-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} --- | This module provides functions that allow treating Text values as series of Utf16 codepoints +-- | This module provides functions that allow treating Text values as series of UTF-16 codepoints -- instead of characters. module Data.Text.Utf16 - ( CodeUnit - , CodeUnitIndex (..) - , lengthUtf16 - , lowerUtf16 - , lowerCodeUnit - , upperUtf16 - , upperCodeUnit - , isCaseInvariant - , unpackUtf16 - , unsafeCutUtf16 - , unsafeSliceUtf16 - , unsafeIndexUtf16 - , indexTextArray - ) where + ( CodeUnit + , CodeUnitIndex (..) + , indexTextArray + , isCaseInvariant + , lengthUtf16 + , lowerCodeUnit + , lowerUtf16 + , unpackUtf16 + , unsafeCutUtf16 + , unsafeIndexUtf16 + , unsafeSliceUtf16 + , upperCodeUnit + , upperUtf16 + ) where import Prelude hiding (length) @@ -67,7 +67,7 @@ #endif --- | Return a Text as a list of UTF-16 code units. +-- | Return a 'Text' as a list of UTF-16 code units. {-# INLINABLE unpackUtf16 #-} unpackUtf16 :: Text -> [CodeUnit] unpackUtf16 (Text u16data offset length) = @@ -132,7 +132,7 @@ lengthUtf16 = CodeUnitIndex . TextUnsafe.lengthWord16 -- | Return the code unit (not character) with the given index. --- Note: The boudns are not checked. +-- Note: The bounds are not checked. unsafeIndexUtf16 :: Text -> CodeUnitIndex -> CodeUnit {-# INLINE unsafeIndexUtf16 #-} unsafeIndexUtf16 (Text arr off _) (CodeUnitIndex pos) = indexTextArray arr (pos + off) @@ -155,9 +155,9 @@ -- Differences from `Text.toLower` include code points in the BMP that lowercase -- to multiple code points, and code points outside of the BMP. -- --- For example, "Δ°" (U+0130), which `toLower` converts to "i" (U+0069, U+0307), +-- For example, \"Δ°\" (U+0130), which `toLower` converts to \"i\" (U+0069, U+0307), -- is converted into U+0069 only by `lowerUtf16`. --- Also, "π‘’’" (U+118A2), a code point from the Warang City writing system in the +-- Also, \"π‘’’\" (U+118A2), a code point from the Warang City writing system in the -- Supplementary Multilingual Plane, introduced in 2014 to Unicode 7. It would -- be lowercased to U+118C2 by `toLower`, but it is left untouched by -- `lowerUtf16`. @@ -201,14 +201,16 @@ -- also the test in AhoCorasickSpec.hs. | otherwise = fromIntegral $ Char.ord $ Char.toLower $ Char.chr $ fromIntegral cu +-- | Lowercase each individual code unit of a text without changing their index. +-- See also 'lowerUtf16' and 'lowerCodeUnit'. {-# INLINE upperUtf16 #-} upperUtf16 :: Text -> Text upperUtf16 = mapUtf16 upperCodeUnit +-- | Analogous to 'lowerCodeUnit'. {-# INLINE upperCodeUnit #-} upperCodeUnit :: CodeUnit -> CodeUnit upperCodeUnit cu - -- Analogous implementation to lowerCodeUnit | fromIntegral cu >= Char.ord 'a' && fromIntegral cu <= Char.ord 'z' = cu - fromIntegral (Char.ord 'a' - Char.ord 'A') | cu <= 127 = cu @@ -222,6 +224,7 @@ isCaseInvariant :: Text -> Bool isCaseInvariant = Text.all (\c -> Char.toLower c == Char.toUpper c) +-- | Retrieve a code unit from 'Text's internal representation. {-# INLINE indexTextArray #-} indexTextArray :: TextArray.Array -> Int -> CodeUnit indexTextArray array@(TextArray.Array byteArray) index diff --git a/src/Data/Text/Utf8.hs b/src/Data/Text/Utf8.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/Utf8.hs @@ -0,0 +1,433 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2022 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +-- | This module provides functions that allow treating 'Text' values as series of UTF-8 code units +-- instead of characters. Currently, it also contains a stub 'Text' type which treats its internal byte array +-- as UTF-8 encoded. We use this as a placeholder until we can use @text-2@. +module Data.Text.Utf8 + ( CodePoint + , CodeUnit + , CodeUnitIndex (..) + , Text (..) + , lengthUtf8 + , lowerCodePoint + , lowerUtf8 + , toLowerAscii + , unicode2utf8 + , unpackUtf8 + -- * Decoding + -- + -- $decoding + , decode2 + , decode3 + , decode4 + , decodeUtf8 + , stringToByteArray + -- * Indexing + -- + -- $indexing + , indexCodeUnit + , unsafeIndexCodePoint + , unsafeIndexCodePoint' + , unsafeIndexCodeUnit + , unsafeIndexCodeUnit' + -- * Slicing Functions + -- + -- $slicingFunctions + , unsafeCutUtf8 + , unsafeSliceUtf8 + -- * General Functions + -- + -- $generalFunctions + , Data.Text.Utf8.concat + , Data.Text.Utf8.dropWhile + , Data.Text.Utf8.null + , Data.Text.Utf8.readFile + , Data.Text.Utf8.replicate + , indices + , isInfixOf + , pack + , unpack + ) where + +import Control.DeepSeq (NFData, rnf) +import Data.Bits (Bits (shiftL), shiftR, (.&.), (.|.)) +import Data.Char (ord) +import Data.Foldable (for_) +import Data.Hashable (Hashable (hashWithSalt), hashByteArrayWithSalt) +import Data.Primitive.ByteArray (ByteArray (ByteArray), byteArrayFromList, compareByteArrays, + indexByteArray, newByteArray, sizeofByteArray, + unsafeFreezeByteArray, writeByteArray) +import Data.String (IsString (fromString)) +import Data.Word (Word8) +import GHC.Generics (Generic) +import Prelude hiding (length) +#if defined(HAS_AESON) +import Data.Aeson (FromJSON, ToJSON, Value (String), parseJSON, toJSON, withText) +#endif + +import qualified Data.ByteString as BS +import qualified Data.Char as Char +import qualified Data.Text as T + +-- | A UTF-8 code unit is a byte. A Unicode code point can be encoded as up to four code units. +type CodeUnit = Word8 + +-- | A Unicode code point. +type CodePoint = Char + +-- | An index into the raw UTF-8 data of a `Text`. This is not the code point +-- index as conventionally accepted by `Text`, so we wrap it to avoid confusing +-- the two. Incorrect index manipulation can lead to surrogate pairs being +-- sliced, so manipulate indices with care. This type is also used for lengths. +newtype CodeUnitIndex = CodeUnitIndex + { codeUnitIndex :: Int + } + deriving stock (Eq, Ord, Show, Generic, Bounded) +#if defined(HAS_AESON) + deriving newtype (Hashable, Num, NFData, FromJSON, ToJSON) +#else + deriving newtype (Hashable, Num, NFData) +#endif + +data Text + -- | A placeholder data type for UTF-8 encoded text until we can use text-2.0. + = Text + !ByteArray -- ^ Underlying array encoded using UTF-8. + !Int -- ^ Starting position of the UTF-8 sequence in bytes. + !Int -- ^ Length of the UTF-8 sequence in bytes. + +-- This instance, as well as the Show instance above, is necessary for the test suite. +instance Eq Text where + (Text u8data offset length) == (Text u8data' offset' length') = + length == length' && compareByteArrays u8data offset u8data' offset' length == EQ + +instance Ord Text where + compare (Text u8data offset length) (Text u8data' offset' length') = + compareByteArrays u8data offset u8data' offset' (min length length') <> compare length length' + +instance Show Text where + show = show . unpack + +-- Instances required for the Searcher modules etc. + +#if defined(HAS_AESON) +-- NOTE: This is ugly and slow but will be removed once we move to text-2.0. +instance ToJSON Text where + toJSON = String . T.pack . unpack + +instance FromJSON Text where + parseJSON = withText "Data.Text.Utf8.Text" (pure . pack . T.unpack) +#endif + +-- Copied from https://hackage.haskell.org/package/hashable-1.4.0.2/docs/src/Data.Hashable.Class.html#line-746 +instance Hashable Text where + hashWithSalt salt (Text (ByteArray arr) off len) = + hashByteArrayWithSalt arr (off `shiftL` 1) (len `shiftL` 1) (hashWithSalt salt len) + +instance NFData Text where + rnf (Text (ByteArray !_) !_ !_) = () + +instance IsString Text where + fromString = pack + +{-# INLINABLE unpackUtf8 #-} +unpackUtf8 :: Text -> [CodeUnit] +unpackUtf8 (Text u8data offset length) = + let + go _ 0 = [] + go i n = unsafeIndexCodeUnit' u8data (CodeUnitIndex i) : go (i + 1) (n - 1) + in + go offset length + +-- | The return value of this function is not really an index. +-- However the signature is supposed to make it clear that the length is returned in terms of code units, not code points. +lengthUtf8 :: Text -> CodeUnitIndex +lengthUtf8 (Text _ _ !length) = CodeUnitIndex length + +-- | Convert a 'Text' value into a 'T.Text' value. +toUtf16Text :: Text -> T.Text +toUtf16Text (Text u8data off len) = + T.unfoldr go 0 + where + go :: CodeUnitIndex -> Maybe (Char, CodeUnitIndex) + go i + | i >= CodeUnitIndex len = Nothing + | otherwise = + let + (codeUnits, codePoint) = unsafeIndexCodePoint' u8data $ CodeUnitIndex off + i + in + Just (codePoint, i + codeUnits) + +-- | Lower-case the ASCII code points A-Z and leave the rest of ASCII intact. +{-# INLINE toLowerAscii #-} +toLowerAscii :: Char -> Char +toLowerAscii cp + | Char.isAsciiUpper cp = Char.chr (Char.ord cp + 0x20) + | otherwise = cp + +-- TODO: Slow placeholder implementation until we can use text-2.0 +{-# INLINE lowerUtf8 #-} +lowerUtf8 :: Text -> Text +lowerUtf8 = pack . map lowerCodePoint . unpack + +asciiCount :: Int +asciiCount = 128 + +{-# INLINE lowerCodePoint #-} +-- | Lower-Case a UTF-8 codepoint. +-- Uses 'toLowerAscii' for ASCII and 'Char.toLower' otherwise. +lowerCodePoint :: Char -> Char +lowerCodePoint cp + | Char.ord cp < asciiCount = toLowerAscii cp + | otherwise = Char.toLower cp + +-- | Convert a Unicode Code Point 'c' into a list of UTF-8 code units (bytes). +unicode2utf8 :: (Ord a, Num a, Bits a) => a -> [a] +unicode2utf8 c + | c < 0x80 = [c] + | c < 0x800 = [0xc0 .|. (c `shiftR` 6), 0x80 .|. (0x3f .&. c)] + | c < 0x10000 = [0xe0 .|. (c `shiftR` 12), 0x80 .|. (0x3f .&. (c `shiftR` 6)), 0x80 .|. (0x3f .&. c)] + | otherwise = [0xf0 .|. (c `shiftR` 18), 0x80 .|. (0x3f .&. (c `shiftR` 12)), 0x80 .|. (0x3f .&. (c `shiftR` 6)), 0x80 .|. (0x3f .&. c)] + +-- $decoding +-- +-- Functions that turns code unit sequences into code point sequences. + +-- | Decode 2 UTF-8 code units into their code point. +-- The given code units should have the following format: +-- +-- > β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +-- > β”‚1 1 0 x x x x xβ”‚1 0 x x x x x xβ”‚ +-- > β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +{-# INLINE decode2 #-} +decode2 :: CodeUnit -> CodeUnit -> CodePoint +decode2 cu0 cu1 = + Char.chr $ (fromIntegral cu0 .&. 0x1f) `shiftL` 6 .|. fromIntegral cu1 .&. 0x3f + +-- | Decode 3 UTF-8 code units into their code point. +-- The given code units should have the following format: +-- +-- > β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +-- > β”‚1 1 1 0 x x x xβ”‚1 0 x x x x x xβ”‚1 0 x x x x x xβ”‚ +-- > β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +{-# INLINE decode3 #-} +decode3 :: CodeUnit -> CodeUnit -> CodeUnit -> CodePoint +decode3 cu0 cu1 cu2 = + Char.chr $ (fromIntegral cu0 .&. 0xf) `shiftL` 12 .|. (fromIntegral cu1 .&. 0x3f) `shiftL` 6 .|. (fromIntegral cu2 .&. 0x3f) + +-- | Decode 4 UTF-8 code units into their code point. +-- The given code units should have the following format: +-- +-- > β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +-- > β”‚1 1 1 1 0 x x xβ”‚1 0 x x x x x xβ”‚1 0 x x x x x xβ”‚1 0 x x x x x xβ”‚ +-- > β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +{-# INLINE decode4 #-} +decode4 :: CodeUnit -> CodeUnit -> CodeUnit -> CodeUnit -> CodePoint +decode4 cu0 cu1 cu2 cu3 = + Char.chr $ (fromIntegral cu0 .&. 0x7) `shiftL` 18 .|. (fromIntegral cu1 .&. 0x3f) `shiftL` 12 .|. (fromIntegral cu2 .&. 0x3f) `shiftL` 6 .|. (fromIntegral cu3 .&. 0x3f) + +-- | Decode a list of UTF-8 code units into a list of code points. +decodeUtf8 :: [CodeUnit] -> [CodePoint] +decodeUtf8 [] = [] +decodeUtf8 (cu0 : cus) | cu0 < 0xc0 = Char.chr (fromIntegral cu0) : decodeUtf8 cus +decodeUtf8 (cu0 : cu1 : cus) | cu0 < 0xe0 = decode2 cu0 cu1 : decodeUtf8 cus +decodeUtf8 (cu0 : cu1 : cu2 : cus) | cu0 < 0xf0 = decode3 cu0 cu1 cu2 : decodeUtf8 cus +decodeUtf8 (cu0 : cu1 : cu2 : cu3 : cus) | cu0 < 0xf8 = decode4 cu0 cu1 cu2 cu3 : decodeUtf8 cus +decodeUtf8 cus = error $ "Invalid UTF-8 input sequence at " ++ show (take 4 cus) + +stringToByteArray :: String -> ByteArray +stringToByteArray = byteArrayFromList . concatMap char2utf8 + -- See https://en.wikipedia.org/wiki/UTF-8 + where + char2utf8 :: Char -> [Word8] + char2utf8 = map fromIntegral . unicode2utf8 . ord + +-- $indexing +-- +-- 'Text' can be indexed by code units or code points. +-- A 'CodePoint' is a 21-bit Unicode code point and can consist of up to four code units. +-- A 'CodeUnit' is a single byte. + +-- | Decode a code point at the given 'CodeUnitIndex'. +-- Returns garbage if there is no valid code point at that position. +-- Does not perform bounds checking. +-- See 'decode2', 'decode3' and 'decode4' for the expected format of multi-byte code points. +{-# INLINE unsafeIndexCodePoint' #-} +unsafeIndexCodePoint' :: ByteArray -> CodeUnitIndex -> (CodeUnitIndex, CodePoint) +unsafeIndexCodePoint' !u8data (CodeUnitIndex !idx) + | cu0 < 0xc0 = (1, Char.chr $ fromIntegral cu0) + | cu0 < 0xe0 = (2, decode2 cu0 (cuAt 1)) + | cu0 < 0xf0 = (3, decode3 cu0 (cuAt 1) (cuAt 2)) + | otherwise = (4, decode4 cu0 (cuAt 1) (cuAt 2) (cuAt 3)) + where + cuAt !i = unsafeIndexCodeUnit' u8data $ CodeUnitIndex $ idx + i + !cu0 = cuAt 0 + +-- | Does exactly the same thing as 'unsafeIndexCodePoint'', but on 'Text' values. +{-# INLINE unsafeIndexCodePoint #-} +unsafeIndexCodePoint :: Text -> CodeUnitIndex -> (CodeUnitIndex, CodePoint) +unsafeIndexCodePoint (Text !u8data !off !_len) (CodeUnitIndex !index) = + unsafeIndexCodePoint' u8data $ CodeUnitIndex $ off + index + +-- | Get the code unit at the given 'CodeUnitIndex'. +-- Performs bounds checking. +{-# INLINE indexCodeUnit #-} +indexCodeUnit :: Text -> CodeUnitIndex -> CodeUnit +indexCodeUnit !text (CodeUnitIndex !index) + | index < 0 || index >= codeUnitIndex (lengthUtf8 text) = error $ "Index out of bounds " ++ show index + | otherwise = unsafeIndexCodeUnit text $ CodeUnitIndex index + +{-# INLINE unsafeIndexCodeUnit' #-} +unsafeIndexCodeUnit' :: ByteArray -> CodeUnitIndex -> CodeUnit +unsafeIndexCodeUnit' !u8data (CodeUnitIndex !idx) = indexByteArray u8data idx + +{-# INLINE unsafeIndexCodeUnit #-} +unsafeIndexCodeUnit :: Text -> CodeUnitIndex -> CodeUnit +unsafeIndexCodeUnit (Text !u8data !off !_len) (CodeUnitIndex !index) = + unsafeIndexCodeUnit' u8data $ CodeUnitIndex $ off + index + +-- $slicingFunctions +-- +-- 'unsafeCutUtf8' and 'unsafeSliceUtf8' are used to retrieve slices of 'Text' values. +-- @unsafeSliceUtf8 begin length@ returns a substring of length @length@ starting at @begin@. +-- @unsafeSliceUtf8 begin length@ returns a tuple of the "surrounding" substrings. +-- +-- They satisfy the following property: +-- +-- > let (prefix, suffix) = unsafeCutUtf8 begin length t +-- > in concat [prefix, unsafeSliceUtf8 begin length t, suffix] == t +-- +-- The following diagram visualizes the relevant offsets for @begin = CodeUnitIndex 2@, @length = CodeUnitIndex 6@ and @t = \"BCDEFGHIJKL\"@. +-- +-- > off off+len +-- > β”‚ β”‚ +-- > β–Ό β–Ό +-- > ──┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬── +-- > Aβ”‚Bβ”‚Cβ”‚Dβ”‚Eβ”‚Fβ”‚Gβ”‚Hβ”‚Iβ”‚Jβ”‚Kβ”‚Lβ”‚Mβ”‚N +-- > ──┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴── +-- > β–² β–² +-- > β”‚ β”‚ +-- > off+begin off+begin+length +-- > +-- > unsafeSliceUtf8 begin length t == "DEFGHI" +-- > unsafeCutUtf8 begin length t == ("BC", "JKL") +-- +-- The shown array is open at each end because in general, @t@ may be a slice as well. +-- +-- __WARNING__: As their name implies, these functions are not (necessarily) bounds-checked. Use at your own risk. + +-- TODO: Make this more readable once we have text-2.0. +unsafeCutUtf8 :: CodeUnitIndex -- ^ Starting position of substring. + -> CodeUnitIndex -- ^ Length of substring. + -> Text -- ^ Initial string. + -> (Text, Text) +unsafeCutUtf8 (CodeUnitIndex !begin) (CodeUnitIndex !length) (Text !u8data !off !len) = + ( Text u8data off begin + , Text u8data (off + begin + length) (len - begin - length) + ) + +-- TODO: Make this more readable once we have text-2.0. +unsafeSliceUtf8 :: CodeUnitIndex -> CodeUnitIndex -> Text -> Text +unsafeSliceUtf8 (CodeUnitIndex !begin) (CodeUnitIndex !length) (Text !u8data !off !_len) = + Text u8data (off + begin) length + +-- $generalFunctions +-- +-- These functions are available in @text@ as well and should be removed once this library moves to @text-2@. +-- You should be able to use these by doing @import qualified Data.Text.Utf8 as Text@ just like you would with @text@. +-- +-- NOTE: The 'Text' instances for @Show@, @Eq@, @Ord@, @IsString@, @FromJSON@, @ToJSON@ and @Hashable@ in this file also fall in this category. + +-- | TODO: Inefficient placeholder implementation. +concat :: [Text] -> Text +concat = pack . concatMap unpack + +-- | See 'Data.Text.dropWhile'. +dropWhile :: (Char -> Bool) -> Text -> Text +dropWhile predicate text = + let + len = codeUnitIndex (lengthUtf8 text) + go i + | i >= CodeUnitIndex len = i + | otherwise = + let + (codeUnits, codePoint) = unsafeIndexCodePoint text i + in + if predicate codePoint then + go $ i + codeUnits + else + i + + prefixEnd = go 0 + in + unsafeSliceUtf8 prefixEnd (CodeUnitIndex len - prefixEnd) text + +-- | Checks whether a text is the empty string. +null :: Text -> Bool +null (Text _ _ len) = len == 0 + +-- | TODO: Inefficient placeholder implementation. +pack :: String -> Text +pack = go . stringToByteArray + where + go !arr = Text arr 0 $ sizeofByteArray arr + +-- | TODO: Inefficient placeholder implementation. +-- See 'Data.Text.replicate' +replicate :: Int -> Text -> Text +replicate n = pack . Prelude.concat . Prelude.replicate n . unpack + +-- | TODO: Inefficient placeholder implementation. +-- This function implements very basic string search. It's @text@ counterpart is 'Data.Text.Internal.Search.indices', which implements the Boyer-Moore algorithm. +-- Since we have this function only to check whether our own Boyer-Moore implementation works, it would not make much sense to implement it using the same algorithm. +-- Once we can use @text-2@, we can compare our implementation to the official @text@ one which presumably works. +indices :: Text -> Text -> [Int] +indices needle haystack + | needleLen == 0 = [] + | otherwise = go 0 0 + where + needleLen = lengthUtf8 needle + haystackLen = lengthUtf8 haystack + + go startIdx needleIdx + -- needle is longer than remaining haystack + | startIdx + needleLen > haystackLen = [] + -- whole needle matched + | needleIdx >= needleLen = codeUnitIndex startIdx : go (startIdx + needleLen) 0 + -- charachter mismatch + | needleCp /= haystackCp = go (startIdx + 1) 0 + -- advance + | otherwise = go startIdx $ needleIdx + codeUnits + where + (codeUnits, needleCp) = unsafeIndexCodePoint needle needleIdx + (_, haystackCp) = unsafeIndexCodePoint haystack $ startIdx + needleIdx + +-- | TODO: Inefficient placeholder implementation. +isInfixOf :: Text -> Text -> Bool +isInfixOf needle haystack = T.isInfixOf (toUtf16Text needle) (toUtf16Text haystack) + +-- | See 'Data.Text.IO.readFile'. +-- TODO: Uses 'Data.ByteString.readFile' and loops through each byte individually. +-- Use 'Data.Primitive.Ptr.copyPtrToMutableByteArray' here if possible. +readFile :: FilePath -> IO Text +readFile path = do + contents <- BS.readFile path + array <- newByteArray $ BS.length contents + for_ [0..BS.length contents - 1] $ \i -> do + writeByteArray array i $ BS.index contents i + array' <- unsafeFreezeByteArray array + pure $ Text array' 0 $ BS.length contents + +-- | TODO: Inefficient placeholder implementation. +unpack :: Text -> String +unpack = decodeUtf8 . unpackUtf8 diff --git a/src/Data/Text/Utf8/AhoCorasick/Automaton.hs b/src/Data/Text/Utf8/AhoCorasick/Automaton.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/Utf8/AhoCorasick/Automaton.hs @@ -0,0 +1,545 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2022 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | An efficient implementation of the Aho-Corasick string matching algorithm. +-- See http://web.stanford.edu/class/archive/cs/cs166/cs166.1166/lectures/02/Small02.pdf +-- for a good explanation of the algorithm. +-- +-- The memory layout of the automaton, and the function that steps it, were +-- optimized to the point where string matching compiles roughly to a loop over +-- the code units in the input text, that keeps track of the current state. +-- Lookup of the next state is either just an array index (for the root state), +-- or a linear scan through a small array (for non-root states). The pointer +-- chases that are common for traversing Haskell data structures have been +-- eliminated. +-- +-- The construction of the automaton has not been optimized that much, because +-- construction time is usually negligible in comparison to matching time. +-- Therefore construction is a two-step process, where first we build the +-- automaton as int maps, which are convenient for incremental construction. +-- Afterwards we pack the automaton into unboxed vectors. +-- +-- This module is a rewrite of the previous version which used an older version of +-- the 'text' package which in turn used UTF-16 internally. +module Data.Text.Utf8.AhoCorasick.Automaton + ( AcMachine (..) + , CaseSensitivity (..) + , CodeUnitIndex (..) + , Match (..) + , Next (..) + , build + , debugBuildDot + , runLower + , runText + , runWithCase + ) where + +import Control.DeepSeq (NFData) +import Data.Bits (Bits (shiftL, shiftR, (.&.), (.|.))) +import Data.Char (chr) +import Data.Foldable (foldl') +import Data.IntMap.Strict (IntMap) +import Data.Word (Word32, Word64) +import GHC.Generics (Generic) + +import qualified Data.Char as Char +import qualified Data.IntMap.Strict as IntMap +import qualified Data.List as List +import qualified Data.Vector as Vector + +import Data.Text.CaseSensitivity (CaseSensitivity (..)) +import Data.Text.Utf8 (CodePoint, CodeUnitIndex (CodeUnitIndex), Text (..)) +import Data.TypedByteArray (Prim, TypedByteArray) + +import qualified Data.Text.Utf8 as Utf8 +import qualified Data.TypedByteArray as TBA + +-- TYPES +-- | A numbered state in the Aho-Corasick automaton. +type State = Int + +-- | A transition is a pair of (code point, next state). The code point is 21 bits, +-- and the state index is 32 bits. The code point is stored in +-- the least significant 32 bits, with the special value 2^21 indicating a +-- wildcard; the "failure" transition. Bits 22 through 31 (starting from zero, +-- both bounds inclusive) are always 0. +-- +-- +-- > Bit 63 (most significant) Bit 0 (least significant) +-- > | | +-- > v v +-- > |<-- goto state -->|<-- 0s -->| |<-- input -->| +-- > |SSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSS|0000000000|W|IIIIIIIIIIIIIIIIIIIII| +-- > | +-- > Wildcard bit (bit 21) +-- +-- If you change this representation, make sure to update 'transitionCodeUnit', +-- 'wildcard', 'transitionState', 'transitionIsWildcard', 'newTransition' and +-- 'newWildcardTransition' as well. Those functions form the interface used to +-- construct and read transitions. +type Transition = Word64 + +type Offset = Word32 + +data Match v = Match + { matchPos :: {-# UNPACK #-} !CodeUnitIndex + -- ^ The code unit index past the last code unit of the match. Note that this + -- is not a code *point* (Haskell `Char`) index; a code point might be encoded + -- as up to four code units. + , matchValue :: v + -- ^ The payload associated with the matched needle. + } + +-- | An Aho-Corasick automaton. +data AcMachine v = AcMachine + { machineValues :: !(Vector.Vector [v]) + -- ^ For every state, the values associated with its needles. If the state is + -- not a match state, the list is empty. + , machineTransitions :: !(TypedByteArray Transition) + -- ^ A packed vector of transitions. For every state, there is a slice of this + -- vector that starts at the offset given by `machineOffsets`, and ends at the + -- first wildcard transition. + , machineOffsets :: !(TypedByteArray Offset) + -- ^ For every state, the index into `machineTransitions` where the transition + -- list for that state starts. + , machineRootAsciiTransitions :: !(TypedByteArray Transition) + -- ^ A lookup table for transitions from the root state, an optimization to + -- avoid having to walk all transitions, at the cost of using a bit of + -- additional memory. + } deriving (Generic) + +instance NFData v => NFData (AcMachine v) + +-- AUTOMATON CONSTRUCTION + +-- | The wildcard value is 2^21, one more than the maximal 21-bit code point. +wildcard :: Integral a => a +wildcard = 0x200000 + +-- | Extract the code unit from a transition. The special wildcard transition +-- will return 0. +transitionCodeUnit :: Transition -> CodePoint +transitionCodeUnit t = Char.chr $ fromIntegral (t .&. 0x1fffff) + +-- | Extract the goto state from a transition. +transitionState :: Transition -> State +transitionState t = fromIntegral (t `shiftR` 32) + +-- | Test if the transition is not for a specific code unit, but the wildcard +-- transition to take if nothing else matches. +transitionIsWildcard :: Transition -> Bool +transitionIsWildcard t = (t .&. wildcard) == wildcard + +newTransition :: CodePoint -> State -> Transition +newTransition input state = + let + input64 = fromIntegral $ Char.ord input :: Word64 + state64 = fromIntegral state :: Word64 + in + (state64 `shiftL` 32) .|. input64 + +newWildcardTransition :: State -> Transition +newWildcardTransition state = + let + state64 = fromIntegral state :: Word64 + in + (state64 `shiftL` 32) .|. wildcard + +-- | Pack transitions for each state into one contiguous array. In order to find +-- the transitions for a specific state, we also produce a vector of start +-- indices. All transition lists are terminated by a wildcard transition, so +-- there is no need to record the length. +packTransitions :: [[Transition]] -> (TypedByteArray Transition, TypedByteArray Offset) +packTransitions transitions = + let + packed = TBA.fromList $ concat transitions + offsets = TBA.fromList $ map fromIntegral $ scanl (+) 0 $ fmap List.length transitions + in + (packed, offsets) + +-- | Construct an Aho-Corasick automaton for the given needles. +-- The automaton uses Unicode code points to match the input. +build :: [(Text, v)] -> AcMachine v +build needlesWithValues = + let + -- Construct the Aho-Corasick automaton using IntMaps, which are a suitable + -- representation when building the automaton. We use int maps rather than + -- hash maps to ensure that the iteration order is the same as that of a + -- vector. + (numStates, transitionMap, initialValueMap) = buildTransitionMap needlesWithValues + fallbackMap = buildFallbackMap transitionMap + valueMap = buildValueMap transitionMap fallbackMap initialValueMap + + -- Convert the map of transitions, and the map of fallback states, into a + -- list of transition lists, where every transition list is terminated by + -- a wildcard transition to the fallback state. + prependTransition ts input state = newTransition (Char.chr input) state : ts + makeTransitions fallback ts = IntMap.foldlWithKey' prependTransition [newWildcardTransition fallback] ts + transitionsList = zipWith makeTransitions (IntMap.elems fallbackMap) (IntMap.elems transitionMap) + + -- Pack the transition lists into one contiguous array, and build the lookup + -- table for the transitions from the root state. + (transitions, offsets) = packTransitions transitionsList + rootTransitions = buildAsciiTransitionLookupTable $ transitionMap IntMap.! 0 + values = Vector.generate numStates (valueMap IntMap.!) + in + AcMachine values transitions offsets rootTransitions + +-- | Build the automaton, and format it as Graphviz Dot, for visual debugging. +debugBuildDot :: [Text] -> String +debugBuildDot needles = + let + (_numStates, transitionMap, initialValueMap) = + buildTransitionMap $ zip needles ([0..] :: [Int]) + fallbackMap = buildFallbackMap transitionMap + valueMap = buildValueMap transitionMap fallbackMap initialValueMap + + dotEdge extra state nextState = + " " ++ show state ++ " -> " ++ show nextState ++ " [" ++ extra ++ "];" + + dotFallbackEdge :: [String] -> State -> State -> [String] + dotFallbackEdge edges state nextState = + dotEdge "style = dashed" state nextState : edges + + dotTransitionEdge :: State -> [String] -> Int -> State -> [String] + dotTransitionEdge state edges input nextState = + dotEdge ("label = \"" ++ showInput input ++ "\"") state nextState : edges + + showInput input = [chr input] + + prependTransitionEdges edges state = + IntMap.foldlWithKey' (dotTransitionEdge state) edges (transitionMap IntMap.! state) + + dotMatchState :: [String] -> State -> [Int] -> [String] + dotMatchState edges _ [] = edges + dotMatchState edges state _ = (" " ++ show state ++ " [shape = doublecircle];") : edges + + dot0 = foldBreadthFirst prependTransitionEdges [] transitionMap + dot1 = IntMap.foldlWithKey' dotFallbackEdge dot0 fallbackMap + dot2 = IntMap.foldlWithKey' dotMatchState dot1 valueMap + in + -- Set rankdir = "LR" to prefer a left-to-right graph, rather than top to + -- bottom. I have dual widescreen monitors and I don't use them in portrait + -- mode. Reverse the instructions because order affects node lay-out, and by + -- prepending we built up a reversed list. + unlines $ ["digraph {", " rankdir = \"LR\";"] ++ reverse dot2 ++ ["}"] + +-- Different int maps that are used during constuction of the automaton. The +-- transition map represents the trie of states, the fallback map contains the +-- fallback (or "failure" or "suffix") edge for every state. +type TransitionMap = IntMap (IntMap State) +type FallbackMap = IntMap State +type ValuesMap v = IntMap [v] + +-- | Build the trie of the Aho-Corasick state machine for all input needles. +buildTransitionMap :: forall v. [(Text, v)] -> (Int, TransitionMap, ValuesMap v) +buildTransitionMap = + let + -- | Inserts a single needle into the given transition and values map. + insertNeedle :: (Int, TransitionMap, ValuesMap v) -> (Text, v) -> (Int, TransitionMap, ValuesMap v) + insertNeedle !acc (!needle, !value) = go stateInitial 0 acc + where + !needleLen = Utf8.lengthUtf8 needle + + go !state !index (!numStates, !transitions, !values) + -- End of the current needle, insert the associated payload value. + -- If a needle occurs multiple times, then at this point we will merge + -- their payload values, so the needle is reported twice, possibly with + -- different payload values. + | index >= needleLen = (numStates, transitions, IntMap.insertWith (++) state [value] values) + go !state !index (!numStates, !transitions, !values) = + let + !transitionsFromState = transitions IntMap.! state + (!codeUnits, !input) = Utf8.unsafeIndexCodePoint needle index + in + case IntMap.lookup (Char.ord input) transitionsFromState of + -- Transition already exists, follow it and continue from there. + Just !nextState -> + go nextState (index + codeUnits) (numStates, transitions, values) + -- Transition for input does not exist at state: + -- Allocate a new state, and insert a transition to it. + -- Also insert an empty transition map for it. + Nothing -> + let + !nextState = numStates + !transitionsFromState' = IntMap.insert (Char.ord input) nextState transitionsFromState + !transitions' + = IntMap.insert state transitionsFromState' + $ IntMap.insert nextState IntMap.empty transitions + in + go nextState (index + codeUnits) (numStates + 1, transitions', values) + + -- Initially, the root state (state 0) exists, and it has no transitions + -- to anywhere. + stateInitial = 0 + initialTransitions = IntMap.singleton stateInitial IntMap.empty + initialValues = IntMap.empty + in + foldl' insertNeedle (1, initialTransitions, initialValues) + +-- Size of the ascii transition lookup table. +asciiCount :: Integral a => a +asciiCount = 128 + +-- | Build a lookup table for the first 128 code points, that can be used for +-- O(1) lookup of a transition, rather than doing a linear scan over all +-- transitions. The fallback goes back to the initial state, state 0. +{-# NOINLINE buildAsciiTransitionLookupTable #-} +buildAsciiTransitionLookupTable :: IntMap State -> TypedByteArray Transition +buildAsciiTransitionLookupTable transitions = TBA.generate asciiCount $ \i -> + case IntMap.lookup i transitions of + Just state -> newTransition (Char.chr i) state + Nothing -> newWildcardTransition 0 + +-- | Traverse the state trie in breadth-first order. +foldBreadthFirst :: (a -> State -> a) -> a -> TransitionMap -> a +foldBreadthFirst f seed transitions = go [0] [] seed + where + -- For the traversal, we keep a queue of states to vitit. Every iteration we + -- take one off the front, and all states reachable from there get added to + -- the back. Rather than using a list for this, we use the functional + -- amortized queue to avoid O(nΒ²) append. This makes a measurable difference + -- when the backlog can grow large. In one of our benchmark inputs for + -- example, we have roughly 160 needles that are 10 characters each (but + -- with some shared prefixes), and the backlog size grows to 148 during + -- construction. Construction time goes down from ~0.80 ms to ~0.35 ms by + -- using the amortized queue. + -- See also section 3.1.1 of Purely Functional Data Structures by Okasaki + -- https://www.cs.cmu.edu/~rwh/theses/okasaki.pdf. + go [] [] !acc = acc + go [] revBacklog !acc = go (reverse revBacklog) [] acc + go (state : backlog) revBacklog !acc = + let + -- Note that the backlog never contains duplicates, because we traverse + -- a trie that only branches out. For every state, there is only one + -- path from the root that leads to it. + extra = IntMap.elems $ transitions IntMap.! state + in + go backlog (extra ++ revBacklog) (f acc state) + +-- | Determine the fallback transition for every state, by traversing the +-- transition trie breadth-first. +buildFallbackMap :: TransitionMap -> FallbackMap +buildFallbackMap transitions = + let + -- Suppose that in state `state`, there is a transition for input `input` + -- to state `nextState`, and we already know the fallback for `state`. Then + -- this function returns the fallback state for `nextState`. + getFallback :: FallbackMap -> State -> Int -> State + -- All the states after the root state (state 0) fall back to the root state. + getFallback _ 0 _ = 0 + getFallback fallbacks !state !input = + let + fallback = fallbacks IntMap.! state + transitionsFromFallback = transitions IntMap.! fallback + in + case IntMap.lookup input transitionsFromFallback of + Just st -> st + Nothing -> getFallback fallbacks fallback input + + insertFallback :: State -> FallbackMap -> Int -> State -> FallbackMap + insertFallback !state fallbacks !input !nextState = + IntMap.insert nextState (getFallback fallbacks state input) fallbacks + + insertFallbacks :: FallbackMap -> State -> FallbackMap + insertFallbacks fallbacks !state = + IntMap.foldlWithKey' (insertFallback state) fallbacks (transitions IntMap.! state) + in + foldBreadthFirst insertFallbacks (IntMap.singleton 0 0) transitions + +-- | Determine which matches to report at every state, by traversing the +-- transition trie breadth-first, and appending all the matches from a fallback +-- state to the matches for the current state. +buildValueMap :: forall v. TransitionMap -> FallbackMap -> ValuesMap v -> ValuesMap v +buildValueMap transitions fallbacks valuesInitial = + let + insertValues :: ValuesMap v -> State -> ValuesMap v + insertValues values !state = + let + fallbackValues = values IntMap.! (fallbacks IntMap.! state) + valuesForState = case IntMap.lookup state valuesInitial of + Just vs -> vs ++ fallbackValues + Nothing -> fallbackValues + in + IntMap.insert state valuesForState values + in + foldBreadthFirst insertValues (IntMap.singleton 0 []) transitions + +-- Define aliases for array indexing so we can turn bounds checks on and off +-- in one place. We ran this code with `Vector.!` (bounds-checked indexing) in +-- production for two months without failing the bounds check, so we have turned +-- the check off for performance now. +{-# INLINE at #-} +at :: forall a. Vector.Vector a -> Int -> a +at = Vector.unsafeIndex + +{-# INLINE uAt #-} +uAt :: Prim a => TypedByteArray a -> Int -> a +uAt = TBA.unsafeIndex + +-- RUNNING THE MACHINE + +-- | Result of handling a match: stepping the automaton can exit early by +-- returning a `Done`, or it can continue with a new accumulator with `Step`. +data Next a = Done !a | Step !a + +-- | Run the automaton, possibly lowercasing the input text on the fly if case +-- insensitivity is desired. See also `runLower`. +-- +-- The code of this function itself is organized as a state machine as well. +-- Each state in the diagram below corresponds to a function defined in +-- `runWithCase`. These functions are written in a way such that GHC identifies them +-- as [join points](https://www.microsoft.com/en-us/research/publication/compiling-without-continuations/). +-- This means that they can be compiled to jumps instead of function calls, which helps performance a lot. +-- +-- @ +-- β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +-- β”‚ β”‚ +-- β”Œβ”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +-- β”‚consumeInputβ”œβ”€β”€β”€β–ΊlookupTransitionβ”œβ”€β”€β”€β–ΊcollectMatchesβ”‚ +-- β””β”€β–²β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β–²β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”˜ +-- β”‚ β”‚ β”‚ β”‚ +-- β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ +-- β”‚ β”‚ +-- β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +-- @ +-- +-- * @consumeInput@ decodes a code point of up to four code units and possibly lowercases it. +-- It passes this code point to @followCodePoint@, which in turn calls @lookupTransition@. +-- * @lookupTransition@ checks whether the given code point matches any transitions at the given state. +-- If so, it follows the transition and calls @collectMatches@. Otherwise, it follows the fallback transition +-- and calls @followCodePoint@ or @consumeInput@. +-- * @collectMatches@ checks whether the current state is accepting and updates the accumulator accordingly. +-- Afterwards it loops back to @consumeInput@. +-- +-- NOTE: @followCodePoint@ is actually inlined into @consumeInput@ by GHC. +-- It is included in the diagram for illustrative reasons only. +-- +-- All of these functions have the arguments @offset@, @remaining@, @state@ and @acc@ which encode the current input +-- position and the accumulator, which contains the matches. If you change any of the functions above, +-- make sure to check the Core dumps afterwards that @offset@, @remaining@ and @state@ were turned +-- into unboxed @Int#@ by GHC. If any of them aren't, the program will constantly allocate and deallocate heap space for them. +-- You can nudge GHC in the right direction by using bang patterns on these arguments. +-- +-- WARNING: Run benchmarks when modifying this function; its performance is +-- fragile. It took many days to discover the current formulation which compiles +-- to fast code; removing the wrong bang pattern could cause a 10% performance +-- regression. +{-# INLINE runWithCase #-} +runWithCase :: forall a v. CaseSensitivity -> a -> (a -> Match v -> Next a) -> AcMachine v -> Text -> a +runWithCase !caseSensitivity !seed !f !machine !text = + consumeInput initialOffset initialRemaining seed initialState + where + initialState = 0 + + Text !u8data !off !len = text + AcMachine !values !transitions !offsets !rootAsciiTransitions = machine + + !initialOffset = CodeUnitIndex off + !initialRemaining = CodeUnitIndex len + + -- NOTE: All of the arguments are strict here, because we want to compile + -- them down to unpacked variables on the stack, or even registers. + + -- When we follow an edge, we look in the transition table and do a + -- linear scan over all transitions until we find the right one, or + -- until we hit the wildcard transition at the end. For 0 or 1 or 2 + -- transitions that is fine, but the initial state often has more + -- transitions, so we have a dedicated lookup table for it, that takes + -- up a bit more space, but provides O(1) lookup of the next state. We + -- only do this for the first 128 code units (all of ascii). + + -- | Consume a code unit sequence that constitutes a full code point. + -- If the code unit at @offset@ is ASCII, we can lower it using 'Utf8.toLowerAscii'. + {-# NOINLINE consumeInput #-} + consumeInput :: CodeUnitIndex -> CodeUnitIndex -> a -> State -> a + consumeInput !_offset 0 !acc !_state = acc + consumeInput !offset !remaining !acc !state = + followCodePoint (offset + codeUnits) (remaining - codeUnits) acc possiblyLoweredCp state + + where + (!codeUnits, !cp) = Utf8.unsafeIndexCodePoint' u8data offset + + !possiblyLoweredCp = case caseSensitivity of + CaseSensitive -> cp + IgnoreCase -> Utf8.lowerCodePoint cp + + {-# INLINE followCodePoint #-} + followCodePoint :: CodeUnitIndex -> CodeUnitIndex -> a -> CodePoint -> State -> a + followCodePoint !offset !remaining !acc !cp !state + | state == initialState && Char.ord cp < asciiCount = lookupRootAsciiTransition offset remaining acc cp + | otherwise = lookupTransition offset remaining acc cp state $ offsets `uAt` state + + -- NOTE: This function can't be inlined since it is self-recursive. + {-# NOINLINE lookupTransition #-} + lookupTransition :: CodeUnitIndex -> CodeUnitIndex -> a -> CodePoint -> State -> Offset -> a + lookupTransition !offset !remaining !acc !cp !state !i + -- There is no transition for the given input. Follow the fallback edge, + -- and try again from that state, etc. If we are in the base state + -- already, then nothing matched, so move on to the next input. + | transitionIsWildcard t = + if state == initialState + then consumeInput offset remaining acc state + else followCodePoint offset remaining acc cp (transitionState t) + -- We found the transition, switch to that new state, possibly matching the rest of cus. + -- NOTE: This comes after wildcard checking, because the code unit of + -- the wildcard transition is 0, which is a valid input. + | transitionCodeUnit t == cp = + collectMatches offset remaining acc (transitionState t) + -- The transition we inspected is not for the current input, and it is not + -- a wildcard either; look at the next transition then. + | otherwise = + lookupTransition offset remaining acc cp state $ i + 1 + + where + !t = transitions `uAt` fromIntegral i + + -- NOTE: there is no `state` argument here, because this case applies only + -- to the root state `stateInitial`. + {-# INLINE lookupRootAsciiTransition #-} + lookupRootAsciiTransition !offset !remaining !acc !cp + -- Given code unit does not match at root ==> Repeat at offset from initial state + | transitionIsWildcard t = consumeInput offset remaining acc initialState + -- Transition matched! + | otherwise = collectMatches offset remaining acc $ transitionState t + where !t = rootAsciiTransitions `uAt` Char.ord cp + + {-# NOINLINE collectMatches #-} + collectMatches !offset !remaining !acc !state = + let + matchedValues = values `at` state + -- Fold over the matched values. If at any point the user-supplied fold + -- function returns `Done`, then we early out. Otherwise continue. + handleMatch !acc' vs = case vs of + [] -> consumeInput offset remaining acc' state + v:more -> case f acc' (Match (offset - initialOffset) v) of + Step newAcc -> handleMatch newAcc more + Done finalAcc -> finalAcc + in + handleMatch acc matchedValues + +-- NOTE: To get full advantage of inlining this function, you probably want to +-- compile the compiling module with -fllvm and the same optimization flags as +-- this module. +{-# INLINE runText #-} +runText :: forall a v. a -> (a -> Match v -> Next a) -> AcMachine v -> Text -> a +runText = runWithCase CaseSensitive + +-- Finds all matches in the lowercased text. This function lowercases the input text +-- on the fly to avoid allocating a second lowercased text array. It is still the +-- responsibility of the caller to lowercase the needles. Needles that contain +-- uppercase code points will not match. +-- +-- NOTE: To get full advantage of inlining this function, you probably want to +-- compile the compiling module with -fllvm and the same optimization flags as +-- this module. +{-# INLINE runLower #-} +runLower :: forall a v. a -> (a -> Match v -> Next a) -> AcMachine v -> Text -> a +runLower = runWithCase IgnoreCase diff --git a/src/Data/Text/Utf8/AhoCorasick/Replacer.hs b/src/Data/Text/Utf8/AhoCorasick/Replacer.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/Utf8/AhoCorasick/Replacer.hs @@ -0,0 +1,219 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2019 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DerivingStrategies #-} + +-- | Implements sequential string replacements based on the Aho-Corasick algorithm. +module Data.Text.Utf8.AhoCorasick.Replacer + ( -- * State machine + Needle + , Payload (..) + , Replacement + , Replacer (..) + , build + , compose + , run + , runWithLimit + ) where + +import Control.DeepSeq (NFData) +import Data.Hashable (Hashable) +import Data.List (sort) +import Data.Maybe (fromJust) +import GHC.Generics (Generic) + +#if defined(HAS_AESON) +import qualified Data.Aeson as AE +#endif + +import Data.Text.CaseSensitivity (CaseSensitivity (..)) +import Data.Text.Utf8 (CodeUnitIndex, Text) +import Data.Text.Utf8.AhoCorasick.Searcher (Searcher) + +import qualified Data.Text.Utf8 as Utf8 +import qualified Data.Text.Utf8.AhoCorasick.Automaton as Aho +import qualified Data.Text.Utf8.AhoCorasick.Searcher as Searcher + +-- | Descriptive type alias for strings to search for. +type Needle = Text + +-- | Descriptive type alias for replacements. +type Replacement = Text + +-- | Priority of a needle. Higher integers indicate higher priorities. +-- Replacement order is such that all matches of priority p are replaced before +-- replacing any matches of priority q where p > q. +type Priority = Int + +data Payload = Payload + { needlePriority :: {-# UNPACK #-} !Priority + , needleLength :: {-# UNPACK #-} !CodeUnitIndex + , needleReplacement :: !Replacement + } +#if defined(HAS_AESON) + deriving (Eq, Generic, Hashable, NFData, Show, AE.FromJSON, AE.ToJSON) +#else + deriving (Eq, Generic, Hashable, NFData, Show) +#endif + +-- | A state machine used for efficient replacements with many different needles. +data Replacer = Replacer + { replacerCaseSensitivity :: CaseSensitivity + , replacerSearcher :: Searcher Payload + } + deriving stock (Show, Eq, Generic) +#if defined(HAS_AESON) + deriving (Hashable, NFData, AE.FromJSON, AE.ToJSON) +#else + deriving (Hashable, NFData) +#endif + +-- | Build an Aho-Corasick automaton that can be used for performing fast +-- sequential replaces. +-- +-- Case-insensitive matching performs per-letter language-agnostic lower-casing. +-- Therefore, it will work in most cases, but not in languages where lower-casing +-- depends on the context of the character in question. +-- +-- We need to revisit this algorithm when we want to implement full Unicode +-- support. +build :: CaseSensitivity -> [(Needle, Replacement)] -> Replacer +build caseSensitivity replaces = Replacer caseSensitivity searcher + where + searcher = Searcher.buildWithValues caseSensitivity $ zipWith mapNeedle [0..] replaces + mapNeedle i (needle, replacement) = + let + needle' = case caseSensitivity of + CaseSensitive -> needle + IgnoreCase -> Utf8.lowerUtf8 needle + in + -- Note that we negate i: earlier needles have a higher priority. We + -- could avoid it and define larger integers to be lower priority, but + -- that made the terminology in this module very confusing. + (needle', Payload (-i) (Utf8.lengthUtf8 needle') replacement) + +-- | Return the composition `replacer2` after `replacer1`, if they have the same +-- case sensitivity. If the case sensitivity differs, Nothing is returned. +compose :: Replacer -> Replacer -> Maybe Replacer +compose (Replacer case1 searcher1) (Replacer case2 searcher2) + | case1 /= case2 = Nothing + | otherwise = + let + -- Replace the priorities of the second machine, so they all come after + -- the first. + renumber i (needle, Payload _ len replacement) = (needle, Payload (-i) len replacement) + needles1 = Searcher.needles searcher1 + needles2 = Searcher.needles searcher2 + searcher = Searcher.buildWithValues case1 $ zipWith renumber [0..] (needles1 ++ needles2) + in + Just $ Replacer case1 searcher + +-- A match collected while running replacements. It is isomorphic to the Match +-- reported by the automaton, but the data is arranged in a more useful way: +-- as the start index and length of the match, and the replacement. +data Match = Match !CodeUnitIndex !CodeUnitIndex !Text deriving (Eq, Ord, Show) + +-- | Apply replacements of all matches. Assumes that the matches are ordered by +-- match position, and that no matches overlap. +replace :: [Match] -> Text -> Text +replace matches haystack = Utf8.concat $ go 0 matches haystack + where + -- At every match, cut the string into three pieces, removing the match. + -- Because a Text is a buffer pointer and (offset, length), cutting does not + -- involve string copies. Only at the very end we piece together the strings + -- again, so Text can allocate a buffer of the right length and memcpy the + -- parts into the new target string. + -- If `k` is a code unit index into the original text, then `k - offset` + -- is an index into `remainder`. In other words, `offset` is the index into + -- the original text where `remainder` starts. + go :: CodeUnitIndex -> [Match] -> Text -> [Text] + go !_offset [] remainder = [remainder] + go !offset ((Match pos len replacement) : ms) remainder = + let + (prefix, suffix) = Utf8.unsafeCutUtf8 (pos - offset) len remainder + in + prefix : replacement : go (pos + len) ms suffix + +-- | Compute the length of the string resulting from applying the replacements. +replacementLength :: [Match] -> Text -> CodeUnitIndex +replacementLength matches initial = go matches (Utf8.lengthUtf8 initial) + where + go [] !acc = acc + go (Match _ matchLen repl : rest) !acc = go rest (acc - matchLen + Utf8.lengthUtf8 repl) + +-- | Given a list of matches sorted on start position, remove matches that start +-- within an earlier match. +removeOverlap :: [Match] -> [Match] +removeOverlap matches = case matches of + [] -> [] + [m] -> [m] + (m0@(Match pos0 len0 _) : m1@(Match pos1 _ _) : ms) -> + if pos1 >= pos0 + len0 + then m0 : removeOverlap (m1:ms) + else removeOverlap (m0:ms) + +-- | When we iterate through all matches, keep track only of the matches with +-- the highest priority: those are the ones that we will replace first. If we +-- find multiple matches with that priority, remember all of them. If we find a +-- match with lower priority, ignore it, because we already have a more +-- important match. Also, if the priority is `threshold` or higher, ignore the +-- match, so we can exclude matches if we already did a round of replacements +-- for that priority. This way we don't have to build a new automaton after +-- every round of replacements. +{-# INLINE prependMatch #-} +prependMatch :: Priority -> (Priority, [Match]) -> Aho.Match Payload -> Aho.Next (Priority, [Match]) +prependMatch !threshold (!pBest, !matches) (Aho.Match pos (Payload pMatch len replacement)) + | pMatch < threshold && pMatch > pBest = Aho.Step (pMatch, [Match (pos - len) len replacement]) + | pMatch < threshold && pMatch == pBest = Aho.Step (pMatch, Match (pos - len) len replacement : matches) + | otherwise = Aho.Step (pBest, matches) + +run :: Replacer -> Text -> Text +run replacer = fromJust . runWithLimit replacer maxBound + +{-# NOINLINE runWithLimit #-} +runWithLimit :: Replacer -> CodeUnitIndex -> Text -> Maybe Text +runWithLimit (Replacer case_ searcher) maxLength = go initialThreshold + where + !automaton = Searcher.automaton searcher + + -- Priorities are 0 or lower, so an initial threshold of 1 keeps all + -- matches. + !initialThreshold = 1 + + -- Needle priorities go from 0 for the highest priority to (-numNeedles + 1) + -- for the lowest priority. That means that if we find a match with + -- minPriority, we don't need to do another pass afterwards, because there + -- are no remaining needles. + !minPriority = 1 - Searcher.numNeedles searcher + + go :: Priority -> Text -> Maybe Text + go !threshold haystack = + let + seed = (minBound :: Priority, []) + matchesWithPriority = case case_ of + CaseSensitive -> Aho.runText seed (prependMatch threshold) automaton haystack + IgnoreCase -> Aho.runLower seed (prependMatch threshold) automaton haystack + in + case matchesWithPriority of + -- No match at the given threshold, there is nothing left to do. + -- Return the input string unmodified. + (_, []) -> Just haystack + -- We found matches at priority p. Remove overlapping matches, then + -- apply all replacements. Next, we need to go again, this time + -- considering only needles with a lower priority than p. As an + -- optimization (which matters mainly for the single needle case), + -- if we find a match at the lowest priority, we don't need another + -- pass. Note that if in `rawMatches` we find only matches of priority + -- p > minPriority, then we do still need another pass, because the + -- replacements could create new matches. + (p, matches) + | replacementLength matches haystack > maxLength -> Nothing + | p == minPriority -> Just $ replace (removeOverlap $ sort matches) haystack + | otherwise -> go p $ replace (removeOverlap $ sort matches) haystack diff --git a/src/Data/Text/Utf8/AhoCorasick/Searcher.hs b/src/Data/Text/Utf8/AhoCorasick/Searcher.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/Utf8/AhoCorasick/Searcher.hs @@ -0,0 +1,178 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2022 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} + +module Data.Text.Utf8.AhoCorasick.Searcher + ( Searcher + , automaton + , build + , buildNeedleIdSearcher + , buildWithValues + , caseSensitivity + , containsAll + , containsAny + , needles + , numNeedles + , setSearcherCaseSensitivity + ) where + +import Control.DeepSeq (NFData) +import Data.Hashable (Hashable (hashWithSalt), Hashed, hashed, unhashed) +import GHC.Generics (Generic) + +#if defined(HAS_AESON) +import Data.Aeson ((.:), (.=)) +import qualified Data.Aeson as AE +#endif + +import qualified Data.IntSet as IS + +import Data.Text.CaseSensitivity (CaseSensitivity (..)) +import Data.Text.Utf8 (Text) + +import qualified Data.Text.Utf8.AhoCorasick.Automaton as Aho + +-- | A set of needles with associated values, and an Aho-Corasick automaton to +-- efficiently find those needles. +-- +-- INVARIANT: searcherAutomaton = Aho.build . searcherNeedles +-- To enforce this invariant, the fields are not exposed from this module. +-- There is a separate constructor function. +-- +-- The purpose of this wrapper is to have a type that is Hashable and Eq, so we +-- can derive those for types that embed the searcher, whithout requiring the +-- automaton itself to be Hashable or Eq, which would be both wasteful and +-- tedious. Because the automaton is fully determined by the needles and +-- associated values, it is sufficient to implement Eq and Hashable in terms of +-- the needles only. +-- +-- We also use Hashed to cache the hash of the needles. +data Searcher v = Searcher + { searcherCaseSensitive :: CaseSensitivity + , searcherNeedles :: Hashed [(Text, v)] + , searcherNumNeedles :: Int + , searcherAutomaton :: Aho.AcMachine v + } deriving (Generic) + +#if defined(HAS_AESON) +instance AE.ToJSON v => AE.ToJSON (Searcher v) where + toJSON s = AE.object + [ "needles" .= needles s + , "caseSensitivity" .= caseSensitivity s + ] + +instance (Hashable v, AE.FromJSON v) => AE.FromJSON (Searcher v) where + parseJSON = AE.withObject "Searcher" $ \o -> buildWithValues <$> o .: "caseSensitivity" <*> o .: "needles" +#endif + +instance Show (Searcher v) where + show _ = "Searcher _ _ _" + +instance Hashable v => Hashable (Searcher v) where + hashWithSalt salt searcher = hashWithSalt salt $ searcherNeedles searcher + {-# INLINE hashWithSalt #-} + +instance Eq v => Eq (Searcher v) where + -- Since we store the length of the needle list anyway, + -- we can use it to early out if there is a length mismatch. + Searcher cx xs nx _ == Searcher cy ys ny _ = (nx, xs, cx) == (ny, ys, cy) + {-# INLINE (==) #-} + +instance NFData v => NFData (Searcher v) + +-- NOTE: Although we could implement Semigroup for every v by just concatenating +-- needle lists, we don't, because this might lead to unexpected results. For +-- example, if v is (Int, a) where the Int is a priority, combining two +-- searchers might want to discard priorities, concatenate the needle lists, and +-- reassign priorities, rather than concatenating the needle lists as-is and +-- possibly having duplicate priorities in the resulting searcher. +instance Semigroup (Searcher ()) where + x <> y + | caseSensitivity x == caseSensitivity y + = buildWithValues (searcherCaseSensitive x) (needles x <> needles y) + | otherwise = error "Combining searchers of different case sensitivity" + {-# INLINE (<>) #-} + +-- | Builds the Searcher for a list of needles +-- The caller is responsible that the needles are lower case in case the IgnoreCase +-- is used for case sensitivity +build :: CaseSensitivity -> [Text] -> Searcher () +build case_ = buildWithValues case_ . fmap (, ()) + +-- | The caller is responsible that the needles are lower case in case the IgnoreCase +-- is used for case sensitivity +buildWithValues :: Hashable v => CaseSensitivity -> [(Text, v)] -> Searcher v +{-# INLINABLE buildWithValues #-} +buildWithValues case_ ns = + Searcher case_ (hashed ns) (length ns) $ Aho.build ns + +needles :: Searcher v -> [(Text, v)] +needles = unhashed . searcherNeedles + +numNeedles :: Searcher v -> Int +numNeedles = searcherNumNeedles + +automaton :: Searcher v -> Aho.AcMachine v +automaton = searcherAutomaton + +caseSensitivity :: Searcher v -> CaseSensitivity +caseSensitivity = searcherCaseSensitive + +-- | Updates the case sensitivity of the searcher. Does not change the +-- capitilization of the needles. The caller should be certain that if IgnoreCase +-- is passed, the needles are already lower case. +setSearcherCaseSensitivity :: CaseSensitivity -> Searcher v -> Searcher v +setSearcherCaseSensitivity case_ searcher = searcher{ + searcherCaseSensitive = case_ + } + +-- | Return whether the haystack contains any of the needles. +-- Case sensitivity depends on the properties of the searcher +-- This function is marked noinline as an inlining boundary. Aho.runText is +-- marked inline, so this function will be optimized to report only whether +-- there is a match, and not construct a list of matches. We don't want this +-- function be inline, to make sure that the conditions of the caller don't +-- affect how this function is optimized. There is little to gain from +-- additional inlining. The pragma is not an optimization in itself, rather it +-- is a defence against fragile optimizer decisions. +{-# NOINLINE containsAny #-} +containsAny :: Searcher () -> Text -> Bool +containsAny !searcher !text = + let + -- On the first match, return True immediately. + f _acc _match = Aho.Done True + in case caseSensitivity searcher of + CaseSensitive -> Aho.runText False f (automaton searcher) text + IgnoreCase -> Aho.runLower False f (automaton searcher) text + +-- | Build a 'Searcher' that returns the needle's index in the needle list when it matches. +buildNeedleIdSearcher :: CaseSensitivity -> [Text] -> Searcher Int +buildNeedleIdSearcher !case_ !ns = + buildWithValues case_ $ zip ns [0..] + +-- | Returns whether the haystack contains all of the needles. +-- This function expects the passed 'Searcher' to be constructed using 'buildNeedleIdAutomaton'. +containsAll :: Searcher Int -> Text -> Bool +containsAll !searcher !haystack = + let + initial = IS.fromDistinctAscList [0..numNeedles searcher - 1] + ac = automaton searcher + + f !acc (Aho.Match _index !needleId) + | IS.null acc' = Aho.Done acc' + | otherwise = Aho.Step acc' + where + !acc' = IS.delete needleId acc + + in IS.null $ case caseSensitivity searcher of + CaseSensitive -> Aho.runText initial f ac haystack + IgnoreCase -> Aho.runLower initial f ac haystack diff --git a/src/Data/Text/Utf8/AhoCorasick/Splitter.hs b/src/Data/Text/Utf8/AhoCorasick/Splitter.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/Utf8/AhoCorasick/Splitter.hs @@ -0,0 +1,189 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2019 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} + +-- | Splitting strings using Aho–Corasick. +module Data.Text.Utf8.AhoCorasick.Splitter + ( Splitter + , automaton + , build + , separator + , split + , splitIgnoreCase + , splitReverse + , splitReverseIgnoreCase + ) where + +import Control.DeepSeq (NFData (..)) +import Data.Function (on) +import Data.Hashable (Hashable (..)) +import Data.List.NonEmpty (NonEmpty ((:|))) +import Data.Text.Utf8 (Text) + +#if defined(HAS_AESON) +import qualified Data.Aeson as AE +#endif + +import qualified Data.List.NonEmpty as NonEmpty + +import Data.Text.Utf8.AhoCorasick.Automaton (AcMachine) + +import qualified Data.Text.Utf8 as Utf8 +import qualified Data.Text.Utf8.AhoCorasick.Automaton as Aho + +-------------------------------------------------------------------------------- +-- Splitter + +-- | Build a splitter once, then use it many times! +data Splitter = + Splitter + { splitterAutomaton :: AcMachine () -- INVARIANT: Exactly one needle. + , splitterSeparator :: Text -- INVARIANT: Equivalent to needle. + } + +#if defined(HAS_AESON) +instance AE.ToJSON Splitter where + toJSON = AE.toJSON . separator + +instance AE.FromJSON Splitter where + parseJSON v = build <$> AE.parseJSON v +#endif + +-- | Construct a splitter with a single separator. +{-# INLINE build #-} +build :: Text -> Splitter +build sep = + let !auto = Aho.build [(sep, ())] in + Splitter auto sep + +-- | Get the automaton that would be used for finding separators. +{-# INLINE automaton #-} +automaton :: Splitter -> AcMachine () +automaton = splitterAutomaton + +-- | What is the separator we are splitting on? +{-# INLINE separator #-} +separator :: Splitter -> Text +separator = splitterSeparator + +-- | Split the given string into strings separated by the separator. +-- +-- If the order of the results is not important, use the faster function +-- 'splitReverse' instead. +{-# INLINE split #-} +split :: Splitter -> Text -> NonEmpty Text +split = (NonEmpty.reverse .) . splitReverse + +-- | Split the given string into strings separated by the separator. +-- +-- If the order of the results is not important, use the faster function +-- 'splitReverseIgnoreCase' instead. +-- +-- The separator is matched case-insensitively, but the splitter must have been +-- constructed with a lowercase needle. +{-# INLINE splitIgnoreCase #-} +splitIgnoreCase :: Splitter -> Text -> NonEmpty Text +splitIgnoreCase = (NonEmpty.reverse .) . splitReverseIgnoreCase + +-- | Like 'split', but return the substrings in reverse order. +{-# INLINE splitReverse #-} +splitReverse :: Splitter -> Text -> NonEmpty Text +splitReverse s t = + finalizeAccum $ + Aho.runText + (zeroAccum (separator s) t) + stepAccum + (automaton s) + t + +-- | Like 'splitIgnoreCase', but return the substrings in reverse order. +{-# INLINE splitReverseIgnoreCase #-} +splitReverseIgnoreCase :: Splitter -> Text -> NonEmpty Text +splitReverseIgnoreCase s t = + finalizeAccum $ + Aho.runLower + (zeroAccum (separator s) t) + stepAccum + (automaton s) + t + +-------------------------------------------------------------------------------- +-- Fold + +-- | The accumulator is used as state when processing the matches from left to +-- right. While the matches are fed to us ordered by end offset, all matches +-- have the same length because there is only one needle. +data Accum = + Accum + { _accumSepLen :: !Aho.CodeUnitIndex -- ^ Length of separator. + , _accumHaystack :: !Text -- ^ Haystack to slice off of. + , accumResult :: ![Text] -- ^ Match-separated strings. + , accumPrevEnd :: !Aho.CodeUnitIndex -- ^ Offset at end of last match. + } + +-- | Finalizing the accumulator does more than just 'accumResult', hence this +-- is a separate function. +{-# INLINE finalizeAccum #-} +finalizeAccum :: Accum -> NonEmpty Text +finalizeAccum (Accum _ hay res prevEnd) = + -- Once we have processed all the matches, there is still the substring after + -- the final match. This substring is always included in the result, even + -- when there were no matches. Hence we can return a non-empty list. + let !str = Utf8.unsafeSliceUtf8 prevEnd (Utf8.lengthUtf8 hay - prevEnd) hay in + str :| res + +-- | The initial accumulator begins at the begin of the haystack. +{-# INLINE zeroAccum #-} +zeroAccum :: Text -> Text -> Accum +zeroAccum sep hay = Accum (Utf8.lengthUtf8 sep) hay [] 0 + +-- | Step the accumulator using the next match. Overlapping matches will be +-- ignored. Overlapping matches may occur when the separator has a non-empty +-- prefix that is also a suffix. +{-# INLINE stepAccum #-} +stepAccum :: Accum -> Aho.Match v -> Aho.Next Accum +stepAccum acc@(Accum sepLen hay res prevEnd) (Aho.Match sepEnd _) + + -- When the match begins before the current offset, it overlaps a match that + -- we processed before, and so we ignore it. + | sepEnd - sepLen < prevEnd = + Aho.Step acc + + -- The match is behind the current offset, so we slice the haystack until the + -- begin of the match and include that as a result. + | otherwise = + let !str = Utf8.unsafeSliceUtf8 prevEnd (sepEnd - sepLen - prevEnd) hay in + Aho.Step acc { accumResult = str : res, accumPrevEnd = sepEnd } + +-------------------------------------------------------------------------------- +-- Instances + +instance Eq Splitter where + {-# INLINE (==) #-} + (==) = (==) `on` separator + +instance Ord Splitter where + {-# INLINE compare #-} + compare = compare `on` separator + +instance Hashable Splitter where + {-# INLINE hashWithSalt #-} + hashWithSalt salt searcher = + salt `hashWithSalt` separator searcher + +instance NFData Splitter where + {-# INLINE rnf #-} + rnf (Splitter searcher sepLength) = + rnf searcher `seq` + rnf sepLength + +instance Show Splitter where + showsPrec p splitter = + showParen (p > 10) $ + showString "build " . + showsPrec 11 (separator splitter) diff --git a/src/Data/Text/Utf8/BoyerMoore/Automaton.hs b/src/Data/Text/Utf8/BoyerMoore/Automaton.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/Utf8/BoyerMoore/Automaton.hs @@ -0,0 +1,335 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2019 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | An efficient implementation of the Boyer-Moore string search algorithm. +-- http://www-igm.univ-mlv.fr/~lecroq/string/node14.html#SECTION00140 +-- https://en.wikipedia.org/wiki/Boyer%E2%80%93Moore_string-search_algorithm +-- +-- This module contains a almost 1:1 translation from the C example code in the +-- wikipedia article. +-- +-- The algorithm here can be potentially improved by including the Galil rule +-- (https://en.wikipedia.org/wiki/Boyer%E2%80%93Moore_string-search_algorithm#The_Galil_rule) +module Data.Text.Utf8.BoyerMoore.Automaton + ( Automaton + , CaseSensitivity (..) + , CodeUnitIndex (..) + , Next (..) + , buildAutomaton + , patternLength + , patternText + , runText + ) where + +import Prelude hiding (length) + +import Control.DeepSeq (NFData) +import Control.Monad (when) +import Control.Monad.ST (runST) +import Data.Hashable (Hashable (..), Hashed, hashed, unhashed) +import GHC.Generics (Generic) + +#if defined(HAS_AESON) +import qualified Data.Aeson as AE +#endif + +import Data.Text.AhoCorasick.Automaton (Next (..)) +import Data.Text.CaseSensitivity (CaseSensitivity (..)) +import Data.Text.Utf8 (CodeUnit, CodeUnitIndex (..), Text) +import Data.TypedByteArray (Prim, TypedByteArray) + +import qualified Data.Text.Utf8 as Utf8 +import qualified Data.TypedByteArray as TBA + +-- | A Boyer-Moore automaton is based on lookup-tables that allow skipping through the haystack. +-- This allows for sub-linear matching in some cases, as we do not have to look at every input +-- character. +-- +-- NOTE: Unlike the AcMachine, a Boyer-Moore automaton only returns non-overlapping matches. +-- This means that a Boyer-Moore automaton is not a 100% drop-in replacement for Aho-Corasick. +-- +-- Returning overlapping matches would degrade the performance to /O(nm)/ in pathological cases like +-- finding @aaaa@ in @aaaaa....aaaaaa@ as for each match it would scan back the whole /m/ characters +-- of the pattern. +data Automaton = Automaton + { automatonPattern :: Hashed Text + , automatonSuffixTable :: SuffixTable + , automatonBadCharTable :: BadCharTable + } + deriving stock (Generic, Show) + deriving anyclass (NFData) + +instance Hashable Automaton where + hashWithSalt salt (Automaton pattern _ _) = hashWithSalt salt pattern + +instance Eq Automaton where + (Automaton pat1 _ _) == (Automaton pat2 _ _) = pat1 == pat2 + +#if defined(HAS_AESON) +instance AE.FromJSON Automaton where + parseJSON v = buildAutomaton <$> AE.parseJSON v + +instance AE.ToJSON Automaton where + toJSON = AE.toJSON . unhashed . automatonPattern +#endif + +buildAutomaton :: Text -> Automaton +buildAutomaton pattern = Automaton (hashed pattern) (buildSuffixTable pattern) (buildBadCharTable pattern) + +-- | Finds all matches in the text, calling the match callback with the *first* +-- matched character of each match of the pattern. +-- +-- NOTE: This is unlike Aho-Corasick, which reports the index of the character +-- right after a match. +-- +-- NOTE: In the UTF-16 version of this module, there is a function 'Data.Text.BoyerMoore.Automaton.runLower' +-- which does lower-case matching. This function does not exist for the UTF-8 version since it is very +-- tricky to skip code points going backwards without preprocessing the whole input first. +-- +-- NOTE: To get full advantage of inlining this function, you probably want to +-- compile the compiling module with -fllvm and the same optimization flags as +-- this module. +runText :: forall a + . a + -> (a -> CodeUnitIndex -> Next a) + -> Automaton + -> Text + -> a +{-# INLINE runText #-} +runText seed f automaton text + | patLen == 0 = seed + | otherwise = go seed (patLen - 1) + where + Automaton patternHashed suffixTable badCharTable = automaton + -- Use needle as identifier since pattern is potentially a keyword + needle = unhashed patternHashed + patLen = Utf8.lengthUtf8 needle + stringLen = Utf8.lengthUtf8 text + + codeUnitAt = Utf8.unsafeIndexCodeUnit text + + {-# INLINE go #-} + go result haystackIndex + | haystackIndex < stringLen = matchLoop result haystackIndex (patLen - 1) + | otherwise = result + + -- Compare the needle back-to-front with the haystack + matchLoop result haystackIndex needleIndex + | needleIndex >= 0 && codeUnitAt haystackIndex == Utf8.unsafeIndexCodeUnit needle needleIndex = + -- Characters match, try the pair before + matchLoop result (haystackIndex - 1) (needleIndex - 1) + -- We found a match (all needle characters matched) + | needleIndex < 0 = + case f result (haystackIndex + 1) of + Done final -> final + -- `haystackIndex` now points to the character just before the match starts + -- Adding `patLen` once points to the last character of the match, + -- Adding `patLen` once more points to the earliest character where + -- we can find a non-overlapping match. + Step intermediate -> go intermediate (haystackIndex + 2 * patLen) + -- We know it's not a match, the characters differ at the current position + | otherwise = + let + -- The bad character table tells us how far we can advance to the right so that the + -- character at the current position in the input string, where matching failed, + -- is lined up with it's rightmost occurrence in the needle. + -- Note: we could end up left of were we started, essentially never making progress, + -- if we were to use this rule alone. + badCharSkip = badCharLookup badCharTable (codeUnitAt haystackIndex) + suffixSkip = suffixLookup suffixTable needleIndex + skip = max badCharSkip suffixSkip + in + go result (haystackIndex + skip) + +-- | Length of the matched pattern measured in UTF-8 code units (bytes). +patternLength :: Automaton -> CodeUnitIndex +patternLength = Utf8.lengthUtf8 . patternText + +-- | Return the pattern that was used to construct the automaton. +patternText :: Automaton -> Text +patternText (Automaton pattern _ _) = unhashed pattern + +-- | The suffix table tells us for each character of the pattern how many characters we can +-- jump ahead if the match fails at that point. +newtype SuffixTable = SuffixTable (TypedByteArray Int) + deriving stock (Generic, Show) + deriving anyclass (NFData) + +-- | Lookup an entry in the suffix table. +suffixLookup :: SuffixTable -> CodeUnitIndex -> CodeUnitIndex +{-# INLINE suffixLookup #-} +suffixLookup (SuffixTable table) = CodeUnitIndex . indexTable table . codeUnitIndex + +buildSuffixTable :: Text -> SuffixTable +buildSuffixTable pattern = runST $ do + let patLen = Utf8.lengthUtf8 pattern + + table <- TBA.newTypedByteArray $ codeUnitIndex patLen + + let + -- Case 1: For each position of the pattern we record the shift that would align the pattern so + -- that it starts at the longest suffix that is at the same time a prefix, if a mismatch would + -- happen at that position. + -- + -- Suppose the length of the pattern is n, a mismatch occurs at position i in the pattern and j + -- in the haystack, then we know that pattern[i+1..n] == haystack[j+1..j+n-i]. That is, we know + -- that the part of the haystack that we already matched is a suffix of the pattern. + -- If the pattern happens to have a prefix that is equal to or a shorter suffix of that matched + -- suffix, we can shift the pattern to the right so that the pattern starts at the longest + -- suffix that we have seen that conincides with a prefix of the pattern. + -- + -- Consider the pattern `ababa`. Then we get + -- + -- p: 0 1 2 3 4 + -- Pattern: a b a b a + -- lastPrefixIndex: 2 2 4 4 5 + -- table: 6 5 6 5 5 + init1 lastPrefixIndex p + | p >= 0 = do + let + prefixIndex + | isPrefix pattern (p + 1) = p + 1 + | otherwise = lastPrefixIndex + TBA.writeTypedByteArray table (codeUnitIndex p) (codeUnitIndex $ prefixIndex + patLen - 1 - p) + init1 prefixIndex (p - 1) + | otherwise = pure () + + -- Case 2: We also have to account for the fact that the matching suffix of the pattern might + -- occur again somewhere within the pattern. In that case, we may not shift as far as if it was + -- a prefix. That is why the `init2` loop is run after `init1`, potentially overwriting some + -- entries with smaller shifts. + init2 p + | p < patLen - 1 = do + let + suffixLen = suffixLength pattern p + when (Utf8.unsafeIndexCodeUnit pattern (p - suffixLen) /= Utf8.unsafeIndexCodeUnit pattern (patLen - 1 - suffixLen)) $ + TBA.writeTypedByteArray table (codeUnitIndex $ patLen - 1 - suffixLen) (codeUnitIndex $ patLen - 1 - p + suffixLen) + init2 (p + 1) + | otherwise = pure () + + init1 (patLen - 1) (patLen - 1) + init2 0 + + SuffixTable <$> TBA.unsafeFreezeTypedByteArray table + + +-- | The bad char table tells us how far we may skip ahead when encountering a certain character +-- in the input string. For example, if there's a character that is not contained in the pattern at +-- all, we can skip ahead until after that character. +data BadCharTable = BadCharTable + { badCharTableEntries :: {-# UNPACK #-} !(TypedByteArray Int) + -- ^ The element type should be CodeUnitIndex, but there's no unboxed vector for that type, and + -- defining it would be a lot of boilerplate. + , badCharTablePatternLen :: CodeUnitIndex + } + deriving stock (Generic, Show) + deriving anyclass (NFData) + +-- | Number of entries in the fixed-size lookup-table of the bad char table. +badcharTableSize :: Int +{-# INLINE badcharTableSize #-} +badcharTableSize = 256 + +-- | Lookup an entry in the bad char table. +badCharLookup :: BadCharTable -> CodeUnit -> CodeUnitIndex +{-# INLINE badCharLookup #-} +badCharLookup (BadCharTable asciiTable _patLen) char = CodeUnitIndex $ indexTable asciiTable intChar + where + intChar = fromIntegral char + +-- | True if the suffix of the @pattern@ starting from @pos@ is a prefix of the pattern +-- For example, @isPrefix \"aabbaa\" 4 == True@. +isPrefix :: Text -> CodeUnitIndex -> Bool +isPrefix needle pos = go 0 + where + suffixLen = Utf8.lengthUtf8 needle - pos + go i + | i < suffixLen = + -- FIXME: Check whether implementing the linter warning kills tco + if Utf8.unsafeIndexCodeUnit needle i == Utf8.unsafeIndexCodeUnit needle (pos + i) + then go (i + 1) + else False + | otherwise = True + +-- | Length of the longest suffix of the pattern ending on @pos@. +-- For example, @suffixLength \"abaacbbaac\" 4 == 4@, because the substring \"baac\" ends at position +-- 4 and is at the same time the longest suffix that does so, having length 4. +suffixLength :: Text -> CodeUnitIndex -> CodeUnitIndex +suffixLength pattern pos = go 0 + where + patLen = Utf8.lengthUtf8 pattern + go i + | Utf8.unsafeIndexCodeUnit pattern (pos - i) == Utf8.unsafeIndexCodeUnit pattern (patLen - 1 - i) && i < pos = go (i + 1) + | otherwise = i + +buildBadCharTable :: Text -> BadCharTable +buildBadCharTable pattern = runST $ do + let patLen = Utf8.lengthUtf8 pattern + + -- Initialize table with the maximum skip distance, which is the length of the pattern. + -- This applies to all characters that are not part of the pattern. + asciiTable <- TBA.replicate badcharTableSize $ codeUnitIndex patLen + + let + -- Fill the bad character table based on the rightmost occurrence of a character in the pattern. + -- Note that there is also a variant of Boyer-Moore that records all positions (see Wikipedia, + -- but that requires even more storage space). + -- Also note that we exclude the last character of the pattern when building the table. + -- This is because + -- + -- 1. If the last character does not occur anywhere else in the pattern and we encounter it + -- during a mismatch, we can advance the pattern to just after that character: + -- + -- Haystack: aaadcdabcdbb + -- Pattern: abcd + -- + -- In the above example, we would match `d` and `c`, but then fail because `d` != `b`. + -- Since `d` only occurs at the very last position of the pattern, we can shift to + -- + -- Haystack: aaadcdabcdbb + -- Pattern: abcd + -- + -- 2. If it does occur anywhere else in the pattern, we can only shift as far as it's necessary + -- to align it with the haystack: + -- + -- Haystack: aaadddabcdbb + -- Pattern: adcd + -- + -- We match `d`, and then there is a mismatch `d` != `c`, which allows us to shift only up to: + + -- Haystack: aaadddabcdbb + -- Pattern: adcd + fillTable !i + -- for(i = 0; i < patLen - 1; i++) { + | i < patLen - 1 = do + let patChar = Utf8.unsafeIndexCodeUnit pattern i + TBA.writeTypedByteArray asciiTable (fromIntegral patChar) (codeUnitIndex $ patLen - 1 - i) + fillTable (i + 1) + | otherwise = pure () + + fillTable 0 + + asciiTableFrozen <- TBA.unsafeFreezeTypedByteArray asciiTable + + pure BadCharTable + { badCharTableEntries = asciiTableFrozen + , badCharTablePatternLen = patLen + } + + +-- Helper functions for easily toggling the safety of this module + +-- | Read from a lookup table at the specified index. +indexTable :: Prim a => TypedByteArray a -> Int -> a +{-# INLINE indexTable #-} +indexTable = TBA.unsafeIndex diff --git a/src/Data/Text/Utf8/BoyerMoore/Replacer.hs b/src/Data/Text/Utf8/BoyerMoore/Replacer.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/Utf8/BoyerMoore/Replacer.hs @@ -0,0 +1,91 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2019 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE FlexibleInstances #-} + +module Data.Text.Utf8.BoyerMoore.Replacer + ( -- Replacer + replaceSingleLimited + ) where + +import Data.Text.Utf8 (Text) +import Data.Text.Utf8.BoyerMoore.Automaton (Automaton, CodeUnitIndex) + +import qualified Data.Text.Utf8 as Text +import qualified Data.Text.Utf8 as Utf8 +import qualified Data.Text.Utf8.BoyerMoore.Automaton as BoyerMoore + +-- | Replace all occurrences matched by the Boyer-Moore automaton +-- with the given replacement text in some haystack. +-- Performs case-sensitive replacement. +replaceSingleLimited + :: Automaton -- ^ Matches the needles + -> Text -- ^ Replacement string + -> Text -- ^ Haystack + -> CodeUnitIndex -- ^ Maximum number of code units in the returned text + -> Maybe Text +replaceSingleLimited needle replacement haystack maxLength + | needleLength == 0 = Just $ if haystackLength == 0 then replacement else haystack + | otherwise = finish $ BoyerMoore.runText initial foundMatch needle haystack + where + needleLength = BoyerMoore.patternLength needle + haystackLength = Utf8.lengthUtf8 haystack + replacementLength = Utf8.lengthUtf8 replacement + + initial = ReplaceState + { rsChunks = [] + , rsPreviousMatchEnd = 0 + , rsLength = 0 + } + + foundMatch rs matchStart = + let + matchEnd = matchStart + needleLength + + -- Slice the part of the haystack between the end of the previous match + -- and the start of the current match + haystackPartLength = matchStart - rsPreviousMatchEnd rs + haystackPart = Utf8.unsafeSliceUtf8 (rsPreviousMatchEnd rs) haystackPartLength haystack + + -- Add the preceding part of the haystack and the replacement in reverse + -- order to the chunk list (all chunks will be reversed at once in the final step). + newChunks = replacement : haystackPart : rsChunks rs + newLength = replacementLength + haystackPartLength + rsLength rs + + newState = ReplaceState + { rsChunks = newChunks + , rsPreviousMatchEnd = matchEnd + , rsLength = newLength + } + in + if newLength > maxLength + then BoyerMoore.Done newState + else BoyerMoore.Step newState + + finish rs = + let + -- Slice the remaining part of the haystack from the end of the last match + -- to the end of the haystack. + haystackPartLength = haystackLength - rsPreviousMatchEnd rs + finalChunks + = Utf8.unsafeSliceUtf8 (rsPreviousMatchEnd rs) haystackPartLength haystack + : rsChunks rs + finalLength = rsLength rs + haystackPartLength + in + if finalLength > maxLength + then Nothing + else Just $ Text.concat $ reverse finalChunks + +-- | Internal accumulator state for performing a replace while stepping an automaton +data ReplaceState = ReplaceState + { rsChunks :: [Text] + -- ^ Chunks of the final text, in reverse order so that we can efficiently prepend + , rsPreviousMatchEnd :: !CodeUnitIndex + -- ^ Index one past the end of the last match. + , rsLength :: !CodeUnitIndex + -- ^ Length of the newly build string so far, measured in CodeUnits + } diff --git a/src/Data/Text/Utf8/BoyerMoore/Searcher.hs b/src/Data/Text/Utf8/BoyerMoore/Searcher.hs new file mode 100644 --- /dev/null +++ b/src/Data/Text/Utf8/BoyerMoore/Searcher.hs @@ -0,0 +1,121 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2019 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} + +module Data.Text.Utf8.BoyerMoore.Searcher + ( Searcher + , automata + , build + , buildNeedleIdSearcher + , buildWithValues + , containsAll + , containsAny + , needles + , numNeedles + ) where + + +import Control.DeepSeq (NFData) +import Data.Bifunctor (first) +import Data.Hashable (Hashable (hashWithSalt), Hashed, hashed, unhashed) +import GHC.Generics (Generic) + +import Data.Text.Utf8 (Text) +import Data.Text.Utf8.BoyerMoore.Automaton (Automaton) + +import qualified Data.Text.Utf8.BoyerMoore.Automaton as BoyerMoore + + +-- | A set of needles with associated values, and Boyer-Moore automata to +-- efficiently find those needles. +-- +-- INVARIANT: searcherAutomaton = BoyerMoore.buildAutomaton . searcherNeedles +-- To enforce this invariant, the fields are not exposed from this module. +-- There is a separate constructor function. +-- +-- The purpose of this wrapper is to have a type that is Hashable and Eq, so we +-- can derive those for the types that embed the searcher, whithout +-- requiring the automaton itself to be Hashable or Eq, which would be both +-- wasteful and tedious. Because the automaton is fully determined by the +-- needles and associated values, it is sufficient to implement Eq and Hashable +-- in terms of the needles only. +-- +-- We also use Hashed to cache the hash of the needles. +data Searcher v = Searcher + { searcherNeedles :: Hashed [(Text, v)] + , searcherNumNeedles :: Int + , searcherAutomata :: [(Automaton, v)] + } deriving (Generic) + +instance Show (Searcher v) where + show _ = "Searcher _ _ _" + +instance Hashable v => Hashable (Searcher v) where + hashWithSalt salt searcher = hashWithSalt salt $ searcherNeedles searcher + {-# INLINE hashWithSalt #-} + +instance Eq v => Eq (Searcher v) where + Searcher xs nx _ == Searcher ys ny _ = (xs, nx) == (ys, ny) + {-# INLINE (==) #-} + +instance NFData v => NFData (Searcher v) + +-- | Builds the Searcher for a list of needles without values. +-- This is useful for just checking whether the haystack contains the needles. +build :: [Text] -> Searcher () +{-# INLINABLE build #-} +build = buildWithValues . flip zip (repeat ()) + +-- | Builds the Searcher for a list of needles. +buildWithValues :: Hashable v => [(Text, v)] -> Searcher v +{-# INLINABLE buildWithValues #-} +buildWithValues ns = + Searcher (hashed ns) (length ns) $ map (first BoyerMoore.buildAutomaton) ns + +needles :: Searcher v -> [(Text, v)] +needles = unhashed . searcherNeedles + +automata :: Searcher v -> [(Automaton, v)] +automata = searcherAutomata + +numNeedles :: Searcher v -> Int +numNeedles = searcherNumNeedles + +-- | Return whether the haystack contains any of the needles. +-- This function is marked noinline as an inlining boundary. BoyerMoore.runText is +-- marked inline, so this function will be optimized to report only whether +-- there is a match, and not construct a list of matches. We don't want this +-- function be inline, to make sure that the conditions of the caller don't +-- affect how this function is optimized. There is little to gain from +-- additional inlining. The pragma is not an optimization in itself, rather it +-- is a defence against fragile optimizer decisions. +{-# NOINLINE containsAny #-} +containsAny :: Searcher () -> Text -> Bool +containsAny !searcher !text = + let + -- On the first match, return True immediately. + f _acc _match = BoyerMoore.Done True + in + any (\(automaton, ()) -> BoyerMoore.runText False f automaton text) (automata searcher) +-- | Build a 'Searcher' that returns the needle's index in the needle list when it matches. + +buildNeedleIdSearcher :: [Text] -> Searcher Int +buildNeedleIdSearcher !ns = + buildWithValues $ zip ns [0..] + +-- | Like 'containsAny', but checks whether all needles match instead. +-- Use 'buildNeedleIdSearcher' to get an appropriate 'Searcher'. +{-# NOINLINE containsAll #-} +containsAll :: Searcher Int -> Text -> Bool +containsAll !searcher !text = + let + -- On the first match, return True immediately. + f _acc _match = BoyerMoore.Done True + in + all (\(automaton, _) -> BoyerMoore.runText False f automaton text) (automata searcher) diff --git a/src/Data/TypedByteArray.hs b/src/Data/TypedByteArray.hs new file mode 100644 --- /dev/null +++ b/src/Data/TypedByteArray.hs @@ -0,0 +1,83 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2022 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Data.TypedByteArray + ( Data.TypedByteArray.replicate + , MutableTypedByteArray + , Prim + , TypedByteArray + , fromList + , generate + , newTypedByteArray + , unsafeFreezeTypedByteArray + , unsafeIndex + , writeTypedByteArray + ) where + +import Control.DeepSeq (NFData (rnf)) +import Control.Monad.Primitive (PrimMonad (PrimState)) +import Control.Monad.ST (runST) +import Data.Primitive (ByteArray (ByteArray), MutableByteArray, Prim, byteArrayFromList, + indexByteArray, newByteArray, sizeOf, unsafeFreezeByteArray, writeByteArray) + +-- | Thin wrapper around 'ByteArray' that makes signatures and indexing nicer to read. +newtype TypedByteArray a = TypedByteArray ByteArray + deriving Show + +-- | Thin wrapper around 'MutableByteArray s' that makes signatures and indexing nicer to read. +newtype MutableTypedByteArray a s = MutableTypedByteArray (MutableByteArray s) + +instance NFData (TypedByteArray a) where + rnf (TypedByteArray (ByteArray !_)) = () + +{-# INLINE newTypedByteArray #-} +newTypedByteArray :: forall a m. (Prim a, PrimMonad m) => Int -> m (MutableTypedByteArray a (PrimState m)) +newTypedByteArray = fmap MutableTypedByteArray . newByteArray . (* sizeOf (undefined :: a)) + +{-# INLINE fromList #-} +fromList :: Prim a => [a] -> TypedByteArray a +fromList = TypedByteArray . byteArrayFromList + +-- | Element index without bounds checking. +{-# INLINE unsafeIndex #-} +unsafeIndex :: Prim a => TypedByteArray a -> Int -> a +unsafeIndex (TypedByteArray arr) = indexByteArray arr + +{-# INLINE generate #-} +-- | Construct a 'TypedByteArray' of the given length by applying the function to each index in @[0..n-1]@. +generate :: Prim a => Int -> (Int -> a) -> TypedByteArray a +generate !n f = runST $ do + -- Allocate enough space for n elements of type a + arr <- newTypedByteArray n + intLoop 0 n $ \i -> i `seq` writeTypedByteArray arr i $ f i + + unsafeFreezeTypedByteArray arr + +replicate :: (Prim a, PrimMonad m) => Int -> a -> m (MutableTypedByteArray a (PrimState m)) +replicate n value = do + arr <- newTypedByteArray n + intLoop 0 n $ \i -> i `seq` writeTypedByteArray arr i value + pure arr + +{-# INLINE writeTypedByteArray #-} +writeTypedByteArray :: (Prim a, PrimMonad m) => MutableTypedByteArray a (PrimState m) -> Int -> a -> m () +writeTypedByteArray (MutableTypedByteArray array) = writeByteArray array + +{-# INLINE unsafeFreezeTypedByteArray #-} +unsafeFreezeTypedByteArray :: PrimMonad m => MutableTypedByteArray a (PrimState m) -> m (TypedByteArray a) +unsafeFreezeTypedByteArray (MutableTypedByteArray array) = TypedByteArray <$> unsafeFreezeByteArray array + +{-# INLINE intLoop #-} +intLoop :: Monad m => Int -> Int -> (Int -> m ()) -> m () +intLoop !iStart !n p = go iStart + where + go !i + | i >= n = pure () + | otherwise = do + p i + go (i + 1) diff --git a/tests/Data/Text/AhoCorasickSpec.hs b/tests/Data/Text/AhoCorasickSpec.hs --- a/tests/Data/Text/AhoCorasickSpec.hs +++ b/tests/Data/Text/AhoCorasickSpec.hs @@ -7,22 +7,24 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} -module Data.Text.AhoCorasickSpec (spec) where +module Data.Text.AhoCorasickSpec + ( spec + ) where import Control.DeepSeq (rnf) import Control.Monad (forM_, unless) import Data.Foldable (foldl') import Data.List.NonEmpty (NonEmpty ((:|))) +import Data.Text (Text) import Data.Word (Word16) -import Test.Hspec (Spec, Expectation, describe, it, parallel, shouldBe) +import GHC.Stack (HasCallStack) +import Prelude hiding (replicate) +import Test.Hspec (Expectation, Spec, describe, it, parallel, shouldBe) import Test.Hspec.Expectations (shouldMatchList, shouldSatisfy) -import Test.Hspec.QuickCheck (modifyMaxSuccess, modifyMaxSize, prop) +import Test.Hspec.QuickCheck (modifyMaxSize, modifyMaxSuccess, prop) import Test.QuickCheck (Arbitrary (arbitrary, shrink), forAll, forAllShrink, (==>)) -import Test.QuickCheck.Instances () import Test.QuickCheck.Gen (Gen) -import GHC.Stack (HasCallStack) -import Data.Text (Text) -import Prelude hiding (replicate) +import Test.QuickCheck.Instances () import qualified Data.Char as Char import qualified Data.Text as Text @@ -36,6 +38,7 @@ import qualified Data.Text.AhoCorasick.Automaton as Aho import qualified Data.Text.AhoCorasick.Replacer as Replacer +import qualified Data.Text.AhoCorasick.Searcher as Searcher import qualified Data.Text.AhoCorasick.Splitter as Splitter import qualified Data.Text.Utf16 as Utf16 @@ -396,6 +399,32 @@ in Replacer.run replacer haystack `shouldBe` expected + describe "Searcher.containsAll" $ do + + prop "never reports true for empty needles" $ \ (haystack :: Text) -> + let + searcher = Searcher.buildNeedleIdSearcher CaseSensitive [""] + in + Searcher.containsAll searcher haystack `shouldBe` False + + prop "is equivalent to sequential Text.isInfixOf calls" $ \ (needles' :: [NonEmptyText]) (haystack :: Text) -> + let + needles = map unNonEmptyText needles' + searcher = Searcher.buildNeedleIdSearcher CaseSensitive needles + in + Searcher.containsAll searcher haystack `shouldBe` all (`Text.isInfixOf` haystack) needles + + prop "is equivalent to sequential Text.isInfixOf calls for case-insensitive matching" $ \ (needles' :: [NonEmptyText]) (haystack :: Text) -> + let + needles = map unNonEmptyText needles' + + lowerNeedles = map Text.toLower needles + lowerHaystack = Text.toLower haystack + + searcher = Searcher.buildNeedleIdSearcher IgnoreCase lowerNeedles + in + Searcher.containsAll searcher haystack `shouldBe` all (`Text.isInfixOf` lowerHaystack) lowerNeedles + describe "Splitter.split" $ it "passes an example" $ @@ -403,3 +432,15 @@ let splitter = Splitter.build separator in Splitter.split splitter "C++bobobCOBOLbobScala" `shouldBe` "C++" :| ["obCOBOL", "Scala"] + +-- helpers + +-- | A newtype for generating non-empty 'Text' values. +newtype NonEmptyText = NonEmptyText { unNonEmptyText :: Text } + +-- | Simply generates and packs non-empty @[Char]@ values. +instance Arbitrary NonEmptyText where + arbitrary = NonEmptyText . Text.pack <$> Gen.listOf1 arbitrary + +instance Show NonEmptyText where + show = show . unNonEmptyText diff --git a/tests/Data/Text/BoyerMooreSpec.hs b/tests/Data/Text/BoyerMooreSpec.hs --- a/tests/Data/Text/BoyerMooreSpec.hs +++ b/tests/Data/Text/BoyerMooreSpec.hs @@ -3,20 +3,22 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} -module Data.Text.BoyerMooreSpec (spec) where +module Data.Text.BoyerMooreSpec + ( spec + ) where import Control.DeepSeq (rnf) import Control.Monad (forM_) import Data.Foldable (for_) -import Test.Hspec (Spec, Expectation, describe, it, parallel, shouldBe) +import Data.Text (Text) +import GHC.Stack (HasCallStack) +import Prelude hiding (replicate) +import Test.Hspec (Expectation, Spec, describe, it, parallel, shouldBe) import Test.Hspec.Expectations (shouldMatchList, shouldSatisfy) import Test.Hspec.QuickCheck (modifyMaxSuccess, prop) import Test.QuickCheck (Arbitrary (arbitrary, shrink), forAll, forAllShrink, (==>)) -import Test.QuickCheck.Instances () import Test.QuickCheck.Gen (Gen) -import GHC.Stack (HasCallStack) -import Data.Text (Text) -import Prelude hiding (replicate) +import Test.QuickCheck.Instances () import qualified Data.Text as Text import qualified Data.Text.Internal.Search as TextSearch @@ -24,12 +26,13 @@ import qualified Test.QuickCheck as QuickCheck import qualified Test.QuickCheck.Gen as Gen -import Data.Text.Orphans () import Data.Text.BoyerMoore.Automaton (CaseSensitivity (..)) +import Data.Text.Orphans () import qualified Data.Text.AhoCorasick.Replacer as AhoReplacer import qualified Data.Text.BoyerMoore.Automaton as BoyerMoore import qualified Data.Text.BoyerMoore.Replacer as Replacer +import qualified Data.Text.BoyerMoore.Searcher as Searcher import qualified Data.Text.Utf16 as Utf16 -- | Test that for a single needle which equals the haystack, we find a single @@ -207,3 +210,42 @@ in actual `shouldBe` Just expected + describe "Searcher" $ do + + describe "containsAny" $ do + + -- For the edge case where a needle is the empty string, + -- 'Text.isInfixOf' and 'Searcher.containsAny' are different: + -- + -- @ + -- Text.isInfixOf "" "abc" == True /= False == Searcher.containsAny (Searcher.build [""]) "abc" + -- @ + -- + -- However, at this point we probably shouldn't break this property. + prop "is equivalent to disjunction of Text.isInfixOf calls*" $ \ (needles :: [Text]) (haystack :: Text) -> + let + searcher = Searcher.build CaseSensitive needles + test needle = + not (Text.null needle) && needle `Text.isInfixOf` haystack + in + Searcher.containsAny searcher haystack `shouldBe` any test needles + + describe "containsAll" $ do + + prop "is equivalent to conjunction of Text.isInfixOf calls*" $ \ (needles :: [Text]) (haystack :: Text) -> + let + searcher = Searcher.buildNeedleIdSearcher CaseSensitive needles + test needle = + not (Text.null needle) && needle `Text.isInfixOf` haystack + in + Searcher.containsAll searcher haystack `shouldBe` all test needles + + prop "performs case-insensitive search as well" $ \ (needles :: [Text]) (haystack :: Text) -> + let + lowerNeedles = map Utf16.lowerUtf16 needles + lowerHaystack = Utf16.lowerUtf16 haystack + searcher = Searcher.buildNeedleIdSearcher IgnoreCase lowerNeedles + test needle = + not (Text.null needle) && needle `Text.isInfixOf` lowerHaystack + in + Searcher.containsAll searcher haystack `shouldBe` all test lowerNeedles diff --git a/tests/Data/Text/Orphans.hs b/tests/Data/Text/Orphans.hs --- a/tests/Data/Text/Orphans.hs +++ b/tests/Data/Text/Orphans.hs @@ -4,7 +4,16 @@ import qualified Test.QuickCheck.Gen as Gen -import Data.Text.AhoCorasick.Automaton (CaseSensitivity (..)) +import Data.Text.Utf8 as Utf8 +import Data.Text.CaseSensitivity (CaseSensitivity (..)) + instance Arbitrary CaseSensitivity where arbitrary = Gen.elements [CaseSensitive, IgnoreCase] + +-- TODO: Slow placeholder implementation until we can use text-2.0 +instance Arbitrary Utf8.Text where + arbitrary = fmap Utf8.pack arbitrary + +instance Arbitrary Utf8.CodeUnitIndex where + arbitrary = fmap Utf8.CodeUnitIndex arbitrary diff --git a/tests/Data/Text/Utf8/AhoCorasickSpec.hs b/tests/Data/Text/Utf8/AhoCorasickSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Text/Utf8/AhoCorasickSpec.hs @@ -0,0 +1,248 @@ +-- Alfred-Margaret: Fast Aho-Corasick string searching +-- Copyright 2022 Channable +-- +-- Licensed under the 3-clause BSD license, see the LICENSE file in the +-- repository root. + +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Data.Text.Utf8.AhoCorasickSpec where + +import Control.Monad (forM_) +import Data.Foldable (foldl') +import Data.List.NonEmpty (NonEmpty ((:|))) +import Data.Primitive (byteArrayFromList) +import Test.Hspec (Expectation, Spec, describe, it, shouldBe) +import Test.Hspec.QuickCheck (modifyMaxSize, prop) +import Test.QuickCheck (Arbitrary (arbitrary, shrink), forAll, forAllShrink) + +import qualified Data.Text as T +import qualified Test.QuickCheck.Gen as Gen + +import Data.Text.CaseSensitivity (CaseSensitivity (..)) +import Data.Text.Orphans () +import Data.Text.Utf8 (Text) + +import qualified Data.Text.Utf8 as Text +import qualified Data.Text.Utf8 as Utf8 +import qualified Data.Text.Utf8.AhoCorasick.Automaton as Aho +import qualified Data.Text.Utf8.AhoCorasick.Replacer as Replacer +import qualified Data.Text.Utf8.AhoCorasick.Searcher as Searcher +import qualified Data.Text.Utf8.AhoCorasick.Splitter as Splitter + +spec :: Spec +spec = do + -- Ensure that helper functions are actually helping + -- Examples are from https://en.wikipedia.org/wiki/UTF-8 + describe "IsString ByteArray" $ do + + it "encodes the dollar sign" $ utf8Test "$" [0x24] + it "encodes the euro sign" $ utf8Test "€" [0xe2, 0x82, 0xac] + it "encodes the pound sign" $ utf8Test "Β£" [0xc2, 0xa3] + it "encodes Hwair" $ utf8Test "𐍈" [0xf0, 0x90, 0x8d, 0x88] + it "encodes all of the above" $ utf8Test "$β‚¬Β£πˆ" [0x24, 0xe2, 0x82, 0xac, 0xc2, 0xa3, 0xf0, 0x90, 0x8d, 0x88] + + describe "runText" $ do + + describe "countMatches" $ do + it "counts the right number of matches in a basic example" $ do + countMatches Aho.CaseSensitive ["abc", "rst", "xyz"] "abcdefghijklmnopqrstuvwxyz" `shouldBe` 3 + + it "counts the right number of matches in an example with 1-, 2-, 3- and 4-code unit code points" $ do + countMatches Aho.CaseSensitive ["$", "Β£"] "$β‚¬Β£πˆ" `shouldBe` 2 + + describe "runLower" $ do + + describe "countMatches" $ do + it "counts the right number of matches in a basic example" $ do + countMatches Aho.IgnoreCase ["abc", "rst", "xyz"] "abcdefghijklmnopqrstuvwxyz" `shouldBe` 3 + + it "does not work with uppercase needles" $ do + countMatches Aho.IgnoreCase ["ABC", "Rst", "xYZ"] "abcdefghijklmnopqrstuvwxyz" `shouldBe` 0 + + it "works with characters that are not in ASCII" $ do + countMatches Aho.IgnoreCase ["groß", "ΓΆffnung", "tΓΌr"] "GroßfrΓ€smaschinenΓΆffnungstΓΌr" `shouldBe` 3 + + modifyMaxSize (const 10) $ describe "Replacer" $ do + + describe "run" $ do + let + genHaystack = fmap Utf8.pack $ Gen.listOf $ Gen.frequency [(40, Gen.elements "abAB"), (1, pure 'Δ°'), (1, arbitrary)] + + -- needles may not be empty, because empty needles are filtered out in an I.ActionReplaceMultiple + genNeedle = fmap Utf8.pack $ Gen.resize 3 $ Gen.listOf1 $ Gen.elements "abAB" + genReplaces = Gen.listOf $ (,) <$> genNeedle <*> arbitrary + shrinkReplaces = filter (not . any (\(needle, _) -> Utf8.null needle)) . shrink + + replace needles haystack = + Replacer.run (Replacer.build Aho.CaseSensitive needles) haystack + + replaceIgnoreCase needles haystack = + Replacer.run (Replacer.build Aho.IgnoreCase needles) haystack + + it "replaces all occurrences" $ do + replace [("A", "B")] "AXAXB" `shouldBe` "BXBXB" + replace [("A", "B"), ("X", "Y")] "AXAXB" `shouldBe` "BYBYB" + replace [("aaa", ""), ("b", "c")] "aaabaaa" `shouldBe` "c" + -- Have a few non-matching needles too. + replace [("A", "B"), ("Q", "r"), ("Z", "")] "AXAXB" `shouldBe` "BXBXB" + + it "replaces only non-overlapping matches" $ do + replace [("aa", "zz"), ("bb", "w")] "aaabbb" `shouldBe` "zzawb" + replace [("aaa", "")] "aaaaa" `shouldBe` "aa" + + it "replaces all occurrences in priority order" $ do + replace [("A", ""), ("BBBB", "bingo")] "BBABB" `shouldBe` "bingo" + replace [("BB", ""), ("BBBB", "bingo")] "BBBB" `shouldBe` "" + + it "replaces needles that contain a surrogate pair" $ + replace [("\x1f574", "levitating man in business suit")] + "the \x1f574" `shouldBe` "the levitating man in business suit" + + + it "replaces all occurrences case-insensitively" $ do + replaceIgnoreCase [("A", "B")] "AXAXB" `shouldBe` "BXBXB" + replaceIgnoreCase [("A", "B")] "axaxb" `shouldBe` "BxBxb" + replaceIgnoreCase [("a", "b")] "AXAXB" `shouldBe` "bXbXB" + + replaceIgnoreCase [("A", "B"), ("X", "Y")] "AXAXB" `shouldBe` "BYBYB" + replaceIgnoreCase [("A", "B"), ("X", "Y")] "axaxb" `shouldBe` "BYBYb" + replaceIgnoreCase [("a", "b"), ("x", "y")] "AXAXB" `shouldBe` "bybyB" + + it "matches replacements case-insensitively" $ + replaceIgnoreCase [("foo", "BAR"), ("bar", "BAZ")] "Foo" `shouldBe` "BAZ" + + it "matches replacements case-insensitively for non-ascii characters" $ do + replaceIgnoreCase [("Γ©clair", "lightning")] "Γ‰clair" `shouldBe` "lightning" + -- Note: U+0319 is an uppercase alpha, which looks exactly like A, but it + -- is a different code point. + replaceIgnoreCase [("bΓ¨ta", "Ξ±"), ("\x0391", "alpha")] "BÈTA" `shouldBe` "alpha" + + it "matches surrogate pairs case-insensitively" $ do + -- We can't lowercase a levivating man in business suit, but that should + -- not affect whether we match it or not. + replaceIgnoreCase [("\x1f574", "levitating man in business suit")] "the \x1f574" + `shouldBe` "the levitating man in business suit" + + prop "satisfies (run . compose a b) == (run b (run a))" $ + forAllShrink genHaystack shrink $ \haystack -> + forAll arbitrary $ \case_ -> + forAllShrink genReplaces shrinkReplaces $ \replaces1 -> + forAllShrink genReplaces shrinkReplaces $ \replaces2 -> + let + rm1 = Replacer.build case_ replaces1 + rm2 = Replacer.build case_ replaces2 + Just rm12 = Replacer.compose rm1 rm2 + in + Replacer.run rm2 (Replacer.run rm1 haystack) + `shouldBe` Replacer.run rm12 haystack + + prop "is identity for empty needles" $ \case_ haystack -> + let replacerId = Replacer.build case_ [] + in Replacer.run replacerId haystack `shouldBe` haystack + + prop "is equivalent to sequential Text.replace calls" $ + forAllShrink genHaystack shrink $ \haystack -> + forAllShrink genReplaces shrinkReplaces $ \replaces -> + let + replacer = Replacer.build Aho.CaseSensitive replaces + -- TODO: Remove conversions once we move to text-2.0 + replaceText agg (needle, replacement) = Utf8.pack $ T.unpack $ T.replace (T.pack $ Utf8.unpack needle) (T.pack $ Utf8.unpack replacement) (T.pack $ Utf8.unpack agg) + expected = foldl' replaceText haystack replaces + in + Replacer.run replacer haystack `shouldBe` expected + + describe "Searcher" $ do + + describe "containsAny" $ do + + it "gives the right values for the examples in the README" $ do + let needles = ["tshirt", "shirts", "shorts"] + let searcher = Searcher.build Aho.CaseSensitive needles + + Searcher.containsAny searcher "short tshirts" `shouldBe` True + Searcher.containsAny searcher "long shirt" `shouldBe` False + Searcher.containsAny searcher "Short TSHIRTS" `shouldBe` False + + let searcher' = Searcher.build Aho.IgnoreCase needles + + Searcher.containsAny searcher' "Short TSHIRTS" `shouldBe` True + + it "works with the the first line of the illiad" $ do + let illiad = "Ἄνδρα ΞΌΞΏΞΉ ἔννΡπΡ, ΞœΞΏαΏ¦ΟƒΞ±, πολύτροπον, ὃς μάλα πολλὰ" + needleSets = [(["ΞΌΞΏΞΉ"], True), (["α½ˆΞ΄Ο…ΟƒΞ΅ΟΟ‚"], False)] + + forM_ needleSets $ \(needles, expectedResult) -> do + let searcher = Searcher.build Aho.CaseSensitive needles + Searcher.containsAny searcher illiad `shouldBe` expectedResult + + describe "containsAll" $ do + + prop "never reports true for empty needles" $ \ (haystack :: Text) -> + let + searcher = Searcher.buildNeedleIdSearcher CaseSensitive [""] + in + Searcher.containsAll searcher haystack `shouldBe` False + + prop "is equivalent to sequential Text.isInfixOf calls for non-empty needles" $ \ (needles' :: [NonEmptyText]) (haystack :: Text) -> + let + needles = map unNonEmptyText needles' + searcher = Searcher.buildNeedleIdSearcher CaseSensitive needles + in + Searcher.containsAll searcher haystack `shouldBe` all (`Text.isInfixOf` haystack) needles + + prop "is equivalent to sequential Text.isInfixOf calls for case-insensitive matching for non-empty needles" $ \ (needles' :: [NonEmptyText]) (haystack :: Text) -> + let + needles = map unNonEmptyText needles' + + lowerNeedles = map Utf8.lowerUtf8 needles + lowerHaystack = Utf8.lowerUtf8 haystack + + searcher = Searcher.buildNeedleIdSearcher IgnoreCase lowerNeedles + in + Searcher.containsAll searcher haystack `shouldBe` all (`Text.isInfixOf` lowerHaystack) lowerNeedles + + describe "Splitter" $ do + + describe "split" $ do + + it "passes an example" $ do + let separator = "bob" + splitter = Splitter.build separator + + Splitter.split splitter "C++bobobCOBOLbobScala" `shouldBe` "C++" :| ["obCOBOL", "Scala"] + + it "neatly splits the first line of the illiad" $ do + let splitter = Splitter.build ", " + + Splitter.split splitter "Ἄνδρα ΞΌΞΏΞΉ ἔννΡπΡ, ΞœΞΏαΏ¦ΟƒΞ±, πολύτροπον, ὃς μάλα πολλὰ" `shouldBe` + "Ἄνδρα ΞΌΞΏΞΉ ἔννΡπΡ" :| ["ΞœΞΏαΏ¦ΟƒΞ±", "πολύτροπον", "ὃς μάλα πολλὰ"] + +-- helpers + +utf8Test :: Utf8.Text -> [Utf8.CodeUnit] -> Expectation +utf8Test str byteList = str `shouldBe` Utf8.Text (byteArrayFromList byteList) 0 (length byteList) + +-- From ./benchmark +countMatches :: Aho.CaseSensitivity -> [Utf8.Text] -> Utf8.Text -> Int +{-# NOINLINE countMatches #-} +countMatches caseSensitivity needles haystack = case needles of + [] -> 0 + _ -> + let + ac = Aho.build $ zip needles (repeat ()) + onMatch !n _match = Aho.Step (n + 1) + in + Aho.runWithCase caseSensitivity 0 onMatch ac haystack + +-- | A newtype for generating non-empty 'Text' values. +newtype NonEmptyText = NonEmptyText { unNonEmptyText :: Text } + +-- | Simply generates and packs non-empty @[Char]@ values. +instance Arbitrary NonEmptyText where + arbitrary = NonEmptyText . Text.pack <$> Gen.listOf1 arbitrary + +instance Show NonEmptyText where + show = show . unNonEmptyText diff --git a/tests/Data/Text/Utf8/BoyerMooreSpec.hs b/tests/Data/Text/Utf8/BoyerMooreSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Text/Utf8/BoyerMooreSpec.hs @@ -0,0 +1,262 @@ +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} + +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Data.Text.Utf8.BoyerMooreSpec + ( spec + ) where + +import Control.DeepSeq (rnf) +import Control.Monad (forM_) +import Data.Foldable (for_) +import GHC.Stack (HasCallStack) +import Prelude hiding (replicate) +import Test.Hspec (Expectation, Spec, describe, it, parallel, shouldBe) +import Test.Hspec.Expectations (shouldMatchList, shouldSatisfy) +import Test.Hspec.QuickCheck (modifyMaxSuccess, prop) +import Test.QuickCheck (Arbitrary (arbitrary, shrink), forAllShrink, (==>)) +import Test.QuickCheck.Gen (Gen) +import Test.QuickCheck.Instances () + +-- import qualified Data.Text.Internal.Search as TextSearch +import qualified Test.QuickCheck as QuickCheck +import qualified Test.QuickCheck.Gen as Gen + +import Data.Text.Orphans () +import Data.Text.Utf8 (Text) +import Data.Text.Utf8.BoyerMoore.Automaton (CaseSensitivity (..)) + +import qualified Data.Text.Utf8 as Text +import qualified Data.Text.Utf8 as TextSearch +import qualified Data.Text.Utf8 as Utf8 +import qualified Data.Text.Utf8.AhoCorasick.Replacer as AhoReplacer +import qualified Data.Text.Utf8.BoyerMoore.Automaton as BoyerMoore +import qualified Data.Text.Utf8.BoyerMoore.Replacer as Replacer +import qualified Data.Text.Utf8.BoyerMoore.Searcher as Searcher + +-- | Test that for a single needle which equals the haystack, we find a single +-- match. Does not apply to the empty needle. +needleIsHaystackMatches :: HasCallStack => Text -> Expectation +needleIsHaystackMatches needle = + let + prependMatch ms match = BoyerMoore.Step (Utf8.codeUnitIndex match : ms) + matches = BoyerMoore.runText [] prependMatch (BoyerMoore.buildAutomaton needle) needle + in + matches `shouldBe` [0] + +boyerMatch :: Text -> Text -> [Int] +boyerMatch needle haystack = + let + prependMatch matches match = BoyerMoore.Step (Utf8.codeUnitIndex match : matches) + in + BoyerMoore.runText [] prependMatch (BoyerMoore.buildAutomaton needle) haystack + +-- | Match without a payload, return only the match positions. +matchEndPositions :: Text -> Text -> [Int] +matchEndPositions needle haystack = + let + matches = boyerMatch needle haystack + in + fmap (Utf8.codeUnitIndex (Utf8.lengthUtf8 needle) +) matches + +-- | `matchEndPositions` implemented naively in terms of Text's functionality, +-- which we assume to be correct. +naiveMatchPositions :: Text -> Text -> [Int] +naiveMatchPositions needle haystack = + map toEndPos $ TextSearch.indices needle haystack + where + toEndPos index = Utf8.codeUnitIndex (Utf8.lengthUtf8 needle) + index + +-- | Generate random needles and haystacks, such that the needles have a +-- reasonable probability of occuring in the haystack, which would hardly be the +-- case if we just generated random texts for all of them. We do this by first +-- generating a set of fragments, and then building the haystack and needles by +-- combining these fragments. By doing this, we also get a lot of partial +-- matches, where part of a needle does occur in the haystack, but the full +-- needle does not, and also needles with a shared prefix or suffix. This should +-- fully stress the possible transitions in the automaton. +arbitraryNeedleHaystack :: Gen (Text, Text) +arbitraryNeedleHaystack = do + let + -- Prefer ascii just to have printable test cases, but do include the other + -- generator to cover the entire range of code points. + genChar = Gen.frequency + [ (4, QuickCheck.arbitraryASCIIChar) + , (1, QuickCheck.arbitrary) + ] + genNonEmptyText = do + chars <- Gen.listOf1 genChar + pure $ Text.pack chars + + fragments <- Gen.listOf1 $ Gen.resize 5 genNonEmptyText + let + genFragment = Gen.elements $ filter (not . Text.null) fragments + genSmall = Gen.scale (`div` 3) $ Gen.listOf1 genFragment + genBig = Gen.scale (* 4) $ Gen.listOf1 genFragment + + needle <- fmap Text.concat genSmall + haystack <- fmap Text.concat genBig + pure (needle, haystack) + +spec :: Spec +spec = parallel $ modifyMaxSuccess (const 200) $ do + describe "build" $ do + prop "does not throw exceptions" $ \ (pat :: Text) -> + rnf $ BoyerMoore.buildAutomaton pat + + describe "runText" $ do + + describe "when given a needle equal to the haystack" $ do + + it "reports a single match for a repeated character" $ + forM_ [1..128] $ \n -> + needleIsHaystackMatches $ Text.replicate n "a" + + it "reports a single match for non-BMP data" $ do + -- Include a few code points outside of the Basic Multilingual Plane, + -- which require multible code units to encode. + needleIsHaystackMatches "\x000437b8suffix" + needleIsHaystackMatches "aaa\359339aaa\95759aa\899256aa" + + prop "reports a single match for random needles" $ \needle -> + not (Text.null needle) ==> needleIsHaystackMatches needle + + describe "when given a sliced text (with nonzero internal offset)" $ do + + it "still reports offset relative to the text start" $ + -- The match position should be relative to the start of the text "a". + -- Even if this text is represented as a slice of "bbba" internally. + matchEndPositions "a" (Text.dropWhile (== 'b') "bbba") `shouldMatchList` [1] + + describe "when given non-ascii inputs" $ do + + -- We have a special lookup table for bad character shifts for + -- the first 128 code units, which is always hit for ascii inputs. + -- Also exercise the fallback code path with a different input. + -- The code point Γ© is encoded as two code units in UTF-8. + -- 0 7 13 + -- β”‚ β”‚ β”‚ + -- β–Ό β–Ό β–Ό + -- β”Œβ”€β”€β”€β”¬β”€β”¬β”€β”¬β”€β”¬β”€β”¬β”€β”¬β”€β”¬β”€β”¬β”€β”¬β”€β”¬β”€β”¬β”€β” + -- β”‚ Γ© β”‚cβ”‚lβ”‚aβ”‚iβ”‚rβ”‚eβ”‚cβ”‚lβ”‚aβ”‚iβ”‚rβ”‚ Code Points + -- β”œβ”€β”¬β”€β”Όβ”€β”Όβ”€β”Όβ”€β”Όβ”€β”Όβ”€β”Όβ”€β”Όβ”€β”Όβ”€β”Όβ”€β”Όβ”€β”Όβ”€β”€ + -- β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ Code Units (Bytes) + -- β””β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”˜ + it "reports a match if the haystack contains a character > U+7f" $ do + matchEndPositions "eclair" "Γ©claireclair" `shouldMatchList` [13] + matchEndPositions "Γ©clair" "Γ©claireclair" `shouldMatchList` [7] + matchEndPositions "Γ©clair" "eclairΓ©clair" `shouldMatchList` [13] + + it "reports the correct code unit index for complex characters" $ do + -- Note that the index after the match is 4, even though there is + -- only a single code point. U+1d11e is encoded as four code units: + -- in UTF-8: + -- 0 4 + -- β”‚ β”‚ + -- β–Ό β–Ό + -- β”Œβ”€β”€β”€β”€β”€β”€β”€β” + -- β”‚ π„ž β”‚ Code Points + -- β”œβ”€β”¬β”€β”¬β”€β”¬β”€β”€ + -- β”‚ β”‚ β”‚ β”‚ β”‚ Code Units (Bytes) + -- β””β”€β”΄β”€β”΄β”€β”΄β”€β”˜ + matchEndPositions "π„ž" "π„ž" `shouldMatchList` [4] + + -- A levitating woman in business suit with dark skin tone needs a + -- whopping 5 code points to encode. The first two need 4 code units each to encode, + -- the remaining three need 3 code units each for a total of 17 code units: + -- 0 4 8 17 + -- β”‚ β”‚ β”‚ β”‚ + -- β–Ό β–Ό β–Ό β–Ό + -- β”Œβ”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β” + -- β”‚ 1 β”‚ 2 β”‚ 3 β”‚ 4 β”‚ 5 β”‚ Code Points + -- β”œβ”€β”¬β”€β”¬β”€β”¬β”€β”Όβ”€β”¬β”€β”¬β”€β”¬β”€β”Όβ”€β”¬β”€β”¬β”€β”Όβ”€β”¬β”€β”¬β”€β”Όβ”€β”¬β”€β”¬β”€β”€ + -- β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ Code Units (Bytes) + -- β””β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”΄β”€β”˜ + -- 1. U+1f574: man in business suit levitating (πŸ•΄) + -- 2. U+1f3ff: emoji modifier Fitzpatrick type-6 + -- 3. U+200d: zero width joiner + -- 4. U+2640: female sign (♀) + -- 5. U+fe0f: variation selector-16 + -- A peculiar feature of Unicode emoji, is that the male levivating + -- man in business suit with dark skin tone is a substring of the + -- levivating woman in business suit. And the levivating man in + -- business suit without particular skin tone is a substring of that. + let + examples = + [ ("\x1f574\x1f3ff\x200d\x2640\xfe0f", 17) + , ("\x1f574\x1f3ff", 8) + , ("\x1f574", 4) + ] + for_ examples $ \(needle, endPos) -> + matchEndPositions needle "\x1f574\x1f3ff\x200d\x2640\xfe0f" `shouldMatchList` [endPos] + + describe "when given empty needle" $ do + + it "does not report a match" $ do + matchEndPositions "" "" `shouldMatchList` [] + matchEndPositions "" "foo" `shouldMatchList` [] + + describe "kitchen sink" $ do + it "kitchen sinks" $ do + matchEndPositions "\"\SO]JL\"" "aaaaa\"\SO]JL\"" `shouldMatchList` [11] + matchEndPositions "\"X]JL\"" "aaaaa\"X]JL\"" `shouldMatchList` [11] + + describe "when given random needles and haystacks" $ do + + prop "reports only infixes of the haystack" $ + QuickCheck.forAllShrink arbitraryNeedleHaystack shrink $ \ (needle, haystack) -> + let + matches = boyerMatch needle haystack + sliceMatch startPos len = Utf8.unsafeSliceUtf8 startPos len haystack + in + forM_ matches $ \pos -> do + needle `shouldSatisfy` (`Text.isInfixOf` haystack) + sliceMatch (Utf8.CodeUnitIndex pos) (Utf8.lengthUtf8 needle) `shouldBe` needle + + prop "reports all infixes of the haystack" $ + QuickCheck.forAllShrink arbitraryNeedleHaystack shrink $ \ (needle, haystack) -> + matchEndPositions needle haystack `shouldMatchList` naiveMatchPositions needle haystack + + describe "replaceSingleLimited" $ do + + prop "is equivalent to Aho-Corasick replacer with a single needle" $ + forAllShrink arbitraryNeedleHaystack shrink $ \(needle, haystack) -> + forAllShrink arbitrary shrink $ \replacement -> + let + expected = AhoReplacer.run (AhoReplacer.build CaseSensitive [(needle, replacement)]) haystack + + auto = BoyerMoore.buildAutomaton needle + + actual = Replacer.replaceSingleLimited auto replacement haystack maxBound + in + actual `shouldBe` Just expected + + describe "Searcher" $ do + + describe "containsAny" $ do + + -- For the edge case where a needle is the empty string, + -- 'Text.isInfixOf' and 'Searcher.containsAny' are different: + -- + -- @ + -- Text.isInfixOf "" "abc" == True /= False == Searcher.containsAny (Searcher.build [""]) "abc" + -- @ + -- + -- However, at this point we probably shouldn't break this property. + prop "is equivalent to disjunction of Text.isInfixOf calls*" $ \ (needles :: [Text]) (haystack :: Text) -> + let + searcher = Searcher.build needles + test needle = + not (Text.null needle) && needle `Text.isInfixOf` haystack + in + Searcher.containsAny searcher haystack `shouldBe` any test needles + + describe "containsAll" $ do + prop "is equivalent to conjunction of Text.isInfixOf calls*" $ \ (needles :: [Text]) (haystack :: Text) -> + let + searcher = Searcher.buildNeedleIdSearcher needles + test needle = + not (Text.null needle) && needle `Text.isInfixOf` haystack + in + Searcher.containsAll searcher haystack `shouldBe` all test needles diff --git a/tests/Data/Text/Utf8/Utf8Spec.hs b/tests/Data/Text/Utf8/Utf8Spec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Text/Utf8/Utf8Spec.hs @@ -0,0 +1,106 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Data.Text.Utf8.Utf8Spec where + +import Control.Monad (forM_) +import Test.Hspec (Spec, describe, it, shouldBe, shouldSatisfy) +import Test.Hspec.QuickCheck (prop) +import Test.QuickCheck (Gen, choose, forAllShrink, shrink) + +import qualified Data.Char as Char + +import Data.Text.Orphans () +import Data.Text.Utf8 (stringToByteArray) + +import qualified Data.Text.Utf8 as Utf8 + +spec :: Spec +spec = do + describe "Properties of the BMP in UTF-8" $ do + + describe "Char.toLower" $ do + + {- + it "does not generate common suffixes" $ do + forM_ bmpCodepoints $ flip shouldSatisfy $ \cp -> + let + lowerCp = mapCp Char.toLower cp + in + cp == lowerCp || null (commonSuffix (Utf8.unicode2utf8 cp) (Utf8.unicode2utf8 lowerCp)) + -- Sadly, it "actually does" + -} + + it "is idempotent" $ do + forM_ bmpCodepoints $ flip shouldSatisfy $ \cp -> + Char.toLower cp == Char.toLower (Char.toLower cp) + + describe "toLowerAscii" $ do + + it "is equivalent to Char.toLower on ASCII" $ do + + forM_ asciiCodepoints $ flip shouldSatisfy $ \cp -> + Char.toLower cp == Utf8.toLowerAscii cp + + describe "lowerCodePoint" $ do + + prop "is equivalent to Char.toLower on all of Unicode" $ \c -> + Utf8.lowerCodePoint c `shouldBe` Char.toLower c + + describe "dropWhile" $ do + + it "handles a simple example well" $ do + Utf8.dropWhile (== 'b') "bbba" `shouldBe` "a" + + describe "slicing functions" $ do + + let + -- | Example shown in section "Slicing Functions" in 'Data.Text.Utf8". + slicingExample :: Utf8.Text + slicingExample = Utf8.Text (stringToByteArray "ABCDEFGHIJKLMN") 1 11 + + it "satisfies the example in Data.Text.Utf8" $ do + let begin = Utf8.CodeUnitIndex 2 + let length_ = Utf8.CodeUnitIndex 6 + Utf8.unsafeSliceUtf8 begin length_ slicingExample `shouldBe` "DEFGHI" + Utf8.unsafeCutUtf8 begin length_ slicingExample `shouldBe` ("BC", "JKL") + + prop "unsafeSliceUtf8 and unsafeCutUtf8 are complementary" $ + forAllShrink (arbitrarySlicingIndices slicingExample) shrink $ \ (begin, length_) -> do + let (prefix, suffix) = Utf8.unsafeCutUtf8 begin length_ slicingExample + Utf8.concat [prefix, Utf8.unsafeSliceUtf8 begin length_ slicingExample, suffix] `shouldBe` slicingExample + + describe "Basic Text instances" $ do + + prop "Show Text behaves like Show String" $ \ (str :: String) -> do + show (Utf8.pack str) `shouldBe` show str + + prop "Eq Text behaves like Eq String" $ \ (a :: String) (b :: String) -> do + Utf8.pack a == Utf8.pack b `shouldBe` a == b + + prop "Ord Text behaves like Ord String" $ \ (a :: String) (b :: String) -> do + compare (Utf8.pack a) (Utf8.pack b) `shouldBe` compare a b + +arbitrarySlicingIndices :: Utf8.Text -> Gen (Utf8.CodeUnitIndex, Utf8.CodeUnitIndex) +arbitrarySlicingIndices example = do + let exampleLength = Utf8.codeUnitIndex $ Utf8.lengthUtf8 example + + begin <- choose (0, exampleLength) + length_ <- choose (0, exampleLength - begin) + + pure (Utf8.CodeUnitIndex begin, Utf8.CodeUnitIndex length_) + +asciiCodepoints :: [Char] +asciiCodepoints = map Char.chr [0..0x7f] + +-- | The Basic Multilingual Plane (BMP) contains the Unicode code points +-- 0x0000 through 0xFFFF. +bmpCodepoints :: [Char] +bmpCodepoints = map Char.chr [0..0xffff] + +commonSuffix :: Eq a => [a] -> [a] -> [a] +commonSuffix list list' = reverse $ go (reverse list) (reverse list') + where + go (x:xs) (y:ys) + | x == y = x : go xs ys + go _ _ = [] diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -4,8 +4,14 @@ import Data.Text.AhoCorasickSpec as A import Data.Text.BoyerMooreSpec as B +import Data.Text.Utf8.AhoCorasickSpec as U8A +import Data.Text.Utf8.BoyerMooreSpec as U8B +import Data.Text.Utf8.Utf8Spec as U8 main :: IO () main = hspec $ do describe "Data.Text.AhoCorasick" A.spec describe "Data.Text.BoyerMoore" B.spec + describe "Data.Text.Utf8.AhoCorasick" U8A.spec + describe "Data.Text.Utf8.BoyerMoore" U8B.spec + describe "Data.Text.Utf8" U8.spec