Re-implement `recipModInteger` (#9281)

This also exposes the following two type-specialised modular
exponentiation variants of `recipModInteger` useful for implementing a
`recipModNatural` operation.

  recipModBigNat :: BigNat -> BigNat -> BigNat
  recipModWord   :: Word#  -> Word#  -> Word#

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

4 files changed, 124 insertions, 10 deletions

diff --git a/libraries/integer-gmp2/cbits/wrappers.c b/libraries/integer-gmp2/cbits/wrappers.c
index 52920ec21e..3023816681 100644
--- a/libraries/integer-gmp2/cbits/wrappers.c
+++ b/libraries/integer-gmp2/cbits/wrappers.c
@@ -611,3 +611,75 @@ integer_gmp_powm_word(const mp_limb_t b0, // base
 {
   return integer_gmp_powm1(&b0, !!b0, &e0, !!e0, m0);
 }
+
+
+/* wrapper around mpz_invert()
+ *
+ * Store '(1/X) mod abs(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.
+ *
+ * Returns 0 if inverse does not exist.
+ */
+mp_size_t
+integer_gmp_invert(mp_limb_t rp[], // result
+                   const mp_limb_t xp[], const mp_size_t xn, // base
+                   const mp_limb_t mp[], const mp_size_t mn) // mod
+{
+  if (mp_limb_zero_p(xp,xn)
+      || mp_limb_zero_p(mp,mn)
+      || ((mn == 1 || mn == -1) && mp[0] == 1)) {
+    rp[0] = 0;
+    return 1;
+  }
+
+  const mpz_t x = CONST_MPZ_INIT(xp, xn);
+  const mpz_t m = CONST_MPZ_INIT(mp, mn);
+
+  mpz_t r;
+  mpz_init (r);
+
+  const int inv_exists = mpz_invert(r, x, m);
+
+  const mp_size_t rn = inv_exists ? r[0]._mp_size : 0;
+
+  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_invert() operating on single limbs */
+mp_limb_t
+integer_gmp_invert_word(const mp_limb_t x0, const mp_limb_t m0)
+{
+  if (!x0 || m0<=1) return 0;
+  if (x0 == 1) return 1;
+
+  const mpz_t x = CONST_MPZ_INIT(&x0, 1);
+  const mpz_t m = CONST_MPZ_INIT(&m0, 1);
+
+  mpz_t r;
+  mpz_init (r);
+
+  const int inv_exists = mpz_invert(r, x, m);
+  const mp_size_t rn = inv_exists ? r[0]._mp_size : 0;
+
+  assert (rn == 0 || rn == 1);
+  const mp_limb_t r0 = rn ? r[0]._mp_d[0] : 0;
+
+  mpz_clear (r);
+
+  return r0;
+}
diff --git a/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs b/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
index f7b9513427..9559755f61 100644
--- a/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
+++ b/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
@@ -47,6 +47,7 @@ module GHC.Integer.GMP.Internals
     , lcmInteger
     , sqrInteger
     , powModInteger
+    , recipModInteger
 
       -- ** Additional conversion operations to 'Integer'
     , wordToNegInteger
@@ -98,6 +99,8 @@ module GHC.Integer.GMP.Internals
     , powModBigNat
     , powModBigNatWord
 
+    , recipModBigNat
+
       -- ** 'BigNat' logic operations
     , shiftRBigNat
     , shiftLBigNat
@@ -124,6 +127,7 @@ module GHC.Integer.GMP.Internals
     , gcdInt
     , gcdWord
     , powModWord
+    , recipModWord
 
       -- * Primality tests
     , testPrimeInteger
diff --git a/libraries/integer-gmp2/src/GHC/Integer/Type.hs b/libraries/integer-gmp2/src/GHC/Integer/Type.hs
index 8fe1d15482..6284917910 100644
--- a/libraries/integer-gmp2/src/GHC/Integer/Type.hs
+++ b/libraries/integer-gmp2/src/GHC/Integer/Type.hs
@@ -1340,6 +1340,53 @@ foreign import ccall unsafe "integer_gmp_powm1"
   integer_gmp_powm1# :: ByteArray# -> GmpSize# -> ByteArray# -> GmpSize#
                         -> GmpLimb# -> GmpLimb#
 
+
+-- | \"@'recipModInteger' /x/ /m/@\" computes the inverse of @/x/@ modulo @/m/@. If
+-- the inverse exists, the return value @/y/@ will satisfy @0 < /y/ <
+-- abs(/m/)@, otherwise the result is @0@.
+--
+-- /Since: 0.5.1.0/
+{-# NOINLINE recipModInteger #-}
+recipModInteger :: Integer -> Integer -> Integer
+recipModInteger (S# x#) (S# m#)
+  | isTrue# (x# >=# 0#)
+  = wordToInteger (recipModWord (int2Word# x#) (int2Word# (absI# m#)))
+recipModInteger x m = bigNatToInteger (recipModSBigNat x' m')
+  where
+    x' = integerToSBigNat x
+    m' = absSBigNat (integerToSBigNat m)
+
+-- | Version of 'recipModInteger' operating on 'BigNat's
+--
+-- /Since: 1.0.0.0/
+recipModBigNat :: BigNat -> BigNat -> BigNat
+recipModBigNat x m = inline recipModSBigNat (PosBN x) m
+
+-- | Version of 'recipModInteger' operating on 'Word#'s
+--
+-- /Since: 1.0.0.0/
+foreign import ccall unsafe "integer_gmp_invert_word"
+  recipModWord :: GmpLimb# -> GmpLimb# -> GmpLimb#
+
+-- internal non-exported helper
+recipModSBigNat :: SBigNat -> BigNat -> BigNat
+recipModSBigNat x m@(BN# m#) = runS $ do
+    r@(MBN# r#) <- newBigNat# mn#
+    I# rn_# <- liftIO (integer_gmp_invert# r# x# xn# m# mn#)
+    let rn# = narrowGmpSize# rn_#
+    case rn# ==# mn# of
+        0# -> unsafeShrinkFreezeBigNat# r rn#
+        _  -> unsafeFreezeBigNat# r
+  where
+    !(BN# x#) = absSBigNat x
+    xn# = ssizeofSBigNat# x
+    mn# = sizeofBigNat# m
+
+foreign import ccall unsafe "integer_gmp_invert"
+  integer_gmp_invert# :: MutableByteArray# RealWorld
+                         -> ByteArray# -> GmpSize#
+                         -> ByteArray# -> GmpSize# -> IO GmpSize
+
 ----------------------------------------------------------------------------
 -- Conversions to/from floating point
 
diff --git a/testsuite/tests/lib/integer/integerGmpInternals.hs b/testsuite/tests/lib/integer/integerGmpInternals.hs
index d281b739d4..2f49a755a1 100644
--- a/testsuite/tests/lib/integer/integerGmpInternals.hs
+++ b/testsuite/tests/lib/integer/integerGmpInternals.hs
@@ -17,17 +17,8 @@ import qualified GHC.Integer.GMP.Internals as I
 -- so we use naive reference-implementations instead for the meantime
 -- in order to keep the reference-output untouched.
 
--- FIXME: Lacks GMP2 version
--- stolen from `arithmoi` package
 recipModInteger :: Integer -> Integer -> Integer
-recipModInteger k 0 = if k == 1 || k == (-1) then k else 0
-recipModInteger k m = case gcdExtInteger k' m' of
-                  (1, u) -> if u < 0 then m' + u else u
-                  _      -> 0
-  where
-    m' = abs m
-    k' | k >= m' || k < 0   = k `mod` m'
-       | otherwise          = k
+recipModInteger = I.recipModInteger
 
 -- FIXME: Lacks GMP2 version
 gcdExtInteger :: Integer -> Integer -> (Integer, Integer)