content: {-# LANGUAGE MagicHash #-}
  {-# LANGUAGE NamedFieldPuns #-}
  
  -- | Fast modular multiplication for `Word64` using Montgomery multiplication.
  --
  -- - The modulus value must be odd.
  -- - If the modulus value is known to fit in 32 bits, @AtCoder.Internal.Barrett@ can be faster.
  --
  -- @since 1.2.6.0
  module AtCoder.Extra.Math.Montgomery64
    ( -- * Montgomery64
      Montgomery64,
  
      -- * Constructor
      new,
      fromVal,
  
      -- * Accessor
      umod,
  
      -- * Montgomery form encoding
      encode,
      decode,
      reduce,
  
      -- * Calculations
      addMod,
      subMod,
      mulMod,
      powMod,
      eq,
    )
  where
  
  import AtCoder.Internal.Assert qualified as ACIA
  import Data.Bits (bit, (!>>.))
  import Data.WideWord.Word128 (Word128 (..))
  import Data.Word (Word64)
  import GHC.Exts (Proxy#)
  import GHC.Stack (HasCallStack)
  import GHC.TypeNats (KnownNat, natVal')
  
  -- TODO: provide with newtype for Montgomery form?
  
  -- | Fast modular multiplication for `Word64` using Montgomery64 multiplication.
  --
  -- @since 1.2.6.0
  data Montgomery64 = Montgomery64
    { mM64 :: {-# UNPACK #-} !Word64,
      rM64 :: {-# UNPACK #-} !Word64,
      n2M64 :: {-# UNPACK #-} !Word64
    }
    deriving
      ( -- | NOTE: The represented value can be equal even if the internal values are different, so
        -- use `eq` for comparison.
        --
        -- @since 1.2.6.0
        Eq,
        -- | @since 1.2.6.0
        Show
      )
  
  -- | \(O(1)\) Static, shared storage of `Montgomery64`.
  --
  -- ==== Constraints
  -- - \(m \le 2^{62}\)
  -- - \(m\) is odd
  --
  -- @since 1.2.6.0
  {-# INLINE new #-}
  new :: forall a. (HasCallStack, KnownNat a) => Proxy# a -> Montgomery64
  new p = fromVal . fromIntegral $! natVal' p
  
  -- | \(O(1)\) Creates a `Montgomery64` for a modulus value \(m\) of type `Word64` value.
  --
  -- ==== Constraints
  -- - \(m \le 2^{62}\)
  -- - \(m\) is odd
  --
  -- @since 1.2.6.0
  {-# INLINE fromVal #-}
  fromVal :: (HasCallStack) => Word64 -> Montgomery64
  fromVal m =
    let !m128 :: Word128 = fromIntegral m
        !n2 = word128Lo64 $ (-m128) `mod` m128
        !r = getR m 0
        !_ = ACIA.runtimeAssert (r * m == maxBound) "AtCoder.Extra.Montgomery64.fromVal: internal implementation error"
     in Montgomery64 m r n2
    where
      !_ = ACIA.runtimeAssert (odd m && m <= bit 62) $ "AtCoder.Extra.Montgomery64.fromVal: not given odd modulus value that is less than or equal to 2^62: " ++ show m
      getR :: Word64 -> Int -> Word64
      getR !acc i
        | i >= 5 = -acc
        | otherwise = getR (acc * (2 - m * acc)) (i + 1)
  
  -- | \(O(1)\) Retrieves the modulus \(m\).
  --
  -- @since 1.2.6.0
  {-# INLINE umod #-}
  umod :: Montgomery64 -> Word64
  umod Montgomery64 {mM64} = mM64
  
  -- | \(O(1)\) Converts the given `Word64` to Montgomery form.
  --
  -- @since 1.2.6.0
  {-# INLINE encode #-}
  encode :: Montgomery64 -> Word64 -> Word64
  encode mont@Montgomery64 {n2M64} x = reduce mont $! fromIntegral x * fromIntegral n2M64
  
  -- | \(O(1)\) Retrieves the value from a Montgomery form of value.
  --
  -- @since 1.2.6.0
  {-# INLINE decode #-}
  decode :: Montgomery64 -> Word64 -> Word64
  decode mont@Montgomery64 {mM64} x =
    let !res = reduce mont $! fromIntegral x
     in if res >= mM64 then res - mM64 else res
  
  -- | \(O(1)\) Takes the mod in Montgomery form.
  --
  -- @since 1.2.6.0
  {-# INLINE reduce #-}
  reduce :: Montgomery64 -> Word128 -> Word64
  reduce Montgomery64 {mM64, rM64} x =
    word128Hi64 $!
      (x + fromIntegral (word128Lo64 x * rM64) * fromIntegral mM64)
  
  -- | \(O(1)\) Calculates \(a + b \bmod m\) in the Montgomery form.
  {-# INLINE addMod #-}
  addMod :: Word64 -> Word64 -> Word64 -> Word64
  addMod m a b
    | x' >= m = x' - m
    | otherwise = x'
    where
      !x' = a + b
  
  -- | \(O(1)\) Calculates \(a - b \bmod m\) in the Montgomery form.
  {-# INLINE subMod #-}
  subMod :: Word64 -> Word64 -> Word64 -> Word64
  subMod m a b
    | a >= b = a - b
    | otherwise = a - b + m
  
  -- | \(O(1)\) Calculates \(a^n \bmod m\) in the Montgomery form.
  --
  -- @since 1.2.6.0
  {-# INLINE mulMod #-}
  mulMod :: Montgomery64 -> Word64 -> Word64 -> Word64
  mulMod mont a b = reduce mont $! fromIntegral a * fromIntegral b
  
  -- | \(O(w)\) Calculates \(a^n \bmod m\) in the Montgomery form.
  --
  -- @since 1.2.6.0
  {-# INLINE powMod #-}
  powMod :: (HasCallStack) => Montgomery64 -> Word64 -> Int -> Word64
  powMod mont x0 n0 = inner n0 (encode mont 1) x0
    where
      !_ = ACIA.runtimeAssert (0 <= n0) $ "AtCoder.Extra.Math.Montgomery64.powMod: given negative exponential `n`: " ++ show n0 ++ show "`"
      inner :: Int -> Word64 -> Word64 -> Word64
      inner !n !r !y
        | n == 0 = r
        | otherwise =
            let !r' = if odd n then mulMod mont r y else r
                !y' = mulMod mont y y
             in inner (n !>>. 1) r' y'
  
  -- | \(O(1)\) Compares two values of Montgomery form and returns whether they represent the same
  -- value.
  --
  -- @since 1.2.6.0
  {-# INLINE eq #-}
  eq :: Word64 -> Word64 -> Word64 -> Bool
  eq mM64 a b = a' == b'
    where
      !a' = if a < mM64 then a else a - mM64
      !b' = if b < mM64 then b else b - mM64
  