content: {-# LANGUAGE NamedFieldPuns #-}
  
  module BenchLib.MulMod.Montgomery
    ( Montgomery (mMontgomery),
      new,
      umod,
      mulMod,
      mulModGenerated,
      generate,
      reduce,
    )
  where
  
  -- FIXME: test
  -- TODO: Use MagicHash?
  
  import Data.Bits (bit, (!>>.))
  import Data.Word (Word32, Word64)
  
  -- | Fast modular multiplication by Montgomery multiplication. The modulus value must be odd for
  -- the speed.
  data Montgomery = Montgomery
    { mMontgomery :: {-# UNPACK #-} !Word64,
      -- | R2 == (2^64) % MOD;
      r2Montgomery :: {-# UNPACK #-} !Word64,
      -- | MOD * NEG_INV % (2^32) == (2^32) - 1;
      negInvMontgomery :: {-# UNPACK #-} !Word32
    }
  
  -- | Creates a `Montgomery` for modulus @m@.
  new :: Word64 -> Montgomery
  new m =
    let !negInv = inner 0 0 1 0
        !r2 = bit 32 `rem` m
        !r2' = r2 * r2 `rem` m
     in Montgomery m r2' $ fromIntegral negInv
    where
      inner :: Word64 -> Word64 -> Word64 -> Word64 -> Word64
      -- TODO: more efficient impplementation
      inner !negInv 32 !_ !_ = negInv
      inner !negInv !i !s !t
        -- REMARK: `!>>.` does not work here.
        -- \| even t = inner (negInv + s) (i + 1) (2 * s) ((t + m) !>>. 2)
        -- \| otherwise = inner negInv (i + 1) (2 * s) (t !>>. 2)
        | even t = inner (negInv + s) (i + 1) (2 * s) ((t + m) `div` 2)
        | otherwise = inner negInv (i + 1) (2 * s) (t `div` 2)
  
  -- | Retrieves the modulus \(m\).
  umod :: Montgomery -> Word64
  umod = mMontgomery
  
  -- | Calculates \(a \cdot b \bmod m\).
  mulMod :: Montgomery -> Word64 -> Word64 -> Word64
  mulMod m a b = reduce m $ generate m a * generate m b
  
  -- | Calculates \(a \cdot b \bmod m\).
  mulModGenerated :: Montgomery -> Word64 -> Word64 -> Word64
  mulModGenerated m a b = reduce m $ a * b
  
  -- | Retrieves \( a \cdot b \bbmod m from a \cdot b\).
  generate :: Montgomery -> Word64 -> Word64
  generate m x = reduce m $ x * r2Montgomery m
  
  -- | Retrieves \( a \cdot b \bbmod m from a \cdot b\). Prefer `reduce32`.
  reduce :: Montgomery -> Word64 -> Word64
  reduce Montgomery {mMontgomery, negInvMontgomery} x =
    let !x' = (x + fromIntegral (as32 x * negInvMontgomery) * mMontgomery) !>>. 32
     in if x' < mMontgomery then x' else x' - mMontgomery
    where
      -- TODO: makes sense??
      as32 :: Word64 -> Word32
      as32 = fromIntegral
  