{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}

-- | Rolling hash algorithm implemented as a monoid, typically stored in a segment tree. The type
-- parameters \(b\) and \(p\) represent the B-adic base and the modulus, respectively.
--
-- Combining `RollingHash` with `SegTree` enables \(O(\log |s|)\) string slice creation and
-- \(O(1)\) slice comparison.
--
-- @since 1.1.0.0
module AtCoder.Extra.Monoid.RollingHash
  ( -- * Rolling hash
    RollingHash (..),

    -- * Constructors
    new,
    unsafeNew,
  )
where

import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Exts (proxy#)
import GHC.TypeNats (KnownNat, natVal')

-- | Rolling hash algorithm implemented as a monoid, typically stored in a segment tree. The type
-- parameters \(b\) and \(p\) represent the B-adic base and the modulus, respectively.
--
-- Combining `RollingHash` with `SegTree` enables \(O(\log |s|)\) string slice creation and
-- \(O(1)\) slice comparison.
--
--
-- ==== __Example__
-- It's convenient to define a type alias of `RollingHash`:
--
-- >>> import AtCoder.Extra.Monoid.RollingHash qualified as RH
-- >>> import AtCoder.SegTree qualified as ST
-- >>> import Data.Char (ord)
-- >>> import Data.Semigroup (Dual (..))
-- >>> type RH = RH.RollingHash 100 998244353
--
-- Let's test whether "abcba" is a palindrome:
--
-- >>> seg <- ST.build @_ @RH . VU.map (RH.unsafeNew . ord) $ VU.fromList "abcba"
-- >>> seg' <- ST.build @_ @(Dual RH) . VU.map (Dual . RH.unsafeNew . ord) $ VU.fromList "abcba"
-- >>> hash1 <- ST.prod seg 2 5       --   cba  (left to right)
-- >>> Dual hash2 <- ST.prod seg' 0 3 -- abc    (right to lett)
-- >>> hash1 == hash2
-- True
--
-- @since 1.1.0.0
data RollingHash b p = RollingHash
  { -- | The hash value.
    forall {k} {k} (b :: k) (p :: k). RollingHash b p -> Int
hashRH :: {-# UNPACK #-} !Int,
    -- | \(b^{\mathrm{length}} \bmod p\).
    forall {k} {k} (b :: k) (p :: k). RollingHash b p -> Int
nextDigitRH :: {-# UNPACK #-} !Int
  }
  deriving
    ( -- | @since 1.1.0.0
      RollingHash b p -> RollingHash b p -> Bool
(RollingHash b p -> RollingHash b p -> Bool)
-> (RollingHash b p -> RollingHash b p -> Bool)
-> Eq (RollingHash b p)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (b :: k) k (p :: k).
RollingHash b p -> RollingHash b p -> Bool
$c== :: forall k (b :: k) k (p :: k).
RollingHash b p -> RollingHash b p -> Bool
== :: RollingHash b p -> RollingHash b p -> Bool
$c/= :: forall k (b :: k) k (p :: k).
RollingHash b p -> RollingHash b p -> Bool
/= :: RollingHash b p -> RollingHash b p -> Bool
Eq,
      -- | @since 1.1.0.0
      Int -> RollingHash b p -> ShowS
[RollingHash b p] -> ShowS
RollingHash b p -> String
(Int -> RollingHash b p -> ShowS)
-> (RollingHash b p -> String)
-> ([RollingHash b p] -> ShowS)
-> Show (RollingHash b p)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k (b :: k) k (p :: k). Int -> RollingHash b p -> ShowS
forall k (b :: k) k (p :: k). [RollingHash b p] -> ShowS
forall k (b :: k) k (p :: k). RollingHash b p -> String
$cshowsPrec :: forall k (b :: k) k (p :: k). Int -> RollingHash b p -> ShowS
showsPrec :: Int -> RollingHash b p -> ShowS
$cshow :: forall k (b :: k) k (p :: k). RollingHash b p -> String
show :: RollingHash b p -> String
$cshowList :: forall k (b :: k) k (p :: k). [RollingHash b p] -> ShowS
showList :: [RollingHash b p] -> ShowS
Show
    )

-- | \(O(1)\) Creates a one-length `RollingHash` from an integer.
--
-- @since 1.1.0.0
{-# INLINE new #-}
new :: forall b p. (KnownNat b, KnownNat p) => Int -> RollingHash b p
new :: forall (b :: Nat) (p :: Nat).
(KnownNat b, KnownNat p) =>
Int -> RollingHash b p
new Int
h = Int -> Int -> RollingHash b p
forall {k} {k} (b :: k) (p :: k). Int -> Int -> RollingHash b p
RollingHash (Int
h Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy# p -> Nat
forall (n :: Nat). KnownNat n => Proxy# n -> Nat
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p))) (Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy# b -> Nat
forall (n :: Nat). KnownNat n => Proxy# n -> Nat
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @b)))

-- | \(O(1)\) Creates a one-length `RollingHash` from an integer without taking the mod.
--
-- @since 1.1.0.0
{-# INLINE unsafeNew #-}
unsafeNew :: forall b p. (KnownNat b, KnownNat p) => Int -> RollingHash b p
unsafeNew :: forall (b :: Nat) (p :: Nat).
(KnownNat b, KnownNat p) =>
Int -> RollingHash b p
unsafeNew Int
h = Int -> Int -> RollingHash b p
forall {k} {k} (b :: k) (p :: k). Int -> Int -> RollingHash b p
RollingHash Int
h (Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy# b -> Nat
forall (n :: Nat). KnownNat n => Proxy# n -> Nat
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @b)))

-- | @since 1.1.0.0
instance (KnownNat b, KnownNat p) => Semigroup (RollingHash b p) where
  -- \| \(O(1)\)
  {-# INLINE (<>) #-}
  (RollingHash !Int
digit1 !Int
hash1) <> :: RollingHash b p -> RollingHash b p -> RollingHash b p
<> (RollingHash !Int
digit2 !Int
hash2) = Int -> Int -> RollingHash b p
forall {k} {k} (b :: k) (p :: k). Int -> Int -> RollingHash b p
RollingHash Int
digit' Int
hash'
    where
      !p :: Int
p = Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Nat -> Int) -> Nat -> Int
forall a b. (a -> b) -> a -> b
$ Proxy# p -> Nat
forall (n :: Nat). KnownNat n => Proxy# n -> Nat
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p)
      !digit' :: Int
digit' = Int
digit1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
digit2 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
p
      !hash' :: Int
hash' = (Int
hash1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
digit2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hash2) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
p

-- | @since 1.1.0.0
instance (KnownNat b, KnownNat p) => Monoid (RollingHash b p) where
  {-# INLINE mempty #-}
  mempty :: RollingHash b p
mempty = Int -> Int -> RollingHash b p
forall {k} {k} (b :: k) (p :: k). Int -> Int -> RollingHash b p
RollingHash Int
1 Int
0

type RHRepr = (Int, Int)

-- | @since 1.1.0.0
instance VU.IsoUnbox (RollingHash b p) RHRepr where
  {-# INLINE toURepr #-}
  toURepr :: RollingHash b p -> RHRepr
toURepr (RollingHash Int
a Int
b) = (Int
a, Int
b)
  {-# INLINE fromURepr #-}
  fromURepr :: RHRepr -> RollingHash b p
fromURepr (!Int
a, !Int
b) = Int -> Int -> RollingHash b p
forall {k} {k} (b :: k) (p :: k). Int -> Int -> RollingHash b p
RollingHash Int
a Int
b

-- | @since 1.1.0.0
newtype instance VU.MVector s (RollingHash b p) = MV_RH (VUM.MVector s RHRepr)

-- | @since 1.1.0.0
newtype instance VU.Vector (RollingHash b p) = V_RH (VU.Vector RHRepr)

-- | @since 1.1.0.0
deriving via (RollingHash b p `VU.As` RHRepr) instance VGM.MVector VUM.MVector (RollingHash b p)

-- | @since 1.1.0.0
deriving via (RollingHash b p `VU.As` RHRepr) instance VG.Vector VU.Vector (RollingHash b p)

-- | @since 1.1.0.0
instance VU.Unbox (RollingHash b p)