Re-implement `nextPrimeInteger` predicate (#9281)

This also adds `nextPrimeWord#` and `nextPrimeBigNat` predicates.

`nextPrimeInteger` has been available since `integer-gmp-0.5.1`
(added via f49735486533842cc84df70cafc8d565dffd75db).

3 files changed, 102 insertions, 0 deletions

diff --git a/libraries/integer-gmp2/cbits/wrappers.c b/libraries/integer-gmp2/cbits/wrappers.c
index 6c188a3684..1621d3b3b5 100644
--- a/libraries/integer-gmp2/cbits/wrappers.c
+++ b/libraries/integer-gmp2/cbits/wrappers.c
@@ -474,3 +474,61 @@ integer_gmp_test_prime1(const mp_limb_t limb, const HsInt rep)
 
   return integer_gmp_test_prime(&limb, 1, rep);
 }
+
+/* wrapper around mpz_nextprime()
+ *
+ * Stores next prime (relative to {sp,sn}) in {rp,sn}.
+ * Return value is most significant limb of {rp,sn+1}.
+ */
+mp_limb_t
+integer_gmp_next_prime(mp_limb_t rp[], const mp_limb_t sp[],
+                       const mp_size_t sn)
+{
+  if (!sn) return 2;
+
+  const mpz_t op = {{
+      ._mp_alloc = sn,
+      ._mp_size  = sn,
+      ._mp_d = (mp_limb_t*)sp
+    }};
+
+  mpz_t rop;
+  mpz_init (rop);
+  mpz_nextprime (rop, op);
+
+  const mp_size_t rn = rop[0]._mp_size;
+
+  // copy result into {rp,sn} buffer
+  assert (rn == sn || rn == sn+1);
+  memcpy(rp, rop[0]._mp_d, sn*sizeof(mp_limb_t));
+  const mp_limb_t result = rn>sn ? rop[0]._mp_d[sn] : 0;
+
+  mpz_clear (rop);
+
+  return result;
+}
+
+/* wrapper around mpz_nextprime()
+ *
+ * returns next prime modulo 2^GMP_LIMB_BITS
+ */
+mp_limb_t
+integer_gmp_next_prime1(const mp_limb_t limb)
+{
+  if (limb < 2) return 2;
+
+  const mpz_t op = {{
+      ._mp_alloc = 1,
+      ._mp_size  = 1,
+      ._mp_d = (mp_limb_t*)(&limb)
+    }};
+
+  mpz_t rop;
+  mpz_init (rop);
+  mpz_nextprime (rop, op);
+  assert (rop[0]._mp_size > 0);
+  const mp_limb_t result = rop[0]._mp_d[0];
+  mpz_clear (rop);
+
+  return result;
+}
diff --git a/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs b/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
index 480866bec1..244dac7990 100644
--- a/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
+++ b/libraries/integer-gmp2/src/GHC/Integer/GMP/Internals.hs
@@ -125,6 +125,10 @@ module GHC.Integer.GMP.Internals
     , testPrimeBigNat
     , testPrimeWord#
 
+    , nextPrimeInteger
+    , nextPrimeBigNat
+    , nextPrimeWord#
+
       -- * Import/export functions
       -- ** Compute size of serialisation
     , sizeInBaseBigNat
@@ -323,3 +327,26 @@ foreign import ccall unsafe "integer_gmp_test_prime"
 -- /Since: 1.0.0.0/
 foreign import ccall unsafe "integer_gmp_test_prime1"
   testPrimeWord# :: GmpLimb# -> Int# -> Int#
+
+
+-- | Compute next prime greater than @/n/@ probalistically.
+--
+-- According to the GMP documentation, the underlying function
+-- @mpz_nextprime()@ \"uses a probabilistic algorithm to identify
+-- primes. For practical purposes it's adequate, the chance of a
+-- composite passing will be extremely small.\"
+--
+-- /Since: 0.5.1.0/
+{-# NOINLINE nextPrimeInteger #-}
+nextPrimeInteger :: Integer -> Integer
+nextPrimeInteger (S# i#)
+  | isTrue# (i# ># 1#)    = wordToInteger (nextPrimeWord# (int2Word# i#))
+  | True                  = S# 2#
+nextPrimeInteger (Jp# bn) = Jp# (nextPrimeBigNat bn)
+nextPrimeInteger (Jn# _)  = S# 2#
+
+-- | Version of 'nextPrimeInteger' operating on 'Word#'s
+--
+-- /Since: 1.0.0.0/
+foreign import ccall unsafe "integer_gmp_next_prime1"
+  nextPrimeWord# :: GmpLimb# -> GmpLimb#
diff --git a/libraries/integer-gmp2/src/GHC/Integer/Type.hs b/libraries/integer-gmp2/src/GHC/Integer/Type.hs
index 48c5ed85bc..d771de1cda 100644
--- a/libraries/integer-gmp2/src/GHC/Integer/Type.hs
+++ b/libraries/integer-gmp2/src/GHC/Integer/Type.hs
@@ -1684,6 +1684,23 @@ isValidBigNat# (BN# ba#)
 
     (# szq#, szr# #) = quotRemInt# sz# GMP_LIMB_BYTES#
 
+-- | Version of 'nextPrimeInteger' operating on 'BigNat's
+--
+-- /Since: 1.0.0.0/
+nextPrimeBigNat :: BigNat -> BigNat
+nextPrimeBigNat bn@(BN# ba#) = runS $ do
+    mbn@(MBN# mba#) <- newBigNat# n#
+    (W# c#) <- liftIO (nextPrime# mba# ba# n#)
+    case c# of
+        0## -> unsafeFreezeBigNat# mbn
+        _   -> unsafeSnocFreezeBigNat# mbn c#
+  where
+    n# = sizeofBigNat# bn
+
+foreign import ccall unsafe "integer_gmp_next_prime"
+  nextPrime# :: MutableByteArray# RealWorld -> ByteArray# -> GmpSize#
+                -> IO GmpLimb
+
 ----------------------------------------------------------------------------
 -- monadic combinators for low-level state threading