-- | Aho–Corasick algorithm is a fast dictionary-matching (multi-pattern matching) algorithm.
--
-- ==== __Example__
--
-- >>> import AtCoder.Extra.AhoCorasick qualified as AC
-- >>> 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 (when)
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
    AhoCorasick -> Int
sizeAc :: {-# UNPACK #-} !Int,
    -- | A trie (-like directed graph) of input words: Vertex -> (Char -> Vertex).
    --
    -- @since 1.5.3.0
    AhoCorasick -> Vector (HashMap Int Int)
trieAc :: !(V.Vector (HM.HashMap Int Int)),
    -- | Node data of links to parent vertex.
    --
    -- @since 1.5.3.0
    AhoCorasick -> Vector Int
parentAc :: !(VU.Vector Int),
    -- | Node data that represents completed pattern string or nothing (@-1@).
    --
    -- @since 1.5.3.0
    AhoCorasick -> Vector Int
patternAc :: !(VU.Vector Int),
    -- | Node data of links to the longest suffix vertex.
    --
    -- @since 1.5.3.0
    AhoCorasick -> Vector Int
suffixAc :: !(VU.Vector Int),
    -- | Node data of links to the longest suffix pattern vertex.
    --
    -- @since 1.5.3.0
    AhoCorasick -> Vector Int
outputAc :: !(VU.Vector Int)
  }

-- | \(O(\sum_i |S_i|)\)
--
-- ==== 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 :: HasCallStack => Vector (Vector Int) -> AhoCorasick
build Vector (Vector Int)
patterns
  | Vector (Vector Int) -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
VG.null Vector (Vector Int)
patterns =
      -- root only
      Int
-> Vector (HashMap Int Int)
-> Vector Int
-> Vector Int
-> Vector Int
-> Vector Int
-> AhoCorasick
AhoCorasick
        Int
1
        (HashMap Int Int -> Vector (HashMap Int Int)
forall a. a -> Vector a
V.singleton HashMap Int Int
forall k v. HashMap k v
HM.empty)
        (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
1 (-Int
1))
        (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
1 Int
0)
        (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
1 Int
0)
        (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
1 Int
0)
  | Bool
otherwise =
      let (!Int
nNodes, !Vector Int
patternMap, !Vector (HashMap Int Int)
trie, !Vector Int
parent) = HasCallStack =>
Vector (Vector Int)
-> (Int, Vector Int, Vector (HashMap Int Int), Vector Int)
Vector (Vector Int)
-> (Int, Vector Int, Vector (HashMap Int Int), Vector Int)
buildTrie Vector (Vector Int)
patterns
          (!Vector Int
suffix, !Vector Int
output) = HasCallStack =>
Int
-> Vector (HashMap Int Int)
-> Vector Int
-> (Vector Int, Vector Int)
Int
-> Vector (HashMap Int Int)
-> Vector Int
-> (Vector Int, Vector Int)
runBfs Int
nNodes Vector (HashMap Int Int)
trie Vector Int
patternMap
       in Int
-> Vector (HashMap Int Int)
-> Vector Int
-> Vector Int
-> Vector Int
-> Vector Int
-> AhoCorasick
AhoCorasick Int
nNodes Vector (HashMap Int Int)
trie Vector Int
parent Vector Int
patternMap Vector Int
suffix Vector Int
output

