diff --git a/LICENSE b/LICENSE
new file mode 100644
--- /dev/null
+++ b/LICENSE
@@ -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.
diff --git a/README.md b/README.md
new file mode 100644
--- /dev/null
+++ b/README.md
@@ -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:
+
+
+
+
+
+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
diff --git a/Setup.hs b/Setup.hs
new file mode 100644
--- /dev/null
+++ b/Setup.hs
@@ -0,0 +1,2 @@
+import Distribution.Simple
+main = defaultMain
diff --git a/alfred-margaret.cabal b/alfred-margaret.cabal
new file mode 100644
--- /dev/null
+++ b/alfred-margaret.cabal
@@ -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
+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
diff --git a/src/Data/Text/AhoCorasick/Automaton.hs b/src/Data/Text/AhoCorasick/Automaton.hs
new file mode 100644
--- /dev/null
+++ b/src/Data/Text/AhoCorasick/Automaton.hs
@@ -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
diff --git a/src/Data/Text/AhoCorasick/Replacer.hs b/src/Data/Text/AhoCorasick/Replacer.hs
new file mode 100644
--- /dev/null
+++ b/src/Data/Text/AhoCorasick/Replacer.hs
@@ -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
diff --git a/src/Data/Text/AhoCorasick/Searcher.hs b/src/Data/Text/AhoCorasick/Searcher.hs
new file mode 100644
--- /dev/null
+++ b/src/Data/Text/AhoCorasick/Searcher.hs
@@ -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
diff --git a/tests/AhoCorasickSpec.hs b/tests/AhoCorasickSpec.hs
new file mode 100644
--- /dev/null
+++ b/tests/AhoCorasickSpec.hs
@@ -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