Re-implement `powModInteger` (#9281)

This also exposes the following type-specialised modular exponentiation
variants of `powModInteger` useful for implementing a `powModNatural`
operation.

  powModBigNat     :: BigNat -> BigNat -> BigNat -> BigNat
  powModBigNatWord :: BigNat -> BigNat -> Word#  -> Word#
  powModWord       :: Word#  -> Word#  -> Word#  -> Word#

`powModInteger` has been available since `integer-gmp-0.5.1`
(added via 4d516855241b70eb687d95e3c121428de885e83e)

5 files changed, 215 insertions, 14 deletions

diff --git a/libraries/integer-gmp2/cbits/wrappers.c b/libraries/integer-gmp2/cbits/wrappers.c
index a1a78e0fb1..52920ec21e 100644
--- a/libraries/integer-gmp2/cbits/wrappers.c
+++ b/libraries/integer-gmp2/cbits/wrappers.c
@@ -522,3 +522,92 @@ integer_gmp_next_prime1(const mp_limb_t limb)
 
   return result;
 }
+
+/* wrapper around mpz_powm()
+ *
+ * Store '(B^E) mod M' in {rp,rn}
+ *
+ * rp must have allocated mn limbs; This function's return value is
+ * the actual number rn (0 < rn <= mn) of limbs written to the rp limb-array.
+ *
+ * bn and en are allowed to be negative to denote negative numbers
+ */
+mp_size_t
+integer_gmp_powm(mp_limb_t rp[], // result
+                 const mp_limb_t bp[], const mp_size_t bn, // base
+                 const mp_limb_t ep[], const mp_size_t en, // exponent
+                 const mp_limb_t mp[], const mp_size_t mn) // mod
+{
+  assert(!mp_limb_zero_p(mp,mn));
+
+  if ((mn == 1 || mn == -1) && mp[0] == 1) {
+    rp[0] = 0;
+    return 1;
+  }
+
+  if (mp_limb_zero_p(ep,en)) {
+    rp[0] = 1;
+    return 1;
+  }
+
+  const mpz_t b = CONST_MPZ_INIT(bp, mp_limb_zero_p(bp,bn) ? 0 : bn);
+  const mpz_t e = CONST_MPZ_INIT(ep, mp_limb_zero_p(ep,en) ? 0 : en);
+  const mpz_t m = CONST_MPZ_INIT(mp, mn);
+
+  mpz_t r;
+  mpz_init (r);
+
+  mpz_powm(r, b, e, m);
+
+  const mp_size_t rn = r[0]._mp_size;
+
+  if (rn) {
+    assert(0 < rn && rn <= mn);
+    memcpy(rp, r[0]._mp_d, rn*sizeof(mp_limb_t));
+  }
+
+  mpz_clear (r);
+
+  if (!rn) {
+    rp[0] = 0;
+    return 1;
+  }
+
+  return rn;
+}
+
+/* version of integer_gmp_powm() for single-limb moduli */
+mp_limb_t
+integer_gmp_powm1(const mp_limb_t bp[], const mp_size_t bn, // base
+                  const mp_limb_t ep[], const mp_size_t en, // exponent
+                  const mp_limb_t m0) // mod
+{
+  assert(m0);
+
+  if (m0==1) return 0;
+  if (mp_limb_zero_p(ep,en)) return 1;
+
+  const mpz_t b = CONST_MPZ_INIT(bp, mp_limb_zero_p(bp,bn) ? 0 : bn);
+  const mpz_t e = CONST_MPZ_INIT(ep, en);
+  const mpz_t m = CONST_MPZ_INIT(&m0, !!m0);
+
+  mpz_t r;
+  mpz_init (r);
+  mpz_powm(r, b, e, m);
+
+  assert(r[0]._mp_size == 0 || r[0]._mp_size == 1);
+  const mp_limb_t result = r[0]._mp_size ? r[0]._mp_d[0] : 0;
+
+  mpz_clear (r);
+
+  return result;
+}
+
+/* version of integer_gmp_powm() for single-limb arguments */
+mp_limb_t
+integer_gmp_powm_word(const mp_limb_t b0, // base
+                      const mp_limb_t e0, // exponent
+                      const mp_limb_t m0) // mod
+{
+  return integer_gmp_powm1(&b0, !!b0, &e0, !!e0, m0);
+}
diff --git a/libraries/integer-gmp2/include/HsIntegerGmp.h.in b/libraries/integer-gmp2/include/HsIntegerGmp.h.in
index 0637ba3b6a..ba0767cae7 100644
--- a/libraries/integer-gmp2/include/HsIntegerGmp.h.in
+++ b/libraries/integer-gmp2/include/HsIntegerGmp.h.in
@@ -1,8 +1,6 @@
 #ifndef _HS_INTEGER_GMP_H_
 #define _HS_INTEGER_GMP_H_
 
-#define HAVE_SECURE_POWM @HaveSecurePowm@
-
 /* Whether GMP is embedded into integer-gmp */
 #define GHC_GMP_INTREE     @UseIntreeGmp@
 
diff --git a/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs b/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
index 244dac7990..f7b9513427 100644
--- a/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
+++ b/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
@@ -46,6 +46,7 @@ module GHC.Integer.GMP.Internals
     , gcdInteger
     , lcmInteger
     , sqrInteger
+    , powModInteger
 
       -- ** Additional conversion operations to 'Integer'
     , wordToNegInteger
@@ -94,6 +95,9 @@ module GHC.Integer.GMP.Internals
     , gcdBigNat
     , gcdBigNatWord
 
+    , powModBigNat
+    , powModBigNatWord
+
       -- ** 'BigNat' logic operations
     , shiftRBigNat
     , shiftLBigNat
@@ -119,6 +123,7 @@ module GHC.Integer.GMP.Internals
       -- * Miscellaneous GMP-provided operations
     , gcdInt
     , gcdWord
+    , powModWord
 
       -- * Primality tests
     , testPrimeInteger
diff --git a/libraries/integer-gmp2/src/GHC/Integer/Type.hs b/libraries/integer-gmp2/src/GHC/Integer/Type.hs
index d771de1cda..8fe1d15482 100644
--- a/libraries/integer-gmp2/src/GHC/Integer/Type.hs
+++ b/libraries/integer-gmp2/src/GHC/Integer/Type.hs
@@ -1256,6 +1256,89 @@ gcdBigNat x@(BN# x#) y@(BN# y#)
     nx# = sizeofBigNat# x
     ny# = sizeofBigNat# y
 
+----------------------------------------------------------------------------
+-- modular exponentiation
+
+-- | \"@'powModInteger' /b/ /e/ /m/@\" computes base @/b/@ raised to
+-- exponent @/e/@ modulo @abs(/m/)@.
+--
+-- Negative exponents are supported if an inverse modulo @/m/@
+-- exists.
+--
+-- __Warning__: It's advised to avoid calling this primitive with
+-- negative exponents unless it is guaranteed the inverse exists, as
+-- failure to do so will likely cause program abortion due to a
+-- divide-by-zero fault. See also 'recipModInteger'.
+--
+-- Future versions of @integer_gmp@ may not support negative @/e/@
+-- values anymore.
+--
+-- /Since: 0.5.1.0/
+{-# NOINLINE powModInteger #-}
+powModInteger :: Integer -> Integer -> Integer -> Integer
+powModInteger (S# b#) (S# e#) (S# m#)
+  | isTrue# (b# >=# 0#), isTrue# (e# >=# 0#)
+  = wordToInteger (powModWord (int2Word# b#) (int2Word# e#)
+                              (int2Word# (absI# m#)))
+powModInteger b e m = case m of
+    (S# m#) -> wordToInteger (powModSBigNatWord b' e' (int2Word# (absI# m#)))
+    (Jp# m') -> bigNatToInteger (powModSBigNat b' e' m')
+    (Jn# m') -> bigNatToInteger (powModSBigNat b' e' m')
+  where
+    b' = integerToSBigNat b
+    e' = integerToSBigNat e
+
+-- | Version of 'powModInteger' operating on 'BigNat's
+--
+-- /Since: 1.0.0.0/
+powModBigNat :: BigNat -> BigNat -> BigNat -> BigNat
+powModBigNat b e m = inline powModSBigNat (PosBN b) (PosBN e) m
+
+-- | Version of 'powModInteger' for 'Word#'-sized moduli
+--
+-- /Since: 1.0.0.0/
+powModBigNatWord :: BigNat -> BigNat -> GmpLimb# -> GmpLimb#
+powModBigNatWord b e m# = inline powModSBigNatWord (PosBN b) (PosBN e) m#
+
+-- | Version of 'powModInteger' operating on 'Word#'s
+--
+-- /Since: 1.0.0.0/
+foreign import ccall unsafe "integer_gmp_powm_word"
+  powModWord :: GmpLimb# -> GmpLimb# -> GmpLimb# -> GmpLimb#
+
+-- internal non-exported helper
+powModSBigNat :: SBigNat -> SBigNat -> BigNat -> BigNat
+powModSBigNat b e m@(BN# m#) = runS $ do
+    r@(MBN# r#) <- newBigNat# mn#
+    I# rn_# <- liftIO (integer_gmp_powm# r# b# bn# e# en# m# mn#)
+    let rn# = narrowGmpSize# rn_#
+    case rn# ==# mn# of
+        0# -> unsafeShrinkFreezeBigNat# r rn#
+        _  -> unsafeFreezeBigNat# r
+  where
+    !(BN# b#) = absSBigNat b
+    !(BN# e#) = absSBigNat e
+    bn# = ssizeofSBigNat# b
+    en# = ssizeofSBigNat# e
+    mn# = sizeofBigNat# m
+
+foreign import ccall unsafe "integer_gmp_powm"
+  integer_gmp_powm# :: MutableByteArray# RealWorld
+                       -> ByteArray# -> GmpSize# -> ByteArray# -> GmpSize#
+                       -> ByteArray# -> GmpSize# -> IO GmpSize
+
+-- internal non-exported helper
+powModSBigNatWord :: SBigNat -> SBigNat -> GmpLimb# -> GmpLimb#
+powModSBigNatWord b e m# = integer_gmp_powm1# b# bn# e# en# m#
+  where
+    !(BN# b#) = absSBigNat b
+    !(BN# e#) = absSBigNat e
+    bn# = ssizeofSBigNat# b
+    en# = ssizeofSBigNat# e
+
+foreign import ccall unsafe "integer_gmp_powm1"
+  integer_gmp_powm1# :: ByteArray# -> GmpSize# -> ByteArray# -> GmpSize#
+                        -> GmpLimb# -> GmpLimb#
 
 ----------------------------------------------------------------------------
 -- Conversions to/from floating point
@@ -1744,6 +1827,43 @@ fail :: [Char] -> S s a
 fail s = return (raise# s)
 
 ----------------------------------------------------------------------------
+
+-- | Internal helper type for "signed" 'BigNat's
+--
+-- This is a useful abstraction for operations which support negative
+-- mp_size_t arguments.
+data SBigNat = NegBN !BigNat | PosBN !BigNat
+
+-- | Absolute value of 'SBigNat'
+absSBigNat :: SBigNat -> BigNat
+absSBigNat (NegBN bn) = bn
+absSBigNat (PosBN bn) = bn
+
+-- | /Signed/ limb count. Negative sizes denote negative integers
+ssizeofSBigNat# :: SBigNat -> GmpSize#
+ssizeofSBigNat# (NegBN bn) = negateInt# (sizeofBigNat# bn)
+ssizeofSBigNat# (PosBN bn) = sizeofBigNat# bn
+
+-- | Construct 'SBigNat' from 'Int#' value
+intToSBigNat# :: Int# -> SBigNat
+intToSBigNat# 0#     = PosBN zeroBigNat
+intToSBigNat# 1#     = PosBN oneBigNat
+intToSBigNat# (-1#)  = NegBN oneBigNat
+intToSBigNat# i# | isTrue# (i# ># 0#) = PosBN (wordToBigNat (int2Word# i#))
+                 | True   = PosBN (wordToBigNat (int2Word# (negateInt# i#)))
+
+-- | Convert 'Integer' into 'SBigNat'
+integerToSBigNat :: Integer -> SBigNat
+integerToSBigNat (S#  i#) = intToSBigNat# i#
+integerToSBigNat (Jp# bn) = PosBN bn
+integerToSBigNat (Jn# bn) = NegBN bn
+
+-- | Convert 'SBigNat' into 'Integer'
+sBigNatToInteger :: SBigNat -> Integer
+sBigNatToInteger (NegBN bn) = bigNatToNegInteger bn
+sBigNatToInteger (PosBN bn) = bigNatToInteger bn
+
+----------------------------------------------------------------------------
 -- misc helpers, some of these should rather be primitives exported by ghc-prim
 
 cmpW# :: Word# -> Word# -> Ordering
diff --git a/testsuite/tests/lib/integer/integerGmpInternals.hs b/testsuite/tests/lib/integer/integerGmpInternals.hs
index 5db0b099ae..d281b739d4 100644
--- a/testsuite/tests/lib/integer/integerGmpInternals.hs
+++ b/testsuite/tests/lib/integer/integerGmpInternals.hs
@@ -47,19 +47,8 @@ gcdExtInteger a b = (d, u) -- stolen from `arithmoi` package
 powModSecInteger :: Integer -> Integer -> Integer -> Integer
 powModSecInteger = powModInteger
 
--- FIXME: Lacks GMP2 version
 powModInteger :: Integer -> Integer -> Integer -> Integer
-powModInteger b0 e0 m
-  | e0 >= 0    = go b0 e0 1
-  | otherwise  = error "non-neg exponent required"
-  where
-    go !b e !r
-      | odd e     = go b' e' (r*b `mod` m)
-      | e == 0    = r
-      | otherwise = go b' e' r
-      where
-        b' = b*b `mod` m
-        e' = e   `unsafeShiftR` 1 -- slightly faster than "e `div` 2"
+powModInteger = I.powModInteger
 
 -- FIXME: Lacks GMP2 version
 powInteger :: Integer -> Word -> Integer