-- | \(O(1)\) Returns the number of nodes in the trie.
--
-- @since 1.5.3.0
{-# INLINE size #-}
size :: (HasCallStack) => AhoCorasick -> Int
size :: HasCallStack => AhoCorasick -> Int
size = AhoCorasick -> Int
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 :: HasCallStack => AhoCorasick -> Int -> Int -> Int
next AhoCorasick {Vector (HashMap Int Int)
trieAc :: AhoCorasick -> Vector (HashMap Int Int)
trieAc :: Vector (HashMap Int Int)
trieAc, Vector Int
suffixAc :: AhoCorasick -> Vector Int
suffixAc :: Vector Int
suffixAc} Int
v0 Int
c = Int -> Int
inner Int
v0
  where
    inner :: Int -> Int
inner Int
v = case Int -> HashMap Int Int -> Maybe Int
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup Int
c (Vector (HashMap Int Int)
trieAc Vector (HashMap Int Int) -> Int -> HashMap Int Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v) of
      Just Int
end -> Int
end
      Maybe Int
Nothing
        -- no hope
        | Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 -> Int
0
        -- fallback to the longest match suffix
        | Bool
otherwise -> Int -> Int
inner (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$! Vector Int
suffixAc Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v

-- | \(n\) Applies `next` N times for a given input string.
--
-- ==== Constraints
--
-- @since 1.5.3.0
{-# INLINE nextN #-}
nextN ::
  (HasCallStack) =>
  -- | The automaton.
  AhoCorasick ->
  -- | Current node.
  Int ->
  -- | String.
  VU.Vector Int ->
  -- | Resulting node.
  Int
nextN :: HasCallStack => AhoCorasick -> Int -> Vector Int -> Int
nextN AhoCorasick
ac = (Int -> Int -> Int) -> Int -> Vector Int -> Int
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' (HasCallStack => AhoCorasick -> Int -> Int -> Int
AhoCorasick -> Int -> Int -> Int
next AhoCorasick
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 :: HasCallStack => AhoCorasick -> Vector Int -> Vector (Int, Int)
match ac :: AhoCorasick
ac@AhoCorasick {Vector Int
patternAc :: AhoCorasick -> Vector Int
patternAc :: Vector Int
patternAc, Vector Int
outputAc :: AhoCorasick -> Vector Int
outputAc :: Vector Int
outputAc} =
  (Int -> Int -> Vector (Int, Int))
-> Vector Int -> Vector (Int, Int)
forall (v :: * -> *) a b.
(HasCallStack, Vector v a, Vector v b) =>
(Int -> a -> v b) -> v a -> v b
EV.iconcatMap (\Int
i Int
v -> (Int -> Maybe ((Int, Int), Int)) -> Int -> Vector (Int, Int)
forall a b. Unbox a => (b -> Maybe (a, b)) -> b -> Vector a
VU.unfoldr (Int -> Int -> Maybe ((Int, Int), Int)
f Int
i) Int
v) (Vector Int -> Vector (Int, Int))
-> (Vector Int -> Vector Int) -> Vector Int -> Vector (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int) -> Int -> Vector Int -> Vector Int
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
VU.postscanl' (HasCallStack => AhoCorasick -> Int -> Int -> Int
AhoCorasick -> Int -> Int -> Int
next AhoCorasick
ac) Int
0
  where
    f :: Int -> Int -> Maybe ((Int, Int), Int)
    f :: Int -> Int -> Maybe ((Int, Int), Int)
f Int
_ Int
0 = Maybe ((Int, Int), Int)
forall a. Maybe a
Nothing
    f Int
i Int
v = case Vector Int
patternAc Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v of
      -- NOTE: Do not perform early return, as the initial vertex can be non-pattern
      -1 -> Int -> Int -> Maybe ((Int, Int), Int)
f Int
i (Vector Int
outputAc Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v)
      -- NOTE: Here we use `i + 1`, where [pos - patLen, pos) makes up a half-open interval.
      Int
pat -> ((Int, Int), Int) -> Maybe ((Int, Int), Int)
forall a. a -> Maybe a
Just ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
pat), Vector Int
outputAc Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
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 :: HasCallStack =>
Vector (Vector Int)
-> (Int, Vector Int, Vector (HashMap Int Int), Vector Int)
buildTrie Vector (Vector Int)
patternStrings = (forall s.
 ST s (Int, Vector Int, Vector (HashMap Int Int), Vector Int))
-> (Int, Vector Int, Vector (HashMap Int Int), Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s.
  ST s (Int, Vector Int, Vector (HashMap Int Int), Vector Int))
 -> (Int, Vector Int, Vector (HashMap Int Int), Vector Int))
-> (forall s.
    ST s (Int, Vector Int, Vector (HashMap Int Int), Vector Int))
