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.

7 files changed, 143 insertions, 22 deletions

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#