content: -- | Aho–Corasick algorithm is a fast dictionary-matching (multi-pattern matching) algorithm.
  --
  -- ==== __Example__
  --
  -- >>> import AtCoder.Extra.AhoCorasick qualified as AC
  -- >>> import Data.Vector qualified as V
  -- >>> import Data.Vector.Unboxed qualified as VU
  --
  -- Pattern strings must be given as @V.Vector (VU.Vector Int)@:
  --
  -- >>> let patterns = V.fromList [VU.fromList [0, 1], VU.fromList [0, 2], VU.fromList [2, 3, 4]]
  -- >>> let ac = AC.build patterns
  -- >>> AC.size ac
  -- 7
  --
  -- The automaton could be run manually with `next` or `nextN`:
  --
  -- >>> AC.nextN ac {- empty node -} 0 (VU.fromList [0, 2, 3])
  -- 5
  --
  -- `match` returns a vector of @(endPos, patternId)@:
  --
  -- >>> --                         [.....) pattern 0
  -- >>> --                               [.......) pattern2
  -- >>> AC.match ac $ VU.fromList [0, 1, 2, 3, 4]
  -- [(2,0),(5,2)]
  --
  -- If you need a vector of @(startPos, patternId)@, you must manually map the result:
  --
  -- >>> let f (!end, !patId) = (end - VU.length (patterns V.! patId), patId)
  -- >>> --                                    [.....) pattern 0
  -- >>> --                                          [.......) pattern2
  -- >>> VU.map f . AC.match ac $ VU.fromList [0, 1, 2, 3, 4]
  -- [(0,0),(2,2)]
  --
  -- Note that duplicate patterns are only counted once with `match`.
  --
  -- @since 1.5.3.0
  module AtCoder.Extra.AhoCorasick
    ( AhoCorasick (..),
      build,
      size,
      next,
      nextN,
      match,
    )
  where
  
  -- TODO: Generalize with Hash + Unbox? Int-only implementation is faster though.
  
  import AtCoder.Extra.Vector qualified as EV
  import AtCoder.Internal.Queue qualified as Q
  import Control.Monad.Fix (fix)
  import Control.Monad.ST (runST)
  import Data.Foldable (for_)
  import Data.HashMap.Strict qualified as HM
  import Data.Vector qualified as V
  import Data.Vector.Generic qualified as VG
  import Data.Vector.Generic.Mutable qualified as VGM
  import Data.Vector.Mutable qualified as VM
  import Data.Vector.Unboxed qualified as VU
  import Data.Vector.Unboxed.Mutable qualified as VUM
  import GHC.Stack (HasCallStack)
  
  -- | Aho–Corasick algorithm data.
  --
  -- @since 1.5.3.0
  data AhoCorasick = AhoCorasick
    { -- | The number of nodes in the trie.
      --
      -- @since 1.5.3.0
      sizeAc :: {-# UNPACK #-} !Int,
      -- | A trie (-like directed graph) of input words: Vertex -> (Char -> Vertex).
      --
      -- @since 1.5.3.0
      trieAc :: !(V.Vector (HM.HashMap Int Int)),
      -- | Node data of links to parent vertex.
      --
      -- @since 1.5.3.0
      parentAc :: !(VU.Vector Int),
      -- | Node data that represents completed pattern string or nothing (@-1@).
      --
      -- @since 1.5.3.0
      patternAc :: !(VU.Vector Int),
      -- | Node data of links to the longest suffix vertex.
      --
      -- @since 1.5.3.0
      suffixAc :: !(VU.Vector Int),
      -- | Node data of links to the longest suffix pattern vertex.
      --
      -- @since 1.5.3.0
      outputAc :: !(VU.Vector Int)
    }
  
  -- | \(O(\sum_i |S_i|)\) Build an Aho–Corasick automaton.
  --
  -- ==== Constraints
  -- - \(|S_i| > 0\)
  --
  -- @since 1.5.3.0
  {-# INLINEABLE build #-}
  build ::
    (HasCallStack) =>
    -- | Pattern strings.
    V.Vector (VU.Vector Int) ->
    -- | Aho–Corasick automaton based on a trie.
    AhoCorasick
  build patterns
    | VG.null patterns =
        -- root only
        AhoCorasick
          1
          (V.singleton HM.empty)
          (VU.replicate 1 (-1))
          (VU.replicate 1 0)
          (VU.replicate 1 0)
          (VU.replicate 1 0)
    | otherwise =
        let (!nNodes, !patternMap, !trie, !parent) = buildTrie patterns
            (!suffix, !output) = runBfs nNodes trie patternMap
         in AhoCorasick nNodes trie parent patternMap suffix output
  
  -- | \(O(1)\) Returns the number of nodes in the trie.
  --
  -- @since 1.5.3.0
  {-# INLINE size #-}
  size :: (HasCallStack) => AhoCorasick -> Int
  size = sizeAc
  
  -- | \(O(1)\) Retrieves the next node to visit.
  --
  -- @since 1.5.3.0
  {-# INLINEABLE next #-}
  -- TODO: benchmark INLINE
  next ::
    (HasCallStack) =>
    -- | The automaton.
    AhoCorasick ->
    -- | Current node ID (empty node is @0@).
    Int ->
    -- | Character.
    Int ->
    -- | Next node ID.
    Int
  next AhoCorasick {trieAc, suffixAc} v0 c = inner v0
    where
      inner v = case HM.lookup c (trieAc VG.! v) of
        Just end -> end
        Nothing
          -- no hope
          | v == 0 -> 0
          -- fallback to the longest match suffix
          | otherwise -> inner $! suffixAc VG.! v
  
  -- | \(O(n)\) Applies `next` \(n\) times for the input node.
  --
  -- ==== Constraints
  --
  -- @since 1.5.3.0
  {-# INLINE nextN #-}
  nextN ::
    (HasCallStack) =>
    -- | The automaton.
    AhoCorasick ->
    -- | Current node.
    Int ->
    -- | String.
    VU.Vector Int ->
    -- | Resulting node.
    Int
  nextN ac = VU.foldl' (next ac)
  
  -- | \(O(|T|)\) Runs dictionary matching (multi-pattern matching) in linear time and returns a list
  -- of @(endPos, patId)@, where @[endPos - patLen, endPos)@ corresponds to the interval of original
  -- source slice.
  --
  -- Note that duplicate patterns are counted just once with one of them; if pattern A and B are the
  -- same, their appearence is counted as either A or B.
  --
  -- @since 1.5.3.0
  {-# INLINEABLE match #-}
  match :: (HasCallStack) => AhoCorasick -> VU.Vector Int -> VU.Vector (Int, Int)
  match ac@AhoCorasick {patternAc, outputAc} =
    EV.iconcatMap (\i v -> VU.unfoldr (f i) v) . VU.postscanl' (next ac) 0
    where
      f :: Int -> Int -> Maybe ((Int, Int), Int)
      f _ 0 = Nothing
      f i v = case patternAc VG.! v of
        -- NOTE: Do not perform early return, as the initial vertex can be non-pattern
        -1 -> f i (outputAc VG.! v)
        -- NOTE: Here we use `i + 1`, where [pos - patLen, pos) makes up a half-open interval.
        pat -> Just ((i + 1, pat), outputAc VG.! v)
  
  -- | \(O(\sum_i |S_i| \Gamma)\)
  {-# INLINEABLE buildTrie #-}
  buildTrie :: (HasCallStack) => V.Vector (VU.Vector Int) -> (Int, VU.Vector Int, V.Vector (HM.HashMap Int Int), VU.Vector Int)
  buildTrie patternStrings = runST $ do
    let !nMaxNodes = (1 +) . V.sum $ V.map VG.length patternStrings
  
    -- allocator
    nNodesVec <- VUM.replicate 1 (1 :: Int)
  
    -- components
    nextVec <- VM.replicate nMaxNodes HM.empty
    parentVec <- VUM.replicate nMaxNodes (0 :: Int)
  
    -- create a trie and collect pattern vertices
    patternVerts <-
      (VU.convert <$>) . V.forM patternStrings $
        VG.foldM'
          ( \ !u c -> do
              v0 <- HM.lookup c <$> VGM.read nextVec u
              case v0 of
                Nothing -> do
                  -- allocate a new vertex index
                  v <- VGM.read nNodesVec 0
                  VGM.write nNodesVec 0 $! v + 1
                  -- store the next vertex link
                  VGM.modify nextVec (HM.insert c v) u
                  -- fill the vertex information
                  VGM.write parentVec v u
                  pure v
                Just v -> do
                  pure v
          )
          0
  
    !nNodes <- VGM.read nNodesVec 0
  
    let !pattern = VU.create $ do
          -- We could replace the following with VU.accumulate
          patVec <- VUM.replicate nNodes (-1 :: Int)
          VU.iforM_ patternVerts $ \iPattern v -> do
            VGM.write patVec v iPattern
          pure patVec
    !trie <- VG.take nNodes <$> V.unsafeFreeze nextVec
    !parent <- VG.take nNodes <$> VU.unsafeFreeze parentVec
    pure (nNodes, pattern, trie, parent)
  
  -- | \(O(\sum_i |S_i| \Gamma)\)
  {-# INLINEABLE runBfs #-}
  runBfs :: (HasCallStack) => Int -> V.Vector (HM.HashMap Int Int) -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)
  runBfs nNodes trie patternMap = runST $ do
    suffixVec <- VUM.replicate nNodes (0 :: Int)
    outputVec <- VUM.replicate nNodes (0 :: Int)
  
    que <- Q.new @_ @Int nNodes
    for_ (HM.elems (trie VG.! 0)) $ \v -> do
      Q.pushBack que v
  
    -- TODO: deduplicate with `next`
    let nextM c u = case HM.lookup c (trie VG.! u) of
          Just end -> pure end
          Nothing
            | u == 0 -> pure 0
            | otherwise -> do
                v <- VGM.read suffixVec u
                nextM c v
  
    fix $ \popLoop -> do
      q <- Q.popFront que
      case q of
        Nothing -> pure ()
        Just u -> do
          -- visit neighbors
          for_ (HM.toList (trie VG.! u)) $ \(!c, !v) -> do
            Q.pushBack que v
  
            -- find the longest suffix to continue with `c`
            !suffix <- nextM c =<< VGM.read suffixVec u
            VGM.write suffixVec v suffix
  
            -- find the longest suffix that matches to a pattern
            let suffixPattern = patternMap VG.! suffix
            output <-
              if suffixPattern /= -1
                then pure suffix
                else VGM.read outputVec suffix
            VGM.write outputVec v output
  
          -- loop
          popLoop
  
    (,) <$> VU.unsafeFreeze suffixVec <*> VU.unsafeFreeze outputVec
  