Bignum: implement integerPowMod (#18427)

Incidentally fix powModInteger which was crashing in integer-gmp for
negative exponents when the modular multiplicative inverse for the base
didn't exist. Now we compute it explicitly with integerRecipMod so that
every backend returns the same result without crashing.

13 files changed, 190 insertions, 60 deletions

diff --git a/libraries/ghc-bignum/cbits/gmp_wrappers.c b/libraries/ghc-bignum/cbits/gmp_wrappers.c
index e036f05a9a..5c12346500 100644
--- a/libraries/ghc-bignum/cbits/gmp_wrappers.c
+++ b/libraries/ghc-bignum/cbits/gmp_wrappers.c
@@ -632,8 +632,7 @@ integer_gmp_powm(mp_limb_t rp[], // result
   assert(!mp_limb_zero_p(mp,mn));
 
   if ((mn == 1 || mn == -1) && mp[0] == 1) {
-    rp[0] = 0;
-    return 1;
+    return 0;
   }
 
   if (mp_limb_zero_p(ep,en)) {
@@ -659,11 +658,6 @@ integer_gmp_powm(mp_limb_t rp[], // result
 
   mpz_clear (r);
 
-  if (!rn) {
-    rp[0] = 0;
-    return 1;
-  }
-
   return rn;
 }
 
@@ -720,8 +714,7 @@ integer_gmp_powm_sec(mp_limb_t rp[], // result
   assert(mp[0] & 1);
 
   if ((mn == 1 || mn == -1) && mp[0] == 1) {
-    rp[0] = 0;
-    return 1;
+    return 0;
   }
 
   if (mp_limb_zero_p(ep,en)) {
@@ -753,11 +746,6 @@ integer_gmp_powm_sec(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 9e5181c0bb..b23fd1cb21 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Backend/Check.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Backend/Check.hs
@@ -17,6 +17,7 @@ import GHC.Types
 import GHC.Num.WordArray
 import GHC.Num.Primitives
 import {-# SOURCE #-} GHC.Num.Integer
+import {-# SOURCE #-} GHC.Num.Natural
 import qualified GHC.Num.Backend.Native   as Native
 import qualified GHC.Num.Backend.Selected as Other
 
@@ -472,8 +473,8 @@ integer_gcde a b =
 
 integer_recip_mod
    :: Integer
-   -> Integer
-   -> (# Integer | () #)
+   -> Natural
+   -> (# Natural | () #)
 integer_recip_mod x m =
    let
       !r0 = Other.integer_recip_mod x m
@@ -481,6 +482,20 @@ integer_recip_mod x m =
    in case (# r0, r1 #) of
          (# (# | () #), (# | () #) #) -> r0
          (# (# y0 | #), (# y1 | #) #)
-            | isTrue# (integerEq# y0 y1) -> r0
+            | isTrue# (naturalEq# y0 y1) -> r0
          _ -> case unexpectedValue of
             !_ -> (# | () #)
+
+integer_powmod
+   :: Integer
+   -> Natural
+   -> Natural
+   -> Natural
+integer_powmod b e m =
+   let
+      !r0 = Other.integer_powmod b e m
+      !r1 = Native.integer_powmod b e m
+   in if isTrue# (naturalEq# r0 r1)
+         then r0
+         else case unexpectedValue of
+               !_ -> naturalZero
diff --git a/libraries/ghc-bignum/src/GHC/Num/Backend/FFI.hs b/libraries/ghc-bignum/src/GHC/Num/Backend/FFI.hs
index 98350d4d29..21d74ee3e4 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Backend/FFI.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Backend/FFI.hs
@@ -20,6 +20,8 @@ import GHC.Types
 import GHC.Num.WordArray
 import GHC.Num.Primitives
 import qualified GHC.Num.Backend.Native as Native
+import {-# SOURCE #-} GHC.Num.Natural
+import {-# SOURCE #-} GHC.Num.Integer
 
 default ()
 
@@ -606,9 +608,25 @@ integer_gcde = Native.integer_gcde
 -- with 0 < y < abs m
 integer_recip_mod
    :: Integer
-   -> Integer
-   -> (# Integer | () #)
+   -> Natural
+   -> (# Natural | () #)
 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.
+
+-- | Computes the modular exponentiation.
+--
+-- I.e. y = integer_powmod b e m
+--        = b^e `mod` m
+--
+-- with 0 <= y < abs m
+integer_powmod
+   :: Integer
+   -> Natural
+   -> Natural
+   -> Natural
+integer_powmod = Native.integer_powmod
+   -- 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 462e67445c..8db1dd57b2 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Backend/GMP.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Backend/GMP.hs
@@ -28,6 +28,7 @@ import GHC.Types
 import GHC.Magic (runRW#)
 import {-# SOURCE #-} GHC.Num.Integer
 import {-# SOURCE #-} GHC.Num.BigNat
+import {-# SOURCE #-} GHC.Num.Natural
 
 default ()
 
@@ -354,8 +355,37 @@ bignat_powmod
    -> WordArray#
    -> State# RealWorld
    -> State# RealWorld
-bignat_powmod r b e m s =
-   case ioInt# (integer_gmp_powm# r b (wordArraySize# b) e (wordArraySize# e) m (wordArraySize# m)) s of
+bignat_powmod r b e m s = sbignat_powmod r (wordArraySize# b) b e m s
+
+integer_powmod
+   :: Integer
+   -> Natural
+   -> Natural
+   -> Natural
+integer_powmod b e m = naturalFromBigNat# (withNewWordArray# szm io)
+   where
+      !be = naturalToBigNat# e
+      !bm = naturalToBigNat# m
+      !(# sb, bb #) = integerToBigNatSign# b
+      !szb = bigNatSize# bb
+      !szm = bigNatSize# bm
+      !ssb = case sb of -- signed size of b
+               0# -> szb
+               _  -> negateInt# szb
+
+      io r s = sbignat_powmod r ssb bb be bm s
+
+
+sbignat_powmod
+   :: MutableWordArray# RealWorld
+   -> Int#
+   -> WordArray#
+   -> WordArray#
+   -> WordArray#
+   -> State# RealWorld
+   -> State# RealWorld
+sbignat_powmod r b_signed_size b e m s =
+   case ioInt# (integer_gmp_powm# r b b_signed_size e (wordArraySize# e) m (wordArraySize# m)) s of
       (# s', n #) -> mwaSetSize# r (narrowGmpSize# n) s'
 
 integer_gcde
@@ -425,16 +455,16 @@ integer_gcde a b = case runRW# io of (# _, a #) -> a
 
 integer_recip_mod
    :: Integer
-   -> Integer
-   -> (# Integer | () #)
+   -> Natural
+   -> (# Natural | () #)
 integer_recip_mod x m =
    let
       !(#  sign_x, bx #) = integerToBigNatSign# x
-      !(# _sign_m, bm #) = integerToBigNatSign# m
+      !bm = naturalToBigNat# m
       !br = sbignat_recip_mod sign_x bx bm
    in if isTrue# (bigNatIsZero# br)
          then (# | () #)
-         else (# integerFromBigNat# br | #)
+         else (# naturalFromBigNat# br | #)
 
 
 -- | Return 0 for invalid inputs
diff --git a/libraries/ghc-bignum/src/GHC/Num/Backend/Native.hs b/libraries/ghc-bignum/src/GHC/Num/Backend/Native.hs
index c5af604c17..ee031178bd 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Backend/Native.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Backend/Native.hs
@@ -14,7 +14,7 @@ module GHC.Num.Backend.Native where
 #include "MachDeps.h"
 #include "WordSize.h"
 
-#if defined(BIGNUM_NATIVE) || defined(BIGNUM_CHECK)
+#if defined(BIGNUM_NATIVE) || defined(BIGNUM_CHECK) || defined(BIGNUM_FFI)
 import {-# SOURCE #-} GHC.Num.BigNat
 import {-# SOURCE #-} GHC.Num.Natural
 import {-# SOURCE #-} GHC.Num.Integer
@@ -740,14 +740,37 @@ integer_gcde a b = f (# a,integerOne,integerZero #) (# b,integerZero,integerOne
 
 integer_recip_mod
    :: Integer
-   -> Integer
-   -> (# Integer | () #)
+   -> Natural
+   -> (# Natural | () #)
 integer_recip_mod x m =
-   let m' = integerAbs m
+   let m' = integerFromNatural 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' | #)
+         | g `integerEq` integerOne -> (# integerToNatural (a `integerMod` m') | #)
+                                       -- a `mod` m > 0 because m > 0
          | True                     -> (# | () #)
+
+integer_powmod
+   :: Integer
+   -> Natural
+   -> Natural
+   -> Natural
+integer_powmod b0 e0 m = go b0 e0 integerOne
+   where
+      !m' = integerFromNatural m
+
+      go !b e !r
+        | isTrue# (e `naturalTestBit#` 0##)
+        = go b' e' ((r `integerMul` b) `integerMod` m')
+
+        | naturalIsZero e
+        = integerToNatural r -- r >= 0 by integerMod above
+
+        | True
+        = go b' e' r
+        where
+          b' = (b `integerMul` b) `integerRem` m'
+          e' = e `naturalShiftR#` 1##
diff --git a/libraries/ghc-bignum/src/GHC/Num/Integer.hs b/libraries/ghc-bignum/src/GHC/Num/Integer.hs
index 7f0314f859..6ee2d27901 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Integer.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Integer.hs
@@ -1233,11 +1233,39 @@ integerGcde a b = case integerGcde# a b of
 --
 integerRecipMod#
    :: Integer
-   -> Integer
-   -> (# Integer | () #)
+   -> Natural
+   -> (# Natural | () #)
 integerRecipMod# x m
    | integerIsZero x = (# | () #)
-   | integerIsZero m = (# | () #)
-   | IS 1#    <- m   = (# | () #)
-   | IS (-1#) <- m   = (# | () #)
+   | naturalIsZero m = (# | () #)
+   | naturalIsOne  m = (# | () #)
    | True            = Backend.integer_recip_mod x m
+
+
+-- | Computes the modular exponentiation.
+--
+-- I.e. y = integer_powmod b e m
+--        = b^e `mod` m
+--
+-- with 0 <= y < abs m
+--
+-- If e is negative, we use `integerRecipMod#` to try to find a modular
+-- multiplicative inverse (which may not exist).
+integerPowMod# :: Integer -> Integer -> Natural -> (# Natural | () #)
+integerPowMod# !b !e !m
+   | naturalIsZero m     = (# | () #)
+   | naturalIsOne  m     = (# naturalZero | #)
+   | integerIsZero e     = (# naturalOne  | #)
+   | integerIsZero b     = (# naturalZero | #)
+   | integerIsOne  b     = (# naturalOne  | #)
+     -- when the exponent is negative, try to find the modular multiplicative
+     -- inverse and use it instead
+   | integerIsNegative e = case integerRecipMod# b m of
+      (#    | () #) -> (# | () #)
+      (# b' |    #) -> integerPowMod#
+                        (integerFromNatural b')
+                        (integerNegate e)
+                        m
+
+     -- e > 0 by cases above
+   | True = (# Backend.integer_powmod b (integerToNatural e) m | #)
diff --git a/libraries/ghc-bignum/src/GHC/Num/Integer.hs-boot b/libraries/ghc-bignum/src/GHC/Num/Integer.hs-boot
index e9c2d003ca..80ecd36a34 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Integer.hs-boot
+++ b/libraries/ghc-bignum/src/GHC/Num/Integer.hs-boot
@@ -7,6 +7,7 @@ module GHC.Num.Integer where
 import GHC.Types
 import GHC.Prim
 import {-# SOURCE #-} GHC.Num.BigNat
+import {-# SOURCE #-} GHC.Num.Natural
 
 data Integer
 
@@ -32,3 +33,5 @@ integerQuotRem# :: Integer -> Integer -> (# Integer, Integer #)
 integerToBigNatSign# :: Integer -> (# Int#, BigNat# #)
 integerFromBigNatSign# :: Int# -> BigNat# -> Integer
 integerFromBigNat# :: BigNat# -> Integer
+integerToNatural :: Integer -> Natural
+integerFromNatural :: Natural -> Integer
diff --git a/libraries/ghc-bignum/src/GHC/Num/Natural.hs-boot b/libraries/ghc-bignum/src/GHC/Num/Natural.hs-boot
index 964292fa59..57e1a9d9a9 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Natural.hs-boot
+++ b/libraries/ghc-bignum/src/GHC/Num/Natural.hs-boot
@@ -16,8 +16,12 @@ naturalToWord# :: Natural -> Word#
 naturalFromWord# :: Word# -> Natural
 naturalFromBigNat# :: BigNat# -> Natural
 naturalToBigNat# :: Natural -> BigNat#
+
+naturalZero :: Natural
 naturalMul :: Natural -> Natural -> Natural
 naturalRem :: Natural -> Natural -> Natural
-naturalIsZero :: Natural -> Bool
 naturalShiftR# :: Natural -> Word# -> Natural
+
+naturalIsZero :: Natural -> Bool
 naturalTestBit# :: Natural -> Word# -> Bool#
+naturalEq# :: Natural -> Natural -> Bool#
diff --git a/libraries/integer-gmp/src/GHC/Integer/GMP/Internals.hs b/libraries/integer-gmp/src/GHC/Integer/GMP/Internals.hs
index cc1068ba10..2fcb0750ed 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
+    , powModInteger
     , recipModInteger
 
       -- ** Additional conversion operations to 'Integer'
@@ -189,10 +190,16 @@ sqrInteger = I.integerSqr
 
 {-# DEPRECATED recipModInteger "Use integerRecipMod# instead" #-}
 recipModInteger :: Integer -> Integer -> Integer
-recipModInteger x m = case I.integerRecipMod# x m of
-   (# y |    #) -> y
+recipModInteger x m = case I.integerRecipMod# x (I.integerToNatural m) of
+   (# y |    #) -> I.integerFromNatural y
    (#   | () #) -> 0
 
+{-# DEPRECATED powModInteger "Use integerPowMod# instead" #-}
+powModInteger :: Integer -> Integer -> Integer -> Integer
+powModInteger b e m = case I.integerPowMod# b e (I.integerToNatural m) of
+   (# r  | #) -> I.integerFromNatural r
+   (# | () #) -> 0
+
 {-# DEPRECATED wordToNegInteger "Use integerFromWordNeg# instead" #-}
 wordToNegInteger :: Word# -> Integer
 wordToNegInteger = I.integerFromWordNeg#
diff --git a/testsuite/tests/lib/integer/integerPowMod.hs b/testsuite/tests/lib/integer/integerPowMod.hs
index 497e96cbf9..7a027fcc8a 100644
--- a/testsuite/tests/lib/integer/integerPowMod.hs
+++ b/testsuite/tests/lib/integer/integerPowMod.hs
@@ -1,3 +1,6 @@
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE MagicHash #-}
+
 module Main (main) where
 
 import Data.List (group)
@@ -7,22 +10,33 @@ import Control.Monad
 
 import GHC.Word
 import GHC.Base
-import GHC.Natural
+import GHC.Num.Natural
+import GHC.Num.Integer
+
+integerPowMod :: Integer -> Integer -> Integer -> Maybe Integer
+integerPowMod b e m = case integerPowMod# b e (fromIntegral m) of
+   (# | () #) -> Nothing
+   (# r  | #) -> Just (fromIntegral r)
 
 main :: IO ()
 main = do
-    print $ powModNatural b e m
-    print $ powModNatural b e (m-1)
+    print $ naturalPowMod b e m
+    print $ naturalPowMod b e (m-1)
+
+    print $ integerPowMod b e m
+    print $ integerPowMod b e (m-1)
+
+    print $ integerPowMod b (-e) m
+    print $ integerPowMod b (-e) (m-1)
+
+    print $ integerPowMod (-b) e m
+    print $ integerPowMod (-b) e (m-1)
+
+    print $ integerPowMod (-b) (-e) m
+    print $ integerPowMod (-b) (-e) (m-1)
 
   where
+    b,e,m :: Num a => a
     b = 2988348162058574136915891421498819466320163312926952423791023078876139
     e = 2351399303373464486466122544523690094744975233415544072992656881240319
     m = 10^(40::Int)
-
-    x = 5328841272400314897981163497728751426
-    y = 32052182750761975518649228050096851724
-
-    b1024 = roll (map fromIntegral (take 128 [0x80::Int .. ]))
-
-    roll :: [Word8] -> Integer
-    roll = GHC.Base.foldr (\b a -> a `shiftL` 8 .|. fromIntegral b) 0
diff --git a/testsuite/tests/lib/integer/integerPowMod.stdout b/testsuite/tests/lib/integer/integerPowMod.stdout
index 64a4c568ac..c095505778 100644
--- a/testsuite/tests/lib/integer/integerPowMod.stdout
+++ b/testsuite/tests/lib/integer/integerPowMod.stdout
@@ -1,2 +1,10 @@
 1527229998585248450016808958343740453059
 682382427572745901624116300491295556924
+Just 1527229998585248450016808958343740453059
+Just 682382427572745901624116300491295556924
+Just 9710805908340105786808462779613285007339
+Nothing
+Just 8472770001414751549983191041656259546941
+Just 9317617572427254098375883699508704443075
+Just 289194091659894213191537220386714992661
+Nothing
diff --git a/testsuite/tests/lib/integer/integerRecipMod.hs b/testsuite/tests/lib/integer/integerRecipMod.hs
index e5de646790..aad4a7e33b 100644
--- a/testsuite/tests/lib/integer/integerRecipMod.hs
+++ b/testsuite/tests/lib/integer/integerRecipMod.hs
@@ -11,9 +11,10 @@ import Control.Monad
 import GHC.Word
 import GHC.Base
 import GHC.Num.Integer
+import GHC.Num.Natural
 import qualified GHC.Num.Integer as I
 
-recipModInteger :: Integer -> Integer -> Maybe Integer
+recipModInteger :: Integer -> Natural -> Maybe Natural
 recipModInteger x m = case I.integerRecipMod# x m of
    (# y |    #) -> Just y
    (#   | () #) -> Nothing
@@ -24,16 +25,9 @@ main = do
       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
+   -- modulo == 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
index 205af95fb9..b55b6b4040 100644
--- a/testsuite/tests/lib/integer/integerRecipMod.stdout
+++ b/testsuite/tests/lib/integer/integerRecipMod.stdout
@@ -1,5 +1,3 @@
 [(-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