alfred-margaret (empty) → 1.0.0.0

raw patch · 8 files changed

+1560/−0 lines, 8 filesdep +QuickCheckdep +alfred-margaretdep +basesetup-changed

Dependencies added: QuickCheck, alfred-margaret, base, containers, deepseq, hashable, hspec, hspec-expectations, primitive, quickcheck-instances, text, vector

Files

• LICENSE +30/−0
• README.md +107/−0
• Setup.hs +2/−0
• alfred-margaret.cabal +59/−0
• src/Data/Text/AhoCorasick/Automaton.hs +659/−0
• src/Data/Text/AhoCorasick/Replacer.hs +208/−0
• src/Data/Text/AhoCorasick/Searcher.hs +126/−0
• tests/AhoCorasickSpec.hs +369/−0

@@ -0,0 +1,30 @@+Copyright 2019 Channable++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++    * Redistributions of source code must retain the above copyright+      notice, this list of conditions and the following disclaimer.++    * Redistributions in binary form must reproduce the above+      copyright notice, this list of conditions and the following+      disclaimer in the documentation and/or other materials provided+      with the distribution.++    * Neither the name of the copyright holder nor the names of other+      contributors may be used to endorse or promote products derived+      from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

@@ -0,0 +1,107 @@+# Alfred–Margaret++Alfred–Margaret is a fast implementation of the Aho–Corasick string+searching algorithm in Haskell. It powers many string-related operations+in [Channable][channable].++The library is designed to work with the [`text`][text] package. It matches+directly on the internal UTF-16 representation of `Text` for efficiency. See the+[announcement blog post][blog-post] for a deeper dive into Aho–Corasick, and the+optimizations that make this library fast.++Alfred–Margaret is named after Alfred Aho and Margaret Corasick.++## Performance++Running time to count all matches, in a real-world data set,+comparing [a Java implementation][hankcs] and [a Rust implementation][burntsushi]+against Alfred–Margaret, and against memcopy to establish a lower bound:++<p align="center">+<img+  title="Graph that shows that Alfred–Margaret is fast."+  src="performance.png"+  width="80%"+>+</p>++For the full details of this benchmark, see+[our announcement blog post][blog-post], which includes more details about the+data set, the benchmark setup, and a few things to keep in mind when+interpreting this graph.++## Example++Check if a string contains one of the needles:++```haskell+import qualified Data.Text.AhoCorasick.Searcher as Searcher++searcher = Searcher.build ["tshirt", "shirts", "shorts"]++Searcher.containsAny searcher "short tshirts"+> True++Searcher.containsAny searcher "long shirt"+> False++Searcher.containsAny searcher "Short TSHIRTS"+> False++Searcher.containsAnyIgnoreCase searcher "Short TSHIRTS"+> True+```++Sequentially replace many needles:++```haskell+import Data.Text.AhoCorasick.Automaton (CaseSensitivity (..))+import qualified Data.Text.AhoCorasick.Replacer as Replacer++replacer = Replacer.build CaseSensitive [("tshirt", "banana"), ("shirt", "pear")]++Replacer.run replacer "tshirts for sale"+> "bananas for sale"++Replacer.run replacer "tshirts and shirts for sale"+> "bananas and pears for sale"++Replacer.run replacer "sweatshirts and shirtshirts"+> "sweabananas and shirbananas"++Replacer.run replacer "sweatshirts and shirttshirts"+> "sweabananas and pearbananas"+```++Get all matches, possibly overlapping:++```haskell+import qualified Data.Text.AhoCorasick.Automaton as Aho++pairNeedleWithSelf text = (Aho.unpackUtf16 text, text)+automaton = Aho.build $ fmap pairNeedleWithSelf ["tshirt", "shirts", "shorts"]+allMatches = Aho.runText [] (\matches match -> Aho.Step (match : matches))++allMatches automaton "short tshirts"+> [ Match {matchPos = CodeUnitIndex 13, matchValue = "shirts"}+> , Match {matchPos = CodeUnitIndex 12, matchValue = "tshirt"}+> ]++allMatches automaton "sweatshirts and shirtshirts"+> [ Match {matchPos = CodeUnitIndex 27, matchValue = "shirts"}+> , Match {matchPos = CodeUnitIndex 26, matchValue = "tshirt"}+> , Match {matchPos = CodeUnitIndex 22, matchValue = "shirts"}+> , Match {matchPos = CodeUnitIndex 11, matchValue = "shirts"}+> , Match {matchPos = CodeUnitIndex 10, matchValue = "tshirt"}+> ]+```++## License++Alfred–Margaret is licensed under the 3-clause BSD license.++[channable]:  https://www.channable.com/+[blog-post]:  https://tech.channable.com/posts/2019-03-13-how-we-made-haskell-search-strings-as-fast-as-rust.html+[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

@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain

@@ -0,0 +1,59 @@+name:                alfred-margaret+version:             1.0.0.0+synopsis:            Fast Aho-Corasick string searching+description:         An efficient implementation of the Aho-Corasick+                     string searching algorithm.+homepage:            https://github.com/channable/alfred-margaret+license:             BSD3+license-file:        LICENSE+author:              The Alfred-Margaret authors+maintainer:          Ruud van Asseldonk <ruud@channable.com>+copyright:           2019 Channable+category:            Data, Text+build-type:          Simple+extra-source-files:  README.md+cabal-version:       >=1.10+tested-with:+                     -- Stackage LTS 11.22.+                     GHC == 8.2.2+                     -- Stackage LTS 12.14.+                   , GHC == 8.4.3+                     -- Stackage LTS 12.26.+                   , GHC == 8.4.4+                     -- Stackage LTS 13.10.+                   , GHC == 8.6.3++source-repository head+  type:     git+  location: https://github.com/channable/alfred-margaret++library+  hs-source-dirs:      src+  exposed-modules:     Data.Text.AhoCorasick.Automaton+                     , Data.Text.AhoCorasick.Searcher+                     , Data.Text.AhoCorasick.Replacer+  build-depends:+      base             >= 4.7 && < 5+    , containers       >= 0.6.2 && < 0.7+    , deepseq          >= 1.4.4 && < 1.5+    , hashable         >= 1.3.0 && < 1.4+    , text             >= 1.2.4 && < 1.3+    , primitive        >= 0.7.1 && < 0.8+    , vector           >= 0.12.1 && < 0.13+  ghc-options:         -Wall -Wincomplete-record-updates -Wincomplete-uni-patterns+  default-language:    Haskell2010++test-suite test-suite+  type:                exitcode-stdio-1.0+  main-is:             AhoCorasickSpec.hs+  hs-source-dirs:      tests+  ghc-options:         -Wall -Wincomplete-record-updates -Wno-orphans+  build-depends:       base >= 4.7 && < 5+                     , QuickCheck+                     , alfred-margaret+                     , deepseq+                     , hspec+                     , hspec-expectations+                     , quickcheck-instances+                     , text+  default-language:    Haskell2010

@@ -0,0 +1,659 @@+-- 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.++-- Compile this module with LLVM, rather than with the default code generator.+-- LLVM produces about 20% faster code.+-- We pass -fignore-asserts to improve performance: we ran this code with+-- asserts enabled in production for two months, and in this time, the asserts+-- have not been violated.+{-# OPTIONS_GHC -fllvm -O2 -optlo=-O3 -optlo=-tailcallelim -fignore-asserts #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# 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.+module Data.Text.AhoCorasick.Automaton+  ( AcMachine (..)+  , build+  , runText+  , runLower+  , debugBuildDot+  , CaseSensitivity (..)+  , CodeUnit+  , CodeUnitIndex (..)+  , Match (..)+  , Next (..)+  , lengthUtf16+  , lowerUtf16+  , unpackUtf16+  , unsafeCutUtf16+  , unsafeSliceUtf16+  )+  where++import Prelude hiding (length)++import Control.DeepSeq (NFData)+import Control.Exception (assert)+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 (Word16, Word64)+import GHC.Generics (Generic)+import Data.Primitive.ByteArray (ByteArray (..))++import qualified Data.Char as Char+import qualified Data.IntMap.Strict as IntMap+import qualified Data.List as List+import qualified Data.Text.Array as TextArray+import qualified Data.Text.Unsafe as TextUnsafe+import qualified Data.Vector as Vector+import qualified Data.Vector.Unboxed as UVector+import qualified Data.Vector.Primitive as PVector++data CaseSensitivity+  = CaseSensitive+  | IgnoreCase+  deriving stock (Eq, Generic, Show)+  deriving anyclass (Hashable, NFData)++-- | A numbered state in the Aho-Corasick automaton.+type State = Int++-- | A code unit is a 16-bit integer from which UTF-16 encoded text is built up.+-- The `Text` type is represented as a UTF-16 string.+type CodeUnit = Word16++-- | A transition is a pair of (code unit, next state). The code unit is 16 bits,+-- and the state index is 32 bits. We pack these together as a manually unlifted+-- tuple, because an unboxed Vector of tuples is a tuple of vectors, but we want+-- the elements of the tuple to be adjacent in memory. (The Word64 still needs+-- to be unpacked in the places where it is used.) The code unit is stored in+-- the least significant 32 bits, with the special value 2^16 indicating a+-- wildcard; the "failure" transition. Bit 17 through 31 (starting from zero,+-- both bounds inclusive) are always 0.+--+--  Bit 63 (most significant)                 Bit 0 (least significant)+--  |                                                                 |+--  v                                                                 v+-- |<--       goto state         -->|<-- zeros   -->| |<--   input  -->|+-- |SSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSS|000000000000000|W|IIIIIIIIIIIIIIII|+--                                                   |+--                                                   Wildcard bit (bit 16)+--+type Transition = Word64++-- | An index into the raw UTF-16 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, Generic, Show)+  deriving newtype (Hashable, Ord, Num, Bounded, NFData)++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 two code units.+  , matchValue :: v+  -- ^ The payload associated with the matched needle.+  } deriving (Show, Eq)++-- | 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 :: !(UVector.Vector 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)+  -- ^ For every state, the index into `machineTransitions` where the transition+  -- list for that state starts.+  , machineRootAsciiTransitions :: !(UVector.Vector 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)++-- | The wildcard value is 2^16, one more than the maximal 16-bit code unit.+wildcard :: Integral a => a+wildcard = 0x10000++-- | Extract the code unit from a transition. The special wildcard transition+-- will return 0.+transitionCodeUnit :: Transition -> CodeUnit+transitionCodeUnit t = fromIntegral (t .&. 0xffff)++-- | 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 :: CodeUnit -> State -> Transition+newTransition input state =+  let+    input64 = fromIntegral 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]] -> (UVector.Vector Transition, UVector.Vector Int)+packTransitions transitions =+  let+    packed = UVector.fromList $ concat transitions+    offsets = UVector.fromList $ scanl (+) 0 $ fmap List.length transitions+  in+    (packed, offsets)++-- | Construct an Aho-Corasick automaton for the given needles.+-- Takes a list of code units rather than `Text`, to allow mapping the code+-- units before construction, for example to lowercase individual code points,+-- rather than doing proper case folding (which might change the number of code+-- units).+build :: [([CodeUnit], 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 (fromIntegral 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 :: [[CodeUnit]] -> 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 = \"" ++ show input ++ "\"") state nextState) : edges++    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. [([CodeUnit], v)] -> (Int, TransitionMap, ValuesMap v)+buildTransitionMap =+  let+    go :: State+      -> (Int, TransitionMap, ValuesMap v)+      -> ([CodeUnit], v)+      -> (Int, TransitionMap, ValuesMap v)++    -- 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.+    go !state (!numStates, transitions, values) ([], v) =+      (numStates, transitions, IntMap.insertWith (++) state [v] values)++    -- Follow the edge for the given input from the current state, creating it+    -- if it does not exist.+    go !state (!numStates, transitions, values) (!input : needleTail, vs) =+      let+        transitionsFromState = transitions IntMap.! state+      in+        case IntMap.lookup (fromIntegral input) transitionsFromState of+          Just nextState ->+            go nextState (numStates, transitions, values) (needleTail, vs)+          Nothing ->+            let+              -- Allocate a new state, and insert a transition to it.+              -- Also insert an empty transition map for it.+              nextState = numStates+              transitionsFromState' = IntMap.insert (fromIntegral input) nextState transitionsFromState+              transitions'+                = IntMap.insert state transitionsFromState'+                $ IntMap.insert nextState IntMap.empty+                $ transitions+            in+              go nextState (numStates + 1, transitions', values) (needleTail, vs)++    -- 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+    insertNeedle = go stateInitial+  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 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 ->+  case IntMap.lookup i transitions of+    Just state -> newTransition (fromIntegral 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.+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++-- | 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 `lowerCodeUnit` and `runLower`.+-- 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 =+  let+    Text u16data !initialOffset !initialRemaining = text+    !values = machineValues machine+    !transitions = machineTransitions machine+    !offsets = machineOffsets machine+    !rootAsciiTransitions = machineRootAsciiTransitions machine+    !stateInitial = 0++    -- NOTE: All of the arguments are strict here, because we want to compile+    -- them down to unpacked variables on the stack, or even registers.+    -- The INLINE / NOINLINE annotations here were added to fix a regression we+    -- observed when going from GHC 8.2 to GHC 8.6, and this particular+    -- combination of INLINE and NOINLINE is the fastest one. Removing increases+    -- the benchmark running time by about 9%.++    {-# NOINLINE consumeInput #-}+    consumeInput :: Int -> Int -> a -> State -> a+    consumeInput !offset !remaining !acc !state =+      let+        inputCodeUnit = fromIntegral $ TextArray.unsafeIndex u16data offset+        -- NOTE: Although doing this match here entangles the automaton a bit+        -- with case sensitivity, doing so is faster than passing in a function+        -- that transforms each code unit.+        casedCodeUnit = case caseSensitivity of+          IgnoreCase -> lowerCodeUnit inputCodeUnit+          CaseSensitive -> inputCodeUnit+      in+        case remaining of+          0 -> acc+          _ -> followEdge (offset + 1) (remaining - 1) acc state casedCodeUnit++    {-# INLINE followEdge #-}+    followEdge :: Int -> Int -> a -> State -> CodeUnit -> a+    followEdge !offset !remaining !acc !state !input =+      let+        !tssOffset = offsets `uAt` state+      in+        -- 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).+        if state == stateInitial && input < asciiCount+          then lookupRootAsciiTransition offset remaining acc input+          else lookupTransition offset remaining acc state input tssOffset++    {-# NOINLINE collectMatches #-}+    collectMatches :: Int -> Int -> a -> State -> a+    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 (CodeUnitIndex $ offset - initialOffset) v) of+            Step newAcc -> handleMatch newAcc more+            Done finalAcc -> finalAcc+      in+        handleMatch acc matchedValues++    -- NOTE: there is no `state` argument here, because this case applies only+    -- to the root state `stateInitial`.+    {-# INLINE lookupRootAsciiTransition #-}+    lookupRootAsciiTransition :: Int -> Int -> a -> CodeUnit -> a+    lookupRootAsciiTransition !offset !remaining !acc !input =+      case rootAsciiTransitions `uAt` fromIntegral input of+        t | transitionIsWildcard t -> consumeInput offset remaining acc stateInitial+          | otherwise -> collectMatches offset remaining acc (transitionState t)++    {-# INLINE lookupTransition #-}+    lookupTransition :: Int -> Int -> a -> State -> CodeUnit -> Int -> a+    lookupTransition !offset !remaining !acc !state !input !i =+      case transitions `uAt` i of+        -- 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.+        t | transitionIsWildcard t ->+              if state == stateInitial+                then consumeInput offset remaining acc state+                else followEdge offset remaining acc (transitionState t) input++        -- We found the transition, switch to that new state, collecting matches.+        -- NOTE: This comes after wildcard checking, because the code unit of+        -- the wildcard transition is 0, which is a valid input.+        t | transitionCodeUnit t == input ->+              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.+        _ -> lookupTransition offset remaining acc state input (i + 1)+  in+    consumeInput initialOffset initialRemaining seed stateInitial++-- 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 text+-- on the fly to avoid allocating a second lowercased text array. Lowercasing is+-- applied to individual code units, so the indexes into the lowercased text can+-- be used to index into the original text. 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++-- | Return a Text as a list of UTF-16 code units.+unpackUtf16 :: Text -> [CodeUnit]+unpackUtf16 (Text u16data offset length) =+  let+    go _ 0 = []+    go i n = TextArray.unsafeIndex u16data i : go (i + 1) (n - 1)+  in+    go offset length++-- | Return whether the code unit at the given index starts a surrogate pair.+-- Such a code unit must be followed by a high surrogate in valid UTF-16.+-- Returns false on out of bounds indices.+{-# INLINE isLowSurrogate #-}+isLowSurrogate :: Int -> Text -> Bool+isLowSurrogate !i (Text !u16data !offset !len) =+  let+    w = TextArray.unsafeIndex u16data (offset + i)+  in+    i >= 0 && i < len && w >= 0xd800 && w <= 0xdbff++-- | Return whether the code unit at the given index ends a surrogate pair.+-- Such a code unit must be preceded by a low surrogate in valid UTF-16.+-- Returns false on out of bounds indices.+{-# INLINE isHighSurrogate #-}+isHighSurrogate :: Int -> Text -> Bool+isHighSurrogate !i (Text !u16data !offset !len) =+  let+    w = TextArray.unsafeIndex u16data (offset + i)+  in+    i >= 0 && i < len && w >= 0xdc00 && w <= 0xdfff++-- | Extract a substring from a text, at a code unit offset and length.+-- This is similar to `Text.take length . Text.drop begin`, except that the+-- begin and length are in code *units*, not code points, so we can slice the+-- UTF-16 array, and we don't have to walk the entire text to take surrogate+-- pairs into account. It is the responsibility of the user to not slice+-- surrogate pairs, and to ensure that the length is within bounds, hence this+-- function is unsafe.+{-# INLINE unsafeSliceUtf16 #-}+unsafeSliceUtf16 :: CodeUnitIndex -> CodeUnitIndex -> Text -> Text+unsafeSliceUtf16 (CodeUnitIndex !begin) (CodeUnitIndex !length) !text+  = assert (begin + length <= TextUnsafe.lengthWord16 text)+  $ assert (not $ isHighSurrogate begin text)+  $ assert (not $ isLowSurrogate (begin + length - 1) text)+  $ TextUnsafe.takeWord16 length $ TextUnsafe.dropWord16 begin text++-- | The complement of `unsafeSliceUtf16`: removes the slice, and returns the+-- part before and after. See `unsafeSliceUtf16` for details.+{-# INLINE unsafeCutUtf16 #-}+unsafeCutUtf16 :: CodeUnitIndex -> CodeUnitIndex -> Text -> (Text, Text)+unsafeCutUtf16 (CodeUnitIndex !begin) (CodeUnitIndex !length) !text+  = assert (begin + length <= TextUnsafe.lengthWord16 text)+  $ assert (not $ isHighSurrogate begin text)+  $ assert (not $ isLowSurrogate (begin + length - 1) text)+    ( TextUnsafe.takeWord16 begin text+    , TextUnsafe.dropWord16 (begin + length) text+    )++-- | Return the length of the text, in number of code units.+{-# INLINE lengthUtf16 #-}+lengthUtf16 :: Text -> CodeUnitIndex+lengthUtf16 = CodeUnitIndex . TextUnsafe.lengthWord16++-- | Apply a function to each code unit of a text.+mapUtf16 :: (CodeUnit -> CodeUnit) -> Text -> Text+mapUtf16 f (Text u16data offset length) =+  let+    get !i = f $ TextArray.unsafeIndex u16data (offset + i)+    !(PVector.Vector !offset' !length' !(ByteArray !u16data')) =+      PVector.generate length get+  in+    Text (TextArray.Array u16data') offset' length'++-- | Lowercase each individual code unit of a text without changing their index.+-- This is not a proper case folding, but it does ensure that indices into the+-- lowercased string correspond to indices into the original string.+--+-- 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),+-- is converted into U+0069 only by `lowerUtf16`.+-- 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`.+lowerUtf16 :: Text -> Text+lowerUtf16 = mapUtf16 lowerCodeUnit++{-# INLINE lowerCodeUnit #-}+lowerCodeUnit :: CodeUnit -> CodeUnit+lowerCodeUnit cu =+  if cu >= 0xd800 && cu < 0xe000+     -- This code unit is part of a surrogate pair. Don't touch those, because+     -- we don't have all information required to decode the code point. Note+     -- that alphabets that need to be encoded as surrogate pairs are mostly+     -- archaic and obscure; all of the languages used by our customers have+     -- alphabets in the Basic Multilingual Plane, which does not need surrogate+     -- pairs. Note that the BMP is not just ascii or extended ascii. See also+     -- https://codepoints.net/basic_multilingual_plane.+    then cu+     -- The code unit is a code point on its own (not part of a surrogate pair),+     -- lowercase the code point. These code points, which are all in the BMP,+     -- have the important property that lowercasing them is again a code point+     -- in the BMP, so the output can be encoded in exactly one code unit, just+     -- like the input. This property was verified by exhaustive testing; see+     -- also the test in AhoCorasickSpec.hs.+    else fromIntegral $ Char.ord $ Char.toLower $ Char.chr $ fromIntegral cu

@@ -0,0 +1,208 @@+-- 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.++-- See Automaton.hs for why these GHC flags are here.+{-# OPTIONS_GHC -fllvm -O2 -optlo=-O3 -optlo=-tailcallelim -fno-ignore-asserts #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE DerivingStrategies #-}++-- | Implements sequential string replacements based on the Aho-Corasick algorithm.+module Data.Text.AhoCorasick.Replacer+  ( -- * State machine+    Replacer (..)+  , build+  , compose+  , run+  , runWithLimit+  , Needle+  , Replacement+  , Payload (..)+  ) where++import Control.DeepSeq (NFData)+import Data.Hashable (Hashable)+import Data.List (sort)+import Data.Maybe (fromJust)+import Data.Text (Text)+import GHC.Generics (Generic)++import qualified Data.Text as Text++import Data.Text.AhoCorasick.Automaton (CaseSensitivity (..), CodeUnitIndex)+import Data.Text.AhoCorasick.Searcher (Searcher)++import qualified Data.Text.AhoCorasick.Automaton as Aho+import qualified Data.Text.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+  } deriving (Eq, Generic, Hashable, NFData, Show)++-- | A state machine used for efficient replacements with many different needles.+data Replacer = Replacer+  { replacerCaseSensitivity :: CaseSensitivity+  , replacerSearcher :: Searcher Payload+  }+  deriving stock (Show, Eq, Generic)+  deriving anyclass (Hashable, NFData)++-- | Build an Aho-Corasick automaton that can be used for performing fast+-- sequential replaces.+--+-- Case-insensitive matching performs per-letter language-agnostic case folding.+-- Therefore, it will work in most cases, but not in languages where case folding+-- 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 $ zipWith mapNeedle [0..] replaces+    mapNeedle i (needle, replacement) =+      let+        needle' = case caseSensitivity of+          CaseSensitive -> needle+          IgnoreCase -> Aho.lowerUtf16 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) (Aho.lengthUtf16 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 $ 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 = Text.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) = Aho.unsafeCutUtf16 (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 (Aho.lengthUtf16 initial)+  where+    go [] !acc = acc+    go (Match _ matchLen repl : rest) !acc = go rest (acc - matchLen + Aho.lengthUtf16 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

@@ -0,0 +1,126 @@+-- 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.++-- See AhoCorasick.Automaton for more info about these GHC flags.+-- TL;DR: They make things faster, and we need the flags here because the+-- functions from that module may be inlined into this module.+{-# OPTIONS_GHC -fllvm -O2 -optlo=-O3 -optlo=-tailcallelim -fno-ignore-asserts #-}++{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE FlexibleInstances #-}++module Data.Text.AhoCorasick.Searcher+  ( Searcher+  , build+  , buildWithValues+  , needles+  , numNeedles+  , automaton+  , containsAny+  , containsAnyIgnoreCase+  )+  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.Text.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+  { searcherNeedles :: Hashed [(Text, v)]+  , searcherNumNeedles :: Int+  , searcherAutomaton :: Aho.AcMachine v+  } deriving (Generic)++instance Show (Searcher v) where+  show _ = "Searcher _ _ _"++instance Hashable v => Hashable (Searcher v) where+  hashWithSalt salt searcher = hashWithSalt salt $ searcherNeedles searcher++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 xs nx _ == Searcher ys ny _ = (nx, xs) == (ny, ys)++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 = buildWithValues (needles x <> needles y)++build :: [Text] -> Searcher ()+build = buildWithValues . fmap (\x -> (x, ()))++buildWithValues :: Hashable v => [(Text, v)] -> Searcher v+buildWithValues ns =+  let+    unpack (text, value) = (Aho.unpackUtf16 text, value)+  in+    Searcher (hashed ns) (length ns) $ Aho.build $ fmap unpack ns++needles :: Searcher v -> [(Text, v)]+needles = unhashed . searcherNeedles++numNeedles :: Searcher v -> Int+numNeedles = searcherNumNeedles++automaton :: Searcher v -> Aho.AcMachine v+automaton = searcherAutomaton++-- | Return whether the haystack contains any of the needles.+-- Is case sensitive.+-- 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+    Aho.runText False f (automaton searcher) text++-- | Return whether the haystack contains any of the needles.+-- Is case insensitive. The needles in the searcher should be lowercase.+{-# NOINLINE containsAnyIgnoreCase #-}+containsAnyIgnoreCase :: Searcher () -> Text -> Bool+containsAnyIgnoreCase !searcher !text =+  let+    -- On the first match, return True immediately.+    f _acc _match = Aho.Done True+  in+    Aho.runLower False f (automaton searcher) text

@@ -0,0 +1,369 @@+-- 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 ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module Main where++import Control.DeepSeq (rnf)+import Control.Monad (forM_, unless)+import Data.Foldable (foldl')+import Data.Text (Text)+import Data.Word (Word16)+import GHC.Stack (HasCallStack)+import Prelude hiding (replicate)+import Test.Hspec (Spec, Expectation, describe, it, shouldBe, hspec)+import Test.Hspec.Expectations (shouldMatchList, shouldSatisfy)+import Test.Hspec.QuickCheck (modifyMaxSuccess, modifyMaxSize, prop)+import Test.QuickCheck (Arbitrary (arbitrary, shrink), forAll, forAllShrink, (==>))+import Test.QuickCheck.Gen (Gen)+import Test.QuickCheck.Instances ()++import qualified Data.Char as Char+import qualified Data.Text as Text+import qualified Data.Text.Internal.Search as TextSearch+import qualified Data.Text.Unsafe as TextUnsafe+import qualified Test.QuickCheck as QuickCheck+import qualified Test.QuickCheck.Gen as Gen++import Data.Text.AhoCorasick.Automaton (CaseSensitivity (..))++import qualified Data.Text.AhoCorasick.Automaton as Aho+import qualified Data.Text.AhoCorasick.Replacer as Replacer++instance Arbitrary CaseSensitivity where+  arbitrary = Gen.elements [CaseSensitive, IgnoreCase]++-- | 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+    needleUtf16 = Aho.unpackUtf16 needle+    len = Aho.lengthUtf16 needle+    prependMatch ms match = Aho.Step (match : ms)+    matches = Aho.runText [] prependMatch (Aho.build [(needleUtf16, ())]) needle+  in+    matches `shouldBe` [Aho.Match len ()]++ahoMatch :: [(Text, a)] -> Text -> [Aho.Match a]+ahoMatch needles haystack =+  let+    makeNeedle (text, value) = (Aho.unpackUtf16 text, value)+    needlesUtf16 = fmap makeNeedle needles+    prependMatch matches match = Aho.Step (match : matches)+  in+    Aho.runText [] prependMatch (Aho.build needlesUtf16) haystack++-- | Match without a payload, return only the match positions.+matchPositions :: [Text] -> Text -> [Int]+matchPositions needles haystack =+  let+    withUnit x = (x, ())+    matches = ahoMatch (fmap withUnit needles) haystack+  in+    fmap (Aho.codeUnitIndex . Aho.matchPos) matches++-- | `matchPositions` implemented naively in terms of Text's functionality,+-- which we assume to be correct.+naiveMatchPositions :: [Text] -> Text -> [Int]+naiveMatchPositions needles haystack =+  let+    prependMatch :: [Int] -> Text -> Int -> Text -> [Int]+    prependMatch matches needle offset haystackSlice =+      if Text.null haystack+        then matches+        -- Text.indices returns all non-overlapping occurrences of the needle,+        -- but we want the overlapping ones as well. So we only consider the+        -- first match, and then search again starting from one past the+        -- beginning of the match.+        else case TextSearch.indices needle haystackSlice of+          []  -> matches+          i:_ -> prependMatch (match : matches) needle offset' remainingHaystack+            where+              -- The match index is the index past the end, not the start index.+              match = offset + i + TextUnsafe.lengthWord16 needle+              offset' = offset + i + 1+              remainingHaystack = TextUnsafe.dropWord16 (i + 1) haystackSlice++    prependMatches matches needle = prependMatch matches needle 0 haystack+  in+    foldl' prependMatches [] needles++-- | 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.+arbitraryNeedlesHaystack :: Gen ([Text], Text)+arbitraryNeedlesHaystack = 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++  needles <- Gen.listOf1 (fmap Text.concat genSmall)+  haystack <- fmap Text.concat genBig+  pure (needles, haystack)++main :: IO ()+main = hspec $ describe "Data.Text.AhoCorasick" spec++spec :: Spec+spec = do+  modifyMaxSuccess (const 200) $ do+    describe "build" $ do+      prop "does not throw exceptions" $ \ (kv :: [([Word16], Int)]) ->+        rnf $ Aho.build kv++    describe "unpackUtf16" $ do+      it "unpacks code point U+437b8" $+        -- Note that 0x437b8 lies in the currently unassigned "Plane 5"; the+        -- code point does not currently exist, but that should not bother us.+        -- Check in Python: '\U000437b8'.encode('utf-16be')+        Aho.unpackUtf16 "\x000437b8" `shouldBe` [0xd8cd, 0xdfb8]++      it "unpacks adjacent nulls individually" $ do+        Aho.unpackUtf16 "c\NULe" `shouldBe` [99, 0, 101]+        Aho.unpackUtf16 "bc\NUL\NULe" `shouldBe` [98, 99, 0, 0, 101]++    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 multiple 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.+          matchPositions ["a"] (Text.dropWhile (== 'b') "bbba") `shouldMatchList` [1]++      describe "when given non-ascii inputs" $ do++        -- We have a special lookup table for transitions from the base state+        -- for the first 128 code units, which is always hit for ascii inputs.+        -- Also exercise the fallback code path with a different input.+        it "reports a match if the first haystack character is > U+7f" $ do+          matchPositions ["eclair"] "éclair" `shouldMatchList` []+          matchPositions ["éclair"] "éclair" `shouldMatchList` [6]+          matchPositions ["éclair"] "eclair" `shouldMatchList` []++        it "reports the correct UTF-16 index for surrogate pairs" $ do+          -- Note that the index after the match is 2, even though there is+          -- only a single code point. U+1d11e is encoded as two code units+          -- in UTF-16.+          matchPositions ["𝄞"] "𝄞" `shouldMatchList` [2]++          -- A leviating woman in business suit with dark skin tone needs a+          -- whopping 5 code points to encode, of which the first two need a+          -- surrogate pair in UTF-16, for a total of 7 code units.+          -- U+1f574: man in business suit levitating+          -- U+1f3ff: emoji modifier Fitzpatrick type-6+          -- U+200d:  zero width joiner+          -- U+2640:  female sign+          -- 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.+          matchPositions+            [ "\x1f574\x1f3ff\x200d\x2640\xfe0f"+            , "\x1f574\x1f3ff"+            , "\x1f574"+            ] "\x1f574\x1f3ff\x200d\x2640\xfe0f" `shouldMatchList` [2, 4, 7]++      describe "when given overlapping needles" $ do++        it "finds exactly all matches" $ do+          matchPositions ["foobar", "bar"] "foobar" `shouldMatchList` [6, 6]+          matchPositions ["foobarbaz", "bar"] "xfoobarbazy" `shouldMatchList` [10, 7]+          matchPositions ["foobar", "foo"] "xfoobarbazy" `shouldMatchList` [7, 4]++        it "keeps the value associated with a needle" $ do+          (fmap Aho.matchValue $ ahoMatch [("foo", 'A'), ("bar", 'B')] "foobar")+            `shouldMatchList` ['A', 'B']+          (fmap Aho.matchValue $ ahoMatch [("foo", 'A'), ("bar", 'B')] "foobaz")+            `shouldMatchList` ['A']+          (fmap Aho.matchValue $ ahoMatch [("foo", 'A'), ("bar", 'B')] "foebar")+            `shouldMatchList` ['B']++        it "reports both matches in case of a duplicate needle" $ do+          (fmap Aho.matchValue $ ahoMatch [("foo", 'A'), ("foo", 'B')] "foobar")+            `shouldMatchList` ['A', 'B']++        it "finds all quadratic matches" $+          forM_ ["a", "ab", "abc"] $ \baseText ->+            forM_ [1..33] $ \n ->+              let+                replicate k = Text.replicate k baseText+                needles = fmap replicate [1..n]+                matches = matchPositions needles (replicate n)+              in+                -- The needle of length 1 matches n times, the needle of length+                -- 2 matches n - 1 times, ..., the needle of length n matches+                -- once.+                length matches `shouldBe` sum [1..n]++      describe "when given partially overlapping needles" $ do++        it "finds exactly all matches" $ do+          matchPositions ["ab", "bcd"] "abccd" `shouldMatchList` [2]+          matchPositions ["abc","cde"] "abcdde" `shouldMatchList` [3]+          matchPositions ["c","c\NULe"] "c\NUL\NULe" `shouldMatchList` [1]+          -- The case below is a regression test; it did fail before; it would+          -- report a match at position 5 in addition to position 2.+          matchPositions ["bc","c\NULe"] "bc\NUL\NULe" `shouldMatchList` [2]++      describe "when given empyt needles" $ do++        it "does not report a match" $ do+          matchPositions [""] "" `shouldMatchList` []+          matchPositions [""] "foo" `shouldMatchList` []++      describe "when given random needles and haystacks" $ do++        prop "reports only infixes of the haystack" $+          QuickCheck.forAllShrink arbitraryNeedlesHaystack shrink $ \ (needles, haystack) ->+            let+              dup x = (x, x)+              matches = ahoMatch (fmap dup needles) haystack+              sliceMatch endPos len = Aho.unsafeSliceUtf16 (endPos - len) len haystack+            in+              -- Discard inputs for which there are no matches, to ensure we get+              -- enough coverage for the case where there are matches.+              not (null matches) ==>+                forM_ matches $ \ (Aho.Match pos needle) -> do+                  needle `shouldSatisfy` (`Text.isInfixOf` haystack)+                  sliceMatch pos (Aho.lengthUtf16 needle) `shouldBe` needle++        prop "reports all infixes of the haystack" $+          QuickCheck.forAllShrink arbitraryNeedlesHaystack shrink $ \ (needles, haystack) ->+            matchPositions needles haystack `shouldMatchList` naiveMatchPositions needles haystack++  describe "Char.toLower" $ do++    -- We test that Char.toLower maps the BMP onto itself, because this implies+    -- that changing casing code unit by code unit does not change the number of+    -- code units, which allows us to implement lowercasing in an optimized+    -- manner.+    it "maps the Basic Multilingual Plane onto itself" $+      let+        isSurrogate cu = cu >= 0xd800 && cu < 0xe000+      in+        forM_ [0 .. maxBound :: Aho.CodeUnit] $ \cu -> unless (isSurrogate cu) $+          let+            lower = Char.ord $ Char.toLower $ Char.chr $ fromIntegral cu+          in+            lower `shouldSatisfy` not . isSurrogate++  modifyMaxSize (const 10) $ describe "Replacer.run" $ do+    let+      genHaystack = fmap Text.pack $ Gen.listOf $ Gen.frequency [(40, Gen.elements "abAB"), (1, pure 'İ'), (1, arbitrary)]+      genNeedle = fmap Text.pack $ Gen.resize 3 $ Gen.listOf1 $ Gen.elements "abAB"+      genReplaces = Gen.listOf $ (,) <$> genNeedle <*> arbitrary+      shrinkReplaces = filter (not . any (\(needle, _) -> Text.null needle)) . shrink++      replace needles haystack = Replacer.run (Replacer.build CaseSensitive needles) haystack+      replaceIgnoreCase needles haystack = Replacer.run (Replacer.build 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", "levivating man in business suit")]+        "the \x1f574" `shouldBe` "the levivating 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", "levivating man in business suit")]+        "the \x1f574" `shouldBe` "the levivating 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 CaseSensitive replaces+        replaceText agg (needle, replacement) = Text.replace needle replacement agg+        expected = foldl' replaceText haystack replaces+      in+        Replacer.run replacer haystack `shouldBe` expected