Bignum: implement integerRecipMod (#18427)

14 files changed, 160 insertions, 13 deletions

diff --git a/libraries/ghc-bignum/cbits/gmp_wrappers.c b/libraries/ghc-bignum/cbits/gmp_wrappers.c
index 56075fd1f7..e036f05a9a 100644
--- a/libraries/ghc-bignum/cbits/gmp_wrappers.c
+++ b/libraries/ghc-bignum/cbits/gmp_wrappers.c
@@ -779,8 +779,7 @@ integer_gmp_invert(mp_limb_t rp[], // result
   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;
+    return 0;
   }
 
   const mpz_t x = CONST_MPZ_INIT(xp, xn);
@@ -800,11 +799,6 @@ integer_gmp_invert(mp_limb_t rp[], // result
 
   mpz_clear (r);
 
-  if (!rn) {
-    rp[0] = 0;
-    return 1;
-  }
-
   return rn;
 }
 
diff --git a/libraries/ghc-bignum/src/GHC/Num/Backend/Check.hs b/libraries/ghc-bignum/src/GHC/Num/Backend/Check.hs
index a8ee366312..9e5181c0bb 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Backend/Check.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Backend/Check.hs
@@ -469,3 +469,18 @@ integer_gcde a b =
          then (# g0, x0, y0 #)
          else case unexpectedValue of
             !_ -> (# integerZero, integerZero, integerZero #)
+
+integer_recip_mod
+   :: Integer
+   -> Integer
+   -> (# Integer | () #)
+integer_recip_mod x m =
+   let
+      !r0 = Other.integer_recip_mod x m
+      !r1 = Native.integer_recip_mod x m
+   in case (# r0, r1 #) of
+         (# (# | () #), (# | () #) #) -> r0
+         (# (# y0 | #), (# y1 | #) #)
+            | isTrue# (integerEq# y0 y1) -> r0
+         _ -> case unexpectedValue of
+            !_ -> (# | () #)
diff --git a/libraries/ghc-bignum/src/GHC/Num/Backend/FFI.hs b/libraries/ghc-bignum/src/GHC/Num/Backend/FFI.hs
index 3fddd19098..98350d4d29 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Backend/FFI.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Backend/FFI.hs
@@ -596,3 +596,19 @@ integer_gcde = Native.integer_gcde
    -- for now we use Native's implementation. If some FFI backend user needs a
    -- specific implementation, we'll need to determine a prototype to pass and
    -- return BigNat signs and sizes via FFI.
+
+
+-- | Computes the modular inverse of two non-zero integers.
+--
+-- I.e. y = integer_recip_mod x m
+--        = x^(-1) `mod` m
+--
+-- with 0 < y < abs m
+integer_recip_mod
+   :: Integer
+   -> Integer
+   -> (# Integer | () #)
+integer_recip_mod = Native.integer_recip_mod
+   -- for now we use Native's implementation. If some FFI backend user needs a
+   -- specific implementation, we'll need to determine a prototype to pass and
+   -- return BigNat signs and sizes via FFI.
diff --git a/libraries/ghc-bignum/src/GHC/Num/Backend/GMP.hs b/libraries/ghc-bignum/src/GHC/Num/Backend/GMP.hs
index 10242dfdc7..462e67445c 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Backend/GMP.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Backend/GMP.hs
@@ -423,6 +423,32 @@ integer_gcde a b = case runRW# io of (# _, a #) -> a
 
 
 
+integer_recip_mod
+   :: Integer
+   -> Integer
+   -> (# Integer | () #)
+integer_recip_mod x m =
+   let
+      !(#  sign_x, bx #) = integerToBigNatSign# x
+      !(# _sign_m, bm #) = integerToBigNatSign# m
+      !br = sbignat_recip_mod sign_x bx bm
+   in if isTrue# (bigNatIsZero# br)
+         then (# | () #)
+         else (# integerFromBigNat# br | #)
+
+
+-- | Return 0 for invalid inputs
+sbignat_recip_mod :: Int# -> BigNat# -> BigNat# -> BigNat#
+sbignat_recip_mod sign_x x m = withNewWordArray# szm io
+  where
+    io r s = case ioInt# (integer_gmp_invert# r x ssx m szm) s of
+               (# s, rn #) -> mwaSetSize# r (narrowGmpSize# rn) s
+    !szx  = bigNatSize# x
+    !szm  = bigNatSize# m
+    !ssx = case sign_x of -- signed size of x
+            0# -> szx
+            _  -> negateInt# szx
+
 ----------------------------------------------------------------------
 -- FFI ccall imports
 
@@ -444,6 +470,11 @@ foreign import ccall unsafe "integer_gmp_gcdext" integer_gmp_gcdext#
   -> ByteArray# -> GmpSize#
   -> IO ()
 
+foreign import ccall unsafe "integer_gmp_invert"
+  integer_gmp_invert# :: MutableByteArray# RealWorld
+                         -> ByteArray# -> GmpSize#
+                         -> ByteArray# -> GmpSize# -> IO GmpSize
+
 -- mp_limb_t mpn_add_1 (mp_limb_t *rp, const mp_limb_t *s1p, mp_size_t n,
 --                      mp_limb_t s2limb)
 foreign import ccall unsafe "gmp.h __gmpn_add_1"
diff --git a/libraries/ghc-bignum/src/GHC/Num/Backend/Native.hs b/libraries/ghc-bignum/src/GHC/Num/Backend/Native.hs
index db4196f51f..c5af604c17 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Backend/Native.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Backend/Native.hs
@@ -737,3 +737,17 @@ integer_gcde a b = f (# a,integerOne,integerZero #) (# b,integerZero,integerOne
       | True            = case integerQuotRem# old_g g of
                               !(# q, r #) -> f new (# r , old_s `integerSub` (q `integerMul` s)
                                                         , old_t `integerSub` (q `integerMul` t) #)
+
+integer_recip_mod
+   :: Integer
+   -> Integer
+   -> (# Integer | () #)
+integer_recip_mod x m =
+   let m' = integerAbs m
+   in case integer_gcde x m' of
+      (# g, a, _b #)
+         -- gcd(x,m) = ax+mb = 1
+         -- ==> ax - 1 = -mb
+         -- ==> ax     = 1 [m]
+         | g `integerEq` integerOne -> (# a `integerMod` m' | #)
+         | True                     -> (# | () #)
diff --git a/libraries/ghc-bignum/src/GHC/Num/BigNat.hs-boot b/libraries/ghc-bignum/src/GHC/Num/BigNat.hs-boot
index 37edb97582..a62c07406d 100644
--- a/libraries/ghc-bignum/src/GHC/Num/BigNat.hs-boot
+++ b/libraries/ghc-bignum/src/GHC/Num/BigNat.hs-boot
@@ -5,11 +5,13 @@
 module GHC.Num.BigNat where
 
 import GHC.Num.WordArray
+import GHC.Num.Primitives
 import GHC.Prim
 
 type BigNat# = WordArray#
 data BigNat = BN# { unBigNat :: BigNat# }
 
+bigNatIsZero# :: BigNat# -> Bool#
 bigNatSize# :: BigNat# -> Int#
 bigNatSubUnsafe :: BigNat# -> BigNat# -> BigNat#
 bigNatMulWord# :: BigNat# -> Word# -> BigNat#
diff --git a/libraries/ghc-bignum/src/GHC/Num/Integer.hs b/libraries/ghc-bignum/src/GHC/Num/Integer.hs
index daa7ab8388..7f0314f859 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Integer.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Integer.hs
@@ -244,6 +244,11 @@ integerIsZero :: Integer -> Bool
 integerIsZero (IS 0#) = True
 integerIsZero _       = False
 
+-- | One predicate
+integerIsOne :: Integer -> Bool
+integerIsOne (IS 1#) = True
+integerIsOne _       = False
+
 -- | Not-equal predicate.
 integerNe :: Integer -> Integer -> Bool
 integerNe !x !y = isTrue# (integerNe# x y)
@@ -1217,3 +1222,22 @@ integerGcde
    -> ( Integer, Integer, Integer)
 integerGcde a b = case integerGcde# a b of
    (# g,x,y #) -> (g,x,y)
+
+
+-- | Computes the modular inverse.
+--
+-- I.e. y = integerRecipMod# x m
+--        = x^(-1) `mod` m
+--
+-- with 0 < y < |m|
+--
+integerRecipMod#
+   :: Integer
+   -> Integer
+   -> (# Integer | () #)
+integerRecipMod# x m
+   | integerIsZero x = (# | () #)
+   | integerIsZero m = (# | () #)
+   | IS 1#    <- m   = (# | () #)
+   | IS (-1#) <- m   = (# | () #)
+   | True            = Backend.integer_recip_mod x m
diff --git a/libraries/ghc-bignum/src/GHC/Num/Integer.hs-boot b/libraries/ghc-bignum/src/GHC/Num/Integer.hs-boot
index a07b9b7548..e9c2d003ca 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Integer.hs-boot
+++ b/libraries/ghc-bignum/src/GHC/Num/Integer.hs-boot
@@ -17,11 +17,15 @@ integerEq# :: Integer -> Integer -> Int#
 integerEq :: Integer -> Integer -> Bool
 integerGt :: Integer -> Integer -> Bool
 integerIsZero :: Integer -> Bool
+integerIsOne :: Integer -> Bool
 integerIsNegative :: Integer -> Bool
 
 integerSub :: Integer -> Integer -> Integer
 integerMul :: Integer -> Integer -> Integer
+integerMod :: Integer -> Integer -> Integer
+integerRem :: Integer -> Integer -> Integer
 integerNegate :: Integer -> Integer
+integerAbs :: Integer -> Integer
 integerDivMod# :: Integer -> Integer -> (# Integer, Integer #)
 integerQuotRem# :: Integer -> Integer -> (# Integer, Integer #)
 
diff --git a/libraries/integer-gmp/src/GHC/Integer/GMP/Internals.hs b/libraries/integer-gmp/src/GHC/Integer/GMP/Internals.hs
index a501196e67..cc1068ba10 100644
--- a/libraries/integer-gmp/src/GHC/Integer/GMP/Internals.hs
+++ b/libraries/integer-gmp/src/GHC/Integer/GMP/Internals.hs
@@ -32,6 +32,7 @@ module GHC.Integer.GMP.Internals
     , gcdExtInteger
     , lcmInteger
     , sqrInteger
+    , recipModInteger
 
       -- ** Additional conversion operations to 'Integer'
     , wordToNegInteger
@@ -186,6 +187,12 @@ lcmInteger = I.integerLcm
 sqrInteger :: Integer -> Integer
 sqrInteger = I.integerSqr
 
+{-# DEPRECATED recipModInteger "Use integerRecipMod# instead" #-}
+recipModInteger :: Integer -> Integer -> Integer
+recipModInteger x m = case I.integerRecipMod# x m of
+   (# y |    #) -> y
+   (#   | () #) -> 0
+
 {-# DEPRECATED wordToNegInteger "Use integerFromWordNeg# instead" #-}
 wordToNegInteger :: Word# -> Integer
 wordToNegInteger = I.integerFromWordNeg#
diff --git a/testsuite/tests/lib/integer/all.T b/testsuite/tests/lib/integer/all.T
index 0c997f7ebf..43bcb0f3e4 100644
--- a/testsuite/tests/lib/integer/all.T
+++ b/testsuite/tests/lib/integer/all.T
@@ -8,7 +8,8 @@ test('IntegerConversionRules', [], makefile_test, ['IntegerConversionRules'])
 test('gcdInteger', normal, compile_and_run, [''])
 test('gcdeInteger', normal, compile_and_run, [''])
 test('integerPowMod', [], compile_and_run, [''])
-test('integerGcdExt', [omit_ways(['ghci'])], compile_and_run, [''])
+test('integerGcdExt', [], compile_and_run, [''])
+test('integerRecipMod', [], compile_and_run, [''])
 
 # skip ghci as it doesn't support unboxed tuples
 test('integerImportExport', [omit_ways(['ghci'])], compile_and_run, [''])
diff --git a/testsuite/tests/lib/integer/integerGmpInternals.hs b/testsuite/tests/lib/integer/integerGmpInternals.hs
index 53a776e13d..982e4dcaba 100644
--- a/testsuite/tests/lib/integer/integerGmpInternals.hs
+++ b/testsuite/tests/lib/integer/integerGmpInternals.hs
@@ -12,9 +12,6 @@ import GHC.Base
 import GHC.Num.Integer
 import qualified GHC.Num.Integer as I
 
-recipModInteger :: Integer -> Integer -> Integer
-recipModInteger = I.recipModInteger
-
 -- FIXME: Lacks GMP2 version
 powInteger :: Integer -> Word -> Integer
 powInteger x e = x^e
@@ -25,7 +22,6 @@ main = do
     print $ powInteger 12345 0
     print $ powInteger 12345 1
     print $ powInteger 12345 30
-    print $ [ (x,i) | x <- [-7..71], let i = recipModInteger x (2*3*11*11*17*17), i /= 0 ]
 
     putStrLn "\n# nextPrimeInteger"
     print $ I.nextPrimeInteger b
diff --git a/testsuite/tests/lib/integer/integerGmpInternals.stdout b/testsuite/tests/lib/integer/integerGmpInternals.stdout
index 4a3eac5055..1ffeabb78b 100644
--- a/testsuite/tests/lib/integer/integerGmpInternals.stdout
+++ b/testsuite/tests/lib/integer/integerGmpInternals.stdout
@@ -3,7 +3,6 @@
 1
 12345
 555562377826831043419246079513769804614412256811161773362797946971665712715296306339052301636736176350153982639312744140625
-[(-7,149867),(-5,167851),(-1,209813),(1,1),(5,41963),(7,59947),(13,177535),(19,143557),(23,182447),(25,134281),(29,7235),(31,33841),(35,95915),(37,113413),(41,61409),(43,24397),(47,174101),(49,158431),(53,193979),(59,188477),(61,185737),(65,35507),(67,118999),(71,186173)]
 
 # nextPrimeInteger
 2988348162058574136915891421498819466320163312926952423791023078876343
diff --git a/testsuite/tests/lib/integer/integerRecipMod.hs b/testsuite/tests/lib/integer/integerRecipMod.hs
new file mode 100644
index 0000000000..e5de646790
--- /dev/null
+++ b/testsuite/tests/lib/integer/integerRecipMod.hs
@@ -0,0 +1,39 @@
+{-# LANGUAGE BangPatterns, CPP, MagicHash, UnboxedTuples, TupleSections #-}
+
+module Main (main) where
+
+import Data.List (group)
+import Data.Bits
+import Data.Word
+import Data.Maybe
+import Control.Monad
+
+import GHC.Word
+import GHC.Base
+import GHC.Num.Integer
+import qualified GHC.Num.Integer as I
+
+recipModInteger :: Integer -> Integer -> Maybe Integer
+recipModInteger x m = case I.integerRecipMod# x m of
+   (# y |    #) -> Just y
+   (#   | () #) -> Nothing
+
+main :: IO ()
+main = do
+   let
+      f x = case recipModInteger x (2*3*11*11*17*17) of
+               y -> fmap (x,) y
+
+      g x = case recipModInteger x (-2*3*11*11*17*17) of
+               y -> fmap (x,) y
+
+   -- positive modulo
+   print $ mapMaybe f [-7..71]
+
+   -- negative modulo
+   print $ mapMaybe g [-7..71]
+
+   -- modulo == 1, -1 or 0
+   print (recipModInteger 77 1)
+   print (recipModInteger 77 (-1))
+   print (recipModInteger 77 0)
diff --git a/testsuite/tests/lib/integer/integerRecipMod.stdout b/testsuite/tests/lib/integer/integerRecipMod.stdout
new file mode 100644
index 0000000000..205af95fb9
--- /dev/null
+++ b/testsuite/tests/lib/integer/integerRecipMod.stdout
@@ -0,0 +1,5 @@
+[(-7,149867),(-5,167851),(-1,209813),(1,1),(5,41963),(7,59947),(13,177535),(19,143557),(23,182447),(25,134281),(29,7235),(31,33841),(35,95915),(37,113413),(41,61409),(43,24397),(47,174101),(49,158431),(53,193979),(59,188477),(61,185737),(65,35507),(67,118999),(71,186173)]
+[(-7,149867),(-5,167851),(-1,209813),(1,1),(5,41963),(7,59947),(13,177535),(19,143557),(23,182447),(25,134281),(29,7235),(31,33841),(35,95915),(37,113413),(41,61409),(43,24397),(47,174101),(49,158431),(53,193979),(59,188477),(61,185737),(65,35507),(67,118999),(71,186173)]
+Nothing
+Nothing
+Nothing