{-|
== Z-function

A string matching algorithm. The Z-algorithm generates the function z from a string s, where z(i)
is the length of the longest prefix of s that starts at i. z[0] is set to 0.

Sources:

* https://cp-algorithms.com/string/z-function.html

-}

module ZFunc
    ( zFunc
    , zFuncBS
    ) where

import Control.Monad
import Data.Array.ST
import Data.Array.Unboxed
import qualified Data.ByteString.Char8 as C

-- | Constructs the Z-function. The input sequence should be 0-indexed. O(n).
zFunc :: Eq a => Int -> (Int -> a) -> UArray Int Int
zFunc :: Int -> (Int -> a) -> UArray Int Int
zFunc Int
n Int -> a
at = (forall s. ST s (STUArray s Int Int)) -> UArray Int Int
forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray ((forall s. ST s (STUArray s Int Int)) -> UArray Int Int)
-> (forall s. ST s (STUArray s Int Int)) -> UArray Int Int
forall a b. (a -> b) -> a -> b
$ do
    STUArray s Int Int
z <- (Int, Int) -> Int -> ST s (STUArray s Int Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
0
    let go :: (Int, Int) -> Int -> m (Int, Int)
go lr :: (Int, Int)
lr@(Int
l, Int
r) Int
i = do
            Int
j <- Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) (Int -> Int) -> m Int -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STUArray s Int Int -> Int -> m Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
z (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)
            let k :: Int
k = (Int -> Bool) -> (Int -> Int) -> Int -> Int
forall a. (a -> Bool) -> (a -> a) -> a -> a
until (\Int
k' -> Int
k' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Bool -> Bool -> Bool
|| Int -> a
at (Int
k' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> a
at Int
k') (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)
            STUArray s Int Int -> Int -> Int -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
z Int
i (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)
            (Int, Int) -> m (Int, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int, Int) -> m (Int, Int)) -> (Int, Int) -> m (Int, Int)
forall a b. (a -> b) -> a -> b
$ if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
r then (Int
i, Int
k) else (Int, Int)
lr
    ((Int, Int) -> Int -> ST s (Int, Int))
-> (Int, Int) -> [Int] -> ST s ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ (Int, Int) -> Int -> ST s (Int, Int)
forall (m :: * -> *).
MArray (STUArray s) Int m =>
(Int, Int) -> Int -> m (Int, Int)
go (Int
0, Int
0) [Int
1 .. Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
    STUArray s Int Int -> ST s (STUArray s Int Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure STUArray s Int Int
z

-- | zFunc for a ByteString. O(n).
zFuncBS :: C.ByteString -> UArray Int Int
zFuncBS :: ByteString -> UArray Int Int
zFuncBS ByteString
bs = Int -> (Int -> Char) -> UArray Int Int
forall a. Eq a => Int -> (Int -> a) -> UArray Int Int
zFunc (ByteString -> Int
C.length ByteString
bs) (ByteString -> Int -> Char
C.index ByteString
bs)