{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.RollingHash where

import Data.Proxy
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U
import GHC.Exts
import GHC.TypeLits
import System.Random.Stateful

{- |
modulo @2^61-1@

@b@ should be a primitive root of @2^61-1@
-}
newtype RollingHash (b :: Nat) = RollingHash {forall (b :: Nat). RollingHash b -> Int
getRollingHash :: Int}
  deriving newtype (RollingHash b -> RollingHash b -> Bool
(RollingHash b -> RollingHash b -> Bool)
-> (RollingHash b -> RollingHash b -> Bool) -> Eq (RollingHash b)
forall (b :: Nat). RollingHash b -> RollingHash b -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (b :: Nat). RollingHash b -> RollingHash b -> Bool
== :: RollingHash b -> RollingHash b -> Bool
$c/= :: forall (b :: Nat). RollingHash b -> RollingHash b -> Bool
/= :: RollingHash b -> RollingHash b -> Bool
Eq, Eq (RollingHash b)
Eq (RollingHash b) =>
(RollingHash b -> RollingHash b -> Ordering)
-> (RollingHash b -> RollingHash b -> Bool)
-> (RollingHash b -> RollingHash b -> Bool)
-> (RollingHash b -> RollingHash b -> Bool)
-> (RollingHash b -> RollingHash b -> Bool)
-> (RollingHash b -> RollingHash b -> RollingHash b)
-> (RollingHash b -> RollingHash b -> RollingHash b)
-> Ord (RollingHash b)
RollingHash b -> RollingHash b -> Bool
RollingHash b -> RollingHash b -> Ordering
RollingHash b -> RollingHash b -> RollingHash b
forall (b :: Nat). Eq (RollingHash b)
forall (b :: Nat). RollingHash b -> RollingHash b -> Bool
forall (b :: Nat). RollingHash b -> RollingHash b -> Ordering
forall (b :: Nat). RollingHash b -> RollingHash b -> RollingHash b
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: forall (b :: Nat). RollingHash b -> RollingHash b -> Ordering
compare :: RollingHash b -> RollingHash b -> Ordering
$c< :: forall (b :: Nat). RollingHash b -> RollingHash b -> Bool
< :: RollingHash b -> RollingHash b -> Bool
$c<= :: forall (b :: Nat). RollingHash b -> RollingHash b -> Bool
<= :: RollingHash b -> RollingHash b -> Bool
$c> :: forall (b :: Nat). RollingHash b -> RollingHash b -> Bool
> :: RollingHash b -> RollingHash b -> Bool
$c>= :: forall (b :: Nat). RollingHash b -> RollingHash b -> Bool
>= :: RollingHash b -> RollingHash b -> Bool
$cmax :: forall (b :: Nat). RollingHash b -> RollingHash b -> RollingHash b
max :: RollingHash b -> RollingHash b -> RollingHash b
$cmin :: forall (b :: Nat). RollingHash b -> RollingHash b -> RollingHash b
min :: RollingHash b -> RollingHash b -> RollingHash b
Ord, Int -> RollingHash b -> ShowS
[RollingHash b] -> ShowS
RollingHash b -> String
(Int -> RollingHash b -> ShowS)
-> (RollingHash b -> String)
-> ([RollingHash b] -> ShowS)
-> Show (RollingHash b)
forall (b :: Nat). Int -> RollingHash b -> ShowS
forall (b :: Nat). [RollingHash b] -> ShowS
forall (b :: Nat). RollingHash b -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (b :: Nat). Int -> RollingHash b -> ShowS
showsPrec :: Int -> RollingHash b -> ShowS
$cshow :: forall (b :: Nat). RollingHash b -> String
show :: RollingHash b -> String
$cshowList :: forall (b :: Nat). [RollingHash b] -> ShowS
showList :: [RollingHash b] -> ShowS
Show)

instance Num (RollingHash b) where
  (RollingHash (I# Int#
x#)) + :: RollingHash b -> RollingHash b -> RollingHash b
+ (RollingHash (I# Int#
y#)) = case Int#
x# Int# -> Int# -> Int#
+# Int#
y# of
    Int#
xy# -> case Int#
xy# Int# -> Int# -> Int#
>=# Int#
0x1fffffffffffffff# of
      Int#
b# -> Int -> RollingHash b
forall (b :: Nat). Int -> RollingHash b
RollingHash (Int# -> Int
I# (Int#
xy# Int# -> Int# -> Int#
+# Int#
b# Int# -> Int# -> Int#
-# Int# -> Int# -> Int#
uncheckedIShiftL# Int#
b# Int#
61#))
  {-# INLINE (+) #-}
  (RollingHash (I# Int#
x#)) - :: RollingHash b -> RollingHash b -> RollingHash b
- (RollingHash (I# Int#
y#)) = case Int#
x# Int# -> Int# -> Int#
-# Int#
y# of
    Int#
xy# -> case Int#
xy# Int# -> Int# -> Int#
<# Int#
0# of
      Int#
b# -> Int -> RollingHash b
forall (b :: Nat). Int -> RollingHash b
RollingHash (Int# -> Int
I# (Int#
xy# Int# -> Int# -> Int#
-# Int#
b# Int# -> Int# -> Int#
+# Int# -> Int# -> Int#
uncheckedIShiftL# Int#
b# Int#
61#))
  {-# INLINE (-) #-}
  (RollingHash (I# Int#
x#)) * :: RollingHash b -> RollingHash b -> RollingHash b
* (RollingHash (I# Int#
y#)) = case Word# -> Word# -> (# Word#, Word# #)
timesWord2# (Int# -> Word#
int2Word# Int#
x#) (Int# -> Word#
int2Word# Int#
y#) of
    (# Word#
hi#, Word#
lo# #) ->
      case Int# -> Int# -> Int#
uncheckedIShiftL# (Word# -> Int#
word2Int# Word#
hi#) Int#
3#
        Int# -> Int# -> Int#
+# Int# -> Int# -> Int#
uncheckedIShiftRL# (Word# -> Int#
word2Int# Word#
lo#) Int#
61#
        Int# -> Int# -> Int#
+# Int# -> Int# -> Int#
andI# (Word# -> Int#
word2Int# Word#
lo#) Int#
0x1fffffffffffffff# of
        Int#
qr# -> case Int#
qr# Int# -> Int# -> Int#
># Int#
0x1fffffffffffffff# of
          Int#
b# -> Int -> RollingHash b
forall (b :: Nat). Int -> RollingHash b
RollingHash (Int# -> Int
I# (Int#
qr# Int# -> Int# -> Int#
+# Int#
b# Int# -> Int# -> Int#
-# Int# -> Int# -> Int#
uncheckedIShiftL# Int#
b# Int#
61#))
  {-# INLINE (*) #-}
  abs :: RollingHash b -> RollingHash b
abs = RollingHash b -> RollingHash b
forall a. a -> a
id
  {-# INLINE abs #-}
  signum :: RollingHash b -> RollingHash b
signum = RollingHash b -> RollingHash b -> RollingHash b
forall a b. a -> b -> a
const (Int -> RollingHash b
forall (b :: Nat). Int -> RollingHash b
RollingHash Int
1)
  {-# INLINE signum #-}
  fromInteger :: Integer -> RollingHash b
fromInteger Integer
x = Int -> RollingHash b
forall (b :: Nat). Int -> RollingHash b
RollingHash (Int -> RollingHash b) -> Int -> RollingHash b
forall a b. (a -> b) -> a -> b
$ Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
mod Integer
x Integer
0x1fffffffffffffff)
  {-# INLINE fromInteger #-}

newtype instance U.MVector s (RollingHash b) = MV_RollingHash (U.MVector s Int)
newtype instance U.Vector (RollingHash b) = V_RollingHash (U.Vector Int)
deriving newtype instance GM.MVector U.MVector (RollingHash b)
deriving newtype instance G.Vector U.Vector (RollingHash b)
instance U.Unbox (RollingHash b)

asRollingHashOf :: RollingHash b -> Proxy b -> RollingHash b
asRollingHashOf :: forall (b :: Nat). RollingHash b -> Proxy b -> RollingHash b
asRollingHashOf = RollingHash b -> Proxy b -> RollingHash b
forall a b. a -> b -> a
const

{- |
check if @g@ is a primitive root of @2^61-1@.

>>> isPrimitiveRootRH 2047
True
>>> take 10 $ filter isPrimitiveRootRH [2..]
[37,43,55,69,74,86,110,123,133,138]
-}
isPrimitiveRootRH :: Int -> Bool
isPrimitiveRootRH :: Int -> Bool
isPrimitiveRootRH Int
g = (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Int -> Bool
ok [Int
2, Int
3, Int
5, Int
7, Int
11, Int
13, Int
31, Int
41, Int
61, Int
151, Int
331, Int
1321]
  where
    ok :: Int -> Bool
    ok :: Int -> Bool
ok Int
p = Int -> RollingHash 2047
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
g RollingHash 2047 -> Int -> RollingHash 2047
forall a b. (Num a, Integral b) => a -> b -> a
^ Int -> Int -> Int
forall a. Integral a => a -> a -> a
quot Int
0x1ffffffffffffffe Int
p RollingHash 2047 -> RollingHash 2047 -> Bool
forall a. Eq a => a -> a -> Bool
/= (RollingHash 2047
1 :: RollingHash 2047)

{- |
Generate a primitive root of @2^61-1@.

>>> genPrimitiveRootRH (mkStdGen 123)
1600615663002506808
>>> isPrimitiveRootRH 1600615663002506808
True
-}
genPrimitiveRootRH :: (RandomGen g) => g -> Int
genPrimitiveRootRH :: forall g. RandomGen g => g -> Int
genPrimitiveRootRH = g -> Int
forall g. RandomGen g => g -> Int
go
  where
    go :: t -> Int
go !t
g
      | Int -> Bool
isPrimitiveRootRH Int
x = Int
x
      | Bool
otherwise = t -> Int
go t
g'
      where
        (Int
x, t
g') = (Int, Int) -> t -> (Int, t)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
forall g. RandomGen g => (Int, Int) -> g -> (Int, g)
randomR (Int
2, Int
0x1ffffffffffffffe) t
g

{- |
>>> withPrimitiveRootRH (mkStdGen 123) $ \proxy -> getRollingHash $ 456 `asRollingHashOf` proxy
456
-}
withPrimitiveRootRH ::
  (RandomGen g) =>
  g ->
  (forall b. (KnownNat b) => Proxy b -> a) ->
  a
withPrimitiveRootRH :: forall g a.
RandomGen g =>
g -> (forall (b :: Nat). KnownNat b => Proxy b -> a) -> a
withPrimitiveRootRH g
gen forall (b :: Nat). KnownNat b => Proxy b -> a
f =
  case Integer -> Maybe SomeNat
someNatVal (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ g -> Int
forall g. RandomGen g => g -> Int
genPrimitiveRootRH g
gen) of
    Just (SomeNat Proxy n
proxy) -> Proxy n -> a
forall (b :: Nat). KnownNat b => Proxy b -> a
f Proxy n
proxy
    Maybe SomeNat
Nothing -> String -> a
forall a. HasCallStack => String -> a
error String
"withPrimitiveRootRH failed"
{-# INLINE withPrimitiveRootRH #-}