-> (Int, Vector Int, Vector (HashMap Int Int), Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  let !nMaxNodes :: Int
nMaxNodes = (Int
1 +) (Int -> Int) -> (Vector Int -> Int) -> Vector Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int
forall a. Num a => Vector a -> a
V.sum (Vector Int -> Int) -> Vector Int -> Int
forall a b. (a -> b) -> a -> b
$ (Vector Int -> Int) -> Vector (Vector Int) -> Vector Int
forall a b. (a -> b) -> Vector a -> Vector b
V.map Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector (Vector Int)
patternStrings

  -- allocator
  MVector s Int
nNodesVec <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 (Int
1 :: Int)

  -- components
  MVector s (HashMap Int Int)
nextVec <- Int
-> HashMap Int Int
-> ST s (MVector (PrimState (ST s)) (HashMap Int Int))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
VM.replicate Int
nMaxNodes HashMap Int Int
forall k v. HashMap k v
HM.empty
  MVector s Int
parentVec <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nMaxNodes (Int
0 :: Int)

  -- create a trie and collect pattern vertices
  Vector Int
patternVerts <-
    (Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert <$>) (ST s (Vector Int) -> ST s (Vector Int))
-> ((Vector Int -> ST s Int) -> ST s (Vector Int))
-> (Vector Int -> ST s Int)
-> ST s (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Vector Int)
-> (Vector Int -> ST s Int) -> ST s (Vector Int)
forall (m :: * -> *) a b.
Monad m =>
Vector a -> (a -> m b) -> m (Vector b)
V.forM Vector (Vector Int)
patternStrings ((Vector Int -> ST s Int) -> ST s (Vector Int))
-> (Vector Int -> ST s Int) -> ST s (Vector Int)
forall a b. (a -> b) -> a -> b
$
      (Int -> Int -> ST s Int) -> Int -> Vector Int -> ST s Int
forall (m :: * -> *) (v :: * -> *) b a.
(Monad m, Vector v b) =>
(a -> b -> m a) -> a -> v b -> m a
VG.foldM'
        ( \ !Int
u Int
c -> do
            Maybe Int
v0 <- Int -> HashMap Int Int -> Maybe Int
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup Int
c (HashMap Int Int -> Maybe Int)
-> ST s (HashMap Int Int) -> ST s (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) (HashMap Int Int)
-> Int -> ST s (HashMap Int Int)
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s (HashMap Int Int)
MVector (PrimState (ST s)) (HashMap Int Int)
nextVec Int
u
            case Maybe Int
v0 of
              Maybe Int
Nothing -> do
                -- allocate a new vertex index
                Int
v <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
nNodesVec Int
0
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
nNodesVec Int
0 (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                -- store the next vertex link
                MVector (PrimState (ST s)) (HashMap Int Int)
-> (HashMap Int Int -> HashMap Int Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s (HashMap Int Int)
MVector (PrimState (ST s)) (HashMap Int Int)
nextVec (Int -> Int -> HashMap Int Int -> HashMap Int Int
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert Int
c Int
v) Int
u
                -- fill the vertex information
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
parentVec Int
v Int
u
                Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
v
              Just Int
v -> do
                Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
v
        )
        Int
0

  !Int
nNodes <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
nNodesVec Int
0

  let !pattern :: Vector Int
pattern = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        -- We could replace the following with VU.accumulate
        MVector s Int
patVec <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nNodes (-Int
1 :: Int)
        Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
patternVerts ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
iPattern Int
v -> do
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
patVec Int
v Int
iPattern
        MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
patVec
  !Vector (HashMap Int Int)
trie <- Int -> Vector (HashMap Int Int) -> Vector (HashMap Int Int)
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
VG.take Int
nNodes (Vector (HashMap Int Int) -> Vector (HashMap Int Int))
-> ST s (Vector (HashMap Int Int))
-> ST s (Vector (HashMap Int Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) (HashMap Int Int)
-> ST s (Vector (HashMap Int Int))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector s (HashMap Int Int)
MVector (PrimState (ST s)) (HashMap Int Int)
nextVec
  !Vector Int
parent <- Int -> Vector Int -> Vector Int
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
VG.take Int
nNodes (Vector Int -> Vector Int)
-> ST s (Vector Int) -> ST s (Vector Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
parentVec
  (Int, Vector Int, Vector (HashMap Int Int), Vector Int)
-> ST s (Int, Vector Int, Vector (HashMap Int Int), Vector Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
nNodes, Vector Int
pattern, Vector (HashMap Int Int)
trie, Vector Int
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 :: HasCallStack =>
Int
-> Vector (HashMap Int Int)
-> Vector Int
-> (Vector Int, Vector Int)
runBfs Int
nNodes Vector (HashMap Int Int)
trie Vector Int
patternMap = (forall s. ST s (Vector Int, Vector Int))
-> (Vector Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int, Vector Int))
 -> (Vector Int, Vector Int))
-> (forall s. ST s (Vector Int, Vector Int))
-> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
  MVector s Int
suffixVec <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nNodes (Int
0 :: Int)
  MVector s Int
outputVec <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nNodes (Int
0 :: Int)

  Queue s Int
que <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (Queue (PrimState m) a)
Q.new @_ @Int Int
nNodes
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (HashMap Int Int -> [Int]
forall k v. HashMap k v -> [v]
HM.elems (Vector (HashMap Int Int)
trie Vector (HashMap Int Int) -> Int -> HashMap Int Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
0)) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
v -> do
    Queue (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushBack Queue s Int
Queue (PrimState (ST s)) Int
que Int
v

  -- TODO: deduplicate with `next`
  let nextM :: Int -> Int -> f Int
nextM Int
c Int
u = case Int -> HashMap Int Int -> Maybe Int
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup Int
c (Vector (HashMap Int Int)
trie Vector (HashMap Int Int) -> Int -> HashMap Int Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
u) of
        Just Int
end -> Int -> f Int
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
end
        Maybe Int
Nothing
          | Int
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 -> Int -> f Int
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
          | Bool
otherwise -> do
              Int
v <- MVector (PrimState f) Int -> Int -> f Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState f) Int
suffixVec Int
u
              Int -> Int -> f Int
nextM Int
c Int
v

  (ST s () -> ST s ()) -> ST s ()
forall a. (a -> a) -> a
fix ((ST s () -> ST s ()) -> ST s ())
-> (ST s () -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ST s ()
popLoop -> do
    Maybe Int
q <- Queue (PrimState (ST s)) Int -> ST s (Maybe Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> m (Maybe a)
Q.popFront Queue s Int
Queue (PrimState (ST s)) Int
que
    case Maybe Int
q of
      Maybe Int
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Just Int
u -> do
        -- visit neighbors
        [(Int, Int)] -> ((Int, Int) -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (HashMap Int Int -> [(Int, Int)]
forall k v. HashMap k v -> [(k, v)]
HM.toList (Vector (HashMap Int Int)
trie Vector (HashMap Int Int) -> Int -> HashMap Int Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
u)) (((Int, Int) -> ST s ()) -> ST s ())
-> ((Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
c, !Int
v) -> do
          Queue (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Queue (PrimState m) a -> a -> m ()
Q.pushBack Queue s Int
Queue (PrimState (ST s)) Int
que Int
v

          -- find the longest suffix to continue with `c`
          !Int
suffix <- Int -> Int -> ST s Int
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
Int -> Int -> f Int
nextM Int
c (Int -> ST s Int) -> ST s Int -> ST s Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
suffixVec Int
u
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
suffixVec Int
v Int
suffix

          -- find the longest suffix that matches to a pattern
          let suffixPattern :: Int
suffixPattern = Vector Int
patternMap Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
suffix
          Int
output <-
            if Int
suffixPattern Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1
              then Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
suffix
              else MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
outputVec Int
suffix
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
outputVec Int
v Int
output

        -- loop
        ST s ()
popLoop

  (,) (Vector Int -> Vector Int -> (Vector Int, Vector Int))
-> ST s (Vector Int)
-> ST s (Vector Int -> (Vector Int, Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
suffixVec ST s (Vector Int -> (Vector Int, Vector Int))
-> ST s (Vector Int) -> ST s (Vector Int, Vector Int)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
outputVec