Automaton.hs

-- Alfred-Margaret: Fast Aho-Corasick string searching
-- Copyright 2019 Channable
--
-- Licensed under the 3-clause BSD license, see the LICENSE file in the
-- repository root.

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | An efficient implementation of the Boyer-Moore string search algorithm.
-- http://www-igm.univ-mlv.fr/~lecroq/string/node14.html#SECTION00140
-- https://en.wikipedia.org/wiki/Boyer%E2%80%93Moore_string-search_algorithm
--
-- This module contains a almost 1:1 translation from the C example code in the
-- wikipedia article.
--
-- The algorithm here can be potentially improved by including the Galil rule
-- (https://en.wikipedia.org/wiki/Boyer%E2%80%93Moore_string-search_algorithm#The_Galil_rule)
module Data.Text.BoyerMoore.Automaton
    ( Automaton
    , CaseSensitivity (..)
    , CodeUnitIndex (..)
    , Next (..)
    , buildAutomaton
    , patternLength
    , patternText
    , runText
    ) where

import Prelude hiding (length)

import Control.DeepSeq (NFData)
import Control.Monad (when)
import Control.Monad.ST (runST)
import Data.Hashable (Hashable (..))
import Data.Primitive.Extended
  ( Prim
  , PrimArray
  , indexPrimArray
  , newPrimArray
  , replicateMutablePrimArray
  , unsafeFreezePrimArray
  , writePrimArray
  )
import GHC.Generics (Generic)

#if defined(HAS_AESON)
import qualified Data.Aeson as AE
#endif

import Data.Text.CaseSensitivity (CaseSensitivity (..))
import Data.Text.Utf8 (CodeUnit, CodeUnitIndex (..), Text)

import qualified Data.Text.Utf8 as Utf8

data Next a
  = Done !a
  | Step !a

-- | A Boyer-Moore automaton is based on lookup-tables that allow skipping through the haystack.
-- This allows for sub-linear matching in some cases, as we do not have to look at every input
-- character.
--
-- NOTE: Unlike the AcMachine, a Boyer-Moore automaton only returns non-overlapping matches.
-- This means that a Boyer-Moore automaton is not a 100% drop-in replacement for Aho-Corasick.
--
-- Returning overlapping matches would degrade the performance to /O(nm)/ in pathological cases like
-- finding @aaaa@ in @aaaaa....aaaaaa@ as for each match it would scan back the whole /m/ characters
-- of the pattern.
data Automaton = Automaton
  { automatonPattern :: !Text
  , automatonPatternHash :: !Int  -- ^ Remember our own hash (similar to what 'Hashed' does but our
                                  -- fields are strict).
  , automatonSuffixTable :: !SuffixTable
  , automatonBadCharTable :: !BadCharTable
  }
  deriving stock (Generic, Show)
  deriving anyclass (NFData)

instance Hashable Automaton where
  hashWithSalt salt (Automaton _ patternHash _ _) = hashWithSalt salt patternHash

instance Eq Automaton where
  (Automaton pat1 patHash1 _ _) == (Automaton pat2 patHash2 _ _) =
    patHash1 == patHash2 && pat1 == pat2

#if defined(HAS_AESON)
instance AE.FromJSON Automaton where
  parseJSON v = buildAutomaton <$> AE.parseJSON v

instance AE.ToJSON Automaton where
  toJSON = AE.toJSON . automatonPattern
#endif

buildAutomaton :: Text -> Automaton
buildAutomaton pattern =
  Automaton pattern (hash pattern) (buildSuffixTable pattern) (buildBadCharTable pattern)

-- | Finds all matches in the text, calling the match callback with the *first*
-- matched character of each match of the pattern.
--
-- NOTE: This is unlike Aho-Corasick, which reports the index of the character
-- right after a match.
--
-- NOTE: In the UTF-16 version of this module, there is a function 'Data.Text.BoyerMoore.Automaton.runLower'
-- which does lower-case matching. This function does not exist for the UTF-8 version since it is very
-- tricky to skip code points going backwards without preprocessing the whole input first.
--
-- NOTE: To get full advantage of inlining this function, you probably want to
-- compile the compiling module with -fllvm and the same optimization flags as
-- this module.
runText  :: forall a
  . a
  -> (a -> CodeUnitIndex -> Next a)
  -> Automaton
  -> Text
  -> a
{-# INLINE runText #-}
runText seed f automaton text
  | patLen == 0 = seed
  | otherwise = go seed (patLen - 1)
  where
    -- Use needle as identifier since pattern is potentially a keyword
    Automaton needle _ suffixTable badCharTable = automaton
    patLen = Utf8.lengthUtf8 needle
    stringLen = Utf8.lengthUtf8 text

    codeUnitAt = Utf8.unsafeIndexCodeUnit text

    {-# INLINE go #-}
    go result haystackIndex
      | haystackIndex < stringLen = matchLoop result haystackIndex (patLen - 1)
      | otherwise = result

    -- Compare the needle back-to-front with the haystack
    matchLoop result haystackIndex needleIndex
      | needleIndex >= 0 && codeUnitAt haystackIndex == Utf8.unsafeIndexCodeUnit needle needleIndex =
        -- Characters match, try the pair before
        matchLoop result (haystackIndex - 1) (needleIndex - 1)
      -- We found a match (all needle characters matched)
      | needleIndex < 0 =
        case f result (haystackIndex + 1) of
          Done final -> final
          -- `haystackIndex` now points to the character just before the match starts
          -- Adding `patLen` once points to the last character of the match,
          -- Adding `patLen` once more points to the earliest character where
          -- we can find a non-overlapping match.
          Step intermediate -> go intermediate (haystackIndex + 2 * patLen)
      -- We know it's not a match, the characters differ at the current position
      | otherwise =
        let
          -- The bad character table tells us how far we can advance to the right so that the
          -- character at the current position in the input string, where matching failed,
          -- is lined up with it's rightmost occurrence in the needle.
          -- Note: we could end up left of were we started, essentially never making progress,
          -- if we were to use this rule alone.
          badCharSkip = badCharLookup badCharTable (codeUnitAt haystackIndex)
          suffixSkip = suffixLookup suffixTable needleIndex
          skip = max badCharSkip suffixSkip
        in
          go result (haystackIndex + skip)

-- | Length of the matched pattern measured in UTF-8 code units (bytes).
patternLength :: Automaton -> CodeUnitIndex
patternLength = Utf8.lengthUtf8 . patternText

-- | Return the pattern that was used to construct the automaton.
patternText :: Automaton -> Text
patternText = automatonPattern

-- | The suffix table tells us for each character of the pattern how many characters we can
-- jump ahead if the match fails at that point.
newtype SuffixTable = SuffixTable (PrimArray Int)
  deriving stock (Generic, Show)
  deriving anyclass (NFData)

-- | Lookup an entry in the suffix table.
suffixLookup :: SuffixTable -> CodeUnitIndex -> CodeUnitIndex
{-# INLINE suffixLookup #-}
suffixLookup (SuffixTable table) = CodeUnitIndex . indexTable table . codeUnitIndex

buildSuffixTable :: Text -> SuffixTable
buildSuffixTable pattern = runST $ do
  let patLen = Utf8.lengthUtf8 pattern

  table <- newPrimArray $ codeUnitIndex patLen

  let
    -- Case 1: For each position of the pattern we record the shift that would align the pattern so
    -- that it starts at the longest suffix that is at the same time a prefix, if a mismatch would
    -- happen at that position.
    --
    -- Suppose the length of the pattern is n, a mismatch occurs at position i in the pattern and j
    -- in the haystack, then we know that pattern[i+1..n] == haystack[j+1..j+n-i]. That is, we know
    -- that the part of the haystack that we already matched is a suffix of the pattern.
    -- If the pattern happens to have a prefix that is equal to or a shorter suffix of that matched
    -- suffix, we can shift the pattern to the right so that the pattern starts at the longest
    -- suffix that we have seen that conincides with a prefix of the pattern.
    --
    -- Consider the pattern `ababa`. Then we get
    --
    -- p:                0  1  2  3  4
    -- Pattern:          a  b  a  b  a
    -- lastPrefixIndex:  2  2  4  4  5
    -- table:            6  5  6  5  5
    init1 lastPrefixIndex p
      | p >= 0 = do
        let
          prefixIndex
            | isPrefix pattern (p + 1) = p + 1
            | otherwise = lastPrefixIndex
        writePrimArray table (codeUnitIndex p) (codeUnitIndex $ prefixIndex + patLen - 1 - p)
        init1 prefixIndex (p - 1)
      | otherwise = pure ()

    -- Case 2: We also have to account for the fact that the matching suffix of the pattern might
    -- occur again somewhere within the pattern. In that case, we may not shift as far as if it was
    -- a prefix. That is why the `init2` loop is run after `init1`, potentially overwriting some
    -- entries with smaller shifts.
    init2 p
      | p < patLen - 1 = do
        let
          suffixLen = suffixLength pattern p
        when (Utf8.unsafeIndexCodeUnit pattern (p - suffixLen) /= Utf8.unsafeIndexCodeUnit pattern (patLen - 1 - suffixLen)) $
          writePrimArray table (codeUnitIndex $ patLen - 1 - suffixLen) (codeUnitIndex $ patLen - 1 - p + suffixLen)
        init2 (p + 1)
      | otherwise = pure ()

  init1 (patLen - 1) (patLen - 1)
  init2 0

  SuffixTable <$> unsafeFreezePrimArray table


-- | The bad char table tells us how far we may skip ahead when encountering a certain character
-- in the input string. For example, if there's a character that is not contained in the pattern at
-- all, we can skip ahead until after that character.
data BadCharTable = BadCharTable
  { badCharTableEntries :: {-# UNPACK #-} !(PrimArray Int)
    -- ^ The element type should be CodeUnitIndex, but there's no unboxed vector for that type, and
    -- defining it would be a lot of boilerplate.
  , badCharTablePatternLen :: CodeUnitIndex
  }
  deriving stock (Generic, Show)
  deriving anyclass (NFData)

-- | Number of entries in the fixed-size lookup-table of the bad char table.
badcharTableSize :: Int
{-# INLINE badcharTableSize #-}
badcharTableSize = 256

-- | Lookup an entry in the bad char table.
badCharLookup :: BadCharTable -> CodeUnit -> CodeUnitIndex
{-# INLINE badCharLookup #-}
badCharLookup (BadCharTable asciiTable _patLen) char = CodeUnitIndex $ indexTable asciiTable intChar
  where
    intChar = fromIntegral char

-- | True if the suffix of the @pattern@ starting from @pos@ is a prefix of the pattern
-- For example, @isPrefix \"aabbaa\" 4 == True@.
isPrefix :: Text -> CodeUnitIndex -> Bool
isPrefix needle pos = go 0
  where
    suffixLen = Utf8.lengthUtf8 needle - pos
    go i
      | i < suffixLen =
        -- FIXME: Check whether implementing the linter warning kills tco
        if Utf8.unsafeIndexCodeUnit needle i == Utf8.unsafeIndexCodeUnit needle (pos + i)
          then go (i + 1)
          else False
      | otherwise = True

-- | Length of the longest suffix of the pattern ending on @pos@.
-- For example, @suffixLength \"abaacbbaac\" 4 == 4@, because the substring \"baac\" ends at position
-- 4 and is at the same time the longest suffix that does so, having length 4.
suffixLength :: Text -> CodeUnitIndex -> CodeUnitIndex
suffixLength pattern pos = go 0
  where
    patLen = Utf8.lengthUtf8 pattern
    go i
      | Utf8.unsafeIndexCodeUnit pattern (pos - i) == Utf8.unsafeIndexCodeUnit pattern (patLen - 1 - i) && i < pos = go (i + 1)
      | otherwise = i

buildBadCharTable :: Text -> BadCharTable
buildBadCharTable pattern = runST $ do
  let patLen = Utf8.lengthUtf8 pattern

  -- Initialize table with the maximum skip distance, which is the length of the pattern.
  -- This applies to all characters that are not part of the pattern.
  asciiTable <- replicateMutablePrimArray badcharTableSize $ codeUnitIndex patLen

  let
    -- Fill the bad character table based on the rightmost occurrence of a character in the pattern.
    -- Note that there is also a variant of Boyer-Moore that records all positions (see Wikipedia,
    -- but that requires even more storage space).
    -- Also note that we exclude the last character of the pattern when building the table.
    -- This is because
    --
    -- 1. If the last character does not occur anywhere else in the pattern and we encounter it
    --    during a mismatch, we can advance the pattern to just after that character:
    --
    --    Haystack: aaadcdabcdbb
    --    Pattern:    abcd
    --
    --    In the above example, we would match `d` and `c`, but then fail because `d` != `b`.
    --    Since `d` only occurs at the very last position of the pattern, we can shift to
    --
    --    Haystack: aaadcdabcdbb
    --    Pattern:      abcd
    --
    -- 2. If it does occur anywhere else in the pattern, we can only shift as far as it's necessary
    --    to align it with the haystack:
    --
    --    Haystack: aaadddabcdbb
    --    Pattern:    adcd
    --
    --    We match `d`, and then there is a mismatch `d` != `c`, which allows us to shift only up to:

    --    Haystack: aaadddabcdbb
    --    Pattern:     adcd
    fillTable !i
      -- for(i = 0; i < patLen - 1; i++) {
      | i < patLen - 1 = do
        let patChar = Utf8.unsafeIndexCodeUnit pattern i
        writePrimArray asciiTable (fromIntegral patChar) (codeUnitIndex $ patLen - 1 - i)
        fillTable (i + 1)
      | otherwise = pure ()

  fillTable 0

  asciiTableFrozen <- unsafeFreezePrimArray asciiTable

  pure BadCharTable
    { badCharTableEntries = asciiTableFrozen
    , badCharTablePatternLen = patLen
    }


-- Helper functions for easily toggling the safety of this module

-- | Read from a lookup table at the specified index.
indexTable :: Prim a => PrimArray a -> Int -> a
{-# INLINE indexTable #-}
indexTable = indexPrimArray