Bignum: allow inlining of naturalEq/Ne/Gt/Lt/Ge/Le/Compare (#20361)

Perform constant folding on bigNatCompare instead.

Some functions of the Enum class for Natural now need to be inlined
explicitly to be specialized at call sites (because `x > lim` for
Natural is inlined and the resulting function is a little too big to
inline). If we don't do this, T17499 runtime allocations regresses by
16%.

5 files changed, 34 insertions, 69 deletions

diff --git a/compiler/GHC/Builtin/Names.hs b/compiler/GHC/Builtin/Names.hs
index 91c9c80cef..5407e16b90 100644
--- a/compiler/GHC/Builtin/Names.hs
+++ b/compiler/GHC/Builtin/Names.hs
@@ -398,11 +398,6 @@ basicKnownKeyNames
 
         naturalToWordName,
         naturalToWordClampName,
-        naturalGeName,
-        naturalLeName,
-        naturalGtName,
-        naturalLtName,
-        naturalCompareName,
         naturalPopCountName,
         naturalShiftRName,
         naturalShiftLName,
@@ -1195,11 +1190,6 @@ integerFromNaturalName
    , integerShiftRName
    , naturalToWordName
    , naturalToWordClampName
-   , naturalGeName
-   , naturalLeName
-   , naturalGtName
-   , naturalLtName
-   , naturalCompareName
    , naturalPopCountName
    , naturalShiftRName
    , naturalShiftLName
@@ -1228,6 +1218,7 @@ integerFromNaturalName
    , naturalSizeInBaseName
    , bignatFromWordListName
    , bignatEqName
+   , bignatCompareName
    :: Name
 
 bnbVarQual, bnnVarQual, bniVarQual :: String -> Unique -> Name
@@ -1238,14 +1229,10 @@ bniVarQual str key = varQual gHC_NUM_INTEGER (fsLit str) key
 -- Types and DataCons
 bignatFromWordListName    = bnbVarQual "bigNatFromWordList#"       bignatFromWordListIdKey
 bignatEqName              = bnbVarQual "bigNatEq#"                 bignatEqIdKey
+bignatCompareName         = bnbVarQual "bigNatCompare"             bignatCompareIdKey
 
 naturalToWordName         = bnnVarQual "naturalToWord#"            naturalToWordIdKey
 naturalToWordClampName    = bnnVarQual "naturalToWordClamp#"       naturalToWordClampIdKey
-naturalGeName             = bnnVarQual "naturalGe#"                naturalGeIdKey
-naturalLeName             = bnnVarQual "naturalLe#"                naturalLeIdKey
-naturalGtName             = bnnVarQual "naturalGt#"                naturalGtIdKey
-naturalLtName             = bnnVarQual "naturalLt#"                naturalLtIdKey
-naturalCompareName        = bnnVarQual "naturalCompare"            naturalCompareIdKey
 naturalPopCountName       = bnnVarQual "naturalPopCount#"          naturalPopCountIdKey
 naturalShiftRName         = bnnVarQual "naturalShiftR#"            naturalShiftRIdKey
 naturalShiftLName         = bnnVarQual "naturalShiftL#"            naturalShiftLIdKey
@@ -2592,14 +2579,8 @@ integerFromNaturalIdKey
    , integerFromWordIdKey
    , integerFromWord64IdKey
    , integerFromInt64IdKey
-   , bignatEqIdKey
    , naturalToWordIdKey
    , naturalToWordClampIdKey
-   , naturalGeIdKey
-   , naturalLeIdKey
-   , naturalGtIdKey
-   , naturalLtIdKey
-   , naturalCompareIdKey
    , naturalPopCountIdKey
    , naturalShiftRIdKey
    , naturalShiftLIdKey
@@ -2627,6 +2608,8 @@ integerFromNaturalIdKey
    , naturalPowModIdKey
    , naturalSizeInBaseIdKey
    , bignatFromWordListIdKey
+   , bignatEqIdKey
+   , bignatCompareIdKey
    :: Unique
 
 integerFromNaturalIdKey    = mkPreludeMiscIdUnique 600
@@ -2675,11 +2658,6 @@ integerFromInt64IdKey      = mkPreludeMiscIdUnique 644
 
 naturalToWordIdKey         = mkPreludeMiscIdUnique 650
 naturalToWordClampIdKey    = mkPreludeMiscIdUnique 651
-naturalGeIdKey             = mkPreludeMiscIdUnique 654
-naturalLeIdKey             = mkPreludeMiscIdUnique 655
-naturalGtIdKey             = mkPreludeMiscIdUnique 656
-naturalLtIdKey             = mkPreludeMiscIdUnique 657
-naturalCompareIdKey        = mkPreludeMiscIdUnique 658
 naturalPopCountIdKey       = mkPreludeMiscIdUnique 659
 naturalShiftRIdKey         = mkPreludeMiscIdUnique 660
 naturalShiftLIdKey         = mkPreludeMiscIdUnique 661
@@ -2709,6 +2687,7 @@ naturalSizeInBaseIdKey     = mkPreludeMiscIdUnique 684
 
 bignatFromWordListIdKey    = mkPreludeMiscIdUnique 690
 bignatEqIdKey              = mkPreludeMiscIdUnique 691
+bignatCompareIdKey         = mkPreludeMiscIdUnique 692
 
 
 ------------------------------------------------------
diff --git a/compiler/GHC/Core/Opt/ConstantFold.hs b/compiler/GHC/Core/Opt/ConstantFold.hs
index b117fa11b2..3d327e69eb 100644
--- a/compiler/GHC/Core/Opt/ConstantFold.hs
+++ b/compiler/GHC/Core/Opt/ConstantFold.hs
@@ -1514,9 +1514,9 @@ isLiteral e = do
         Nothing -> mzero
         Just l  -> pure l
 
--- | Match Integer and Natural literals
+-- | Match BigNat#, Integer and Natural literals
 isBignumLiteral :: CoreExpr -> RuleM Integer
-isBignumLiteral e = isIntegerLiteral e <|> isNaturalLiteral e
+isBignumLiteral e = isNumberLiteral e <|> isIntegerLiteral e <|> isNaturalLiteral e
 
 -- | Match numeric literals
 isNumberLiteral :: CoreExpr -> RuleM Integer
@@ -2056,23 +2056,17 @@ builtinBignumRules =
   , natural_to_word "Natural -> Word# (clamp)" naturalToWordClampName True
 
     -- comparisons (return an unlifted Int#)
-  , integer_cmp "integerEq#" integerEqName (==)
-  , integer_cmp "integerNe#" integerNeName (/=)
-  , integer_cmp "integerLe#" integerLeName (<=)
-  , integer_cmp "integerGt#" integerGtName (>)
-  , integer_cmp "integerLt#" integerLtName (<)
-  , integer_cmp "integerGe#" integerGeName (>=)
-
-  , natural_cmp "naturalLe#" naturalLeName (<=)
-  , natural_cmp "naturalGt#" naturalGtName (>)
-  , natural_cmp "naturalLt#" naturalLtName (<)
-  , natural_cmp "naturalGe#" naturalGeName (>=)
-
-  , bignat_cmp  "bigNatEq#"  bignatEqName (==)
+  , bignum_bin_pred "integerEq#" integerEqName (==)
+  , bignum_bin_pred "integerNe#" integerNeName (/=)
+  , bignum_bin_pred "integerLe#" integerLeName (<=)
+  , bignum_bin_pred "integerGt#" integerGtName (>)
+  , bignum_bin_pred "integerLt#" integerLtName (<)
+  , bignum_bin_pred "integerGe#" integerGeName (>=)
+  , bignum_bin_pred "bigNatEq#"  bignatEqName (==)
 
     -- comparisons (return an Ordering)
   , bignum_compare "integerCompare" integerCompareName
-  , bignum_compare "naturalCompare" naturalCompareName
+  , bignum_compare "bignatCompare" bignatCompareName
 
     -- binary operations
   , integer_binop "integerAdd" integerAddName (+)
@@ -2205,7 +2199,7 @@ builtinBignumRules =
     lit_to_integer str name = mkRule str name 1 $ do
       [a0] <- getArgs
       platform <- getPlatform
-      i <- isNumberLiteral a0 <|> isBignumLiteral a0
+      i <- isBignumLiteral a0
       -- convert any numeric literal into an Integer literal
       pure (mkIntegerExpr platform i)
 
@@ -2231,29 +2225,11 @@ builtinBignumRules =
       platform <- getPlatform
       pure (mkNaturalExpr platform (x - y))
 
-    integer_cmp str name op = mkRule str name 2 $ do
+    bignum_bin_pred str name op = mkRule str name 2 $ do
       platform <- getPlatform
       [a0,a1] <- getArgs
-      x <- isIntegerLiteral a0
-      y <- isIntegerLiteral a1
-      pure $ if x `op` y
-              then trueValInt platform
-              else falseValInt platform
-
-    natural_cmp str name op = mkRule str name 2 $ do
-      platform <- getPlatform
-      [a0,a1] <- getArgs
-      x <- isNaturalLiteral a0
-      y <- isNaturalLiteral a1
-      pure $ if x `op` y
-              then trueValInt platform
-              else falseValInt platform
-
-    bignat_cmp str name op = mkRule str name 2 $ do
-      platform <- getPlatform
-      [a0,a1] <- getArgs
-      x <- isNumberLiteral a0
-      y <- isNumberLiteral a1
+      x <- isBignumLiteral a0
+      y <- isBignumLiteral a1
       pure $ if x `op` y
               then trueValInt platform
               else falseValInt platform
diff --git a/libraries/base/GHC/Enum.hs b/libraries/base/GHC/Enum.hs
index 096e633894..4292592f44 100644
--- a/libraries/base/GHC/Enum.hs
+++ b/libraries/base/GHC/Enum.hs
@@ -989,12 +989,22 @@ instance Enum Natural where
         i = I# (word2Int# w)
     fromEnum _ = errorWithoutStackTrace "fromEnum: out of Int range"
 
+    -- See Note [Stable Unfolding for list producers]
+    {-# INLINE enumFrom #-}
     enumFrom x        = enumDeltaNatural      x 1
+
+    -- See Note [Stable Unfolding for list producers]
+    {-# INLINE enumFromThen #-}
     enumFromThen x y
       | x <= y        = enumDeltaNatural      x (y-x)
       | otherwise     = enumNegDeltaToNatural x (x-y) 0
 
+    -- See Note [Stable Unfolding for list producers]
+    {-# INLINE enumFromTo #-}
     enumFromTo x lim  = enumDeltaToNatural    x 1 lim
+
+    -- See Note [Stable Unfolding for list producers]
+    {-# INLINE enumFromThenTo #-}
     enumFromThenTo x y lim
       | x <= y        = enumDeltaToNatural    x (y-x) lim
       | otherwise     = enumNegDeltaToNatural x (x-y) lim
@@ -1004,12 +1014,16 @@ instance Enum Natural where
 enumDeltaNatural :: Natural -> Natural -> [Natural]
 enumDeltaNatural !x d = x : enumDeltaNatural (x+d) d
 
+-- Inline to specialize
+{-# INLINE enumDeltaToNatural #-}
 enumDeltaToNatural :: Natural -> Natural -> Natural -> [Natural]
 enumDeltaToNatural x0 delta lim = go x0
   where
     go x | x > lim   = []
          | otherwise = x : go (x+delta)
 
+-- Inline to specialize
+{-# INLINE enumNegDeltaToNatural #-}
 enumNegDeltaToNatural :: Natural -> Natural -> Natural -> [Natural]
 enumNegDeltaToNatural x0 ndelta lim = go x0
   where
diff --git a/libraries/ghc-bignum/src/GHC/Num/BigNat.hs b/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
index 1e6acaf9cc..7749ffcbcd 100644
--- a/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
@@ -344,6 +344,7 @@ bigNatCompareWord a (W# b) = bigNatCompareWord# a b
 
 -- | Compare two BigNat
 bigNatCompare :: BigNat# -> BigNat# -> Ordering
+{-# NOINLINE bigNatCompare #-}
 bigNatCompare a b =
    let
       szA = wordArraySize# a
diff --git a/libraries/ghc-bignum/src/GHC/Num/Natural.hs b/libraries/ghc-bignum/src/GHC/Num/Natural.hs
index 3fe09a0625..a494ccee70 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Natural.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Natural.hs
@@ -162,7 +162,6 @@ naturalNe !x !y = isTrue# (naturalNe# x y)
 
 -- | Greater or equal test for Natural
 naturalGe# :: Natural -> Natural -> Bool#
-{-# NOINLINE naturalGe# #-}
 naturalGe# (NS x) (NS y) = x `geWord#` y
 naturalGe# (NS _) (NB _) = 0#
 naturalGe# (NB _) (NS _) = 1#
@@ -174,7 +173,6 @@ naturalGe !x !y = isTrue# (naturalGe# x y)
 
 -- | Lower or equal test for Natural
 naturalLe# :: Natural -> Natural -> Bool#
-{-# NOINLINE naturalLe# #-}
 naturalLe# (NS x) (NS y) = x `leWord#` y
 naturalLe# (NS _) (NB _) = 1#
 naturalLe# (NB _) (NS _) = 0#
@@ -187,7 +185,6 @@ naturalLe !x !y = isTrue# (naturalLe# x y)
 
 -- | Greater test for Natural
 naturalGt# :: Natural -> Natural -> Bool#
-{-# NOINLINE naturalGt# #-}
 naturalGt# (NS x) (NS y) = x `gtWord#` y
 naturalGt# (NS _) (NB _) = 0#
 naturalGt# (NB _) (NS _) = 1#
@@ -199,7 +196,6 @@ naturalGt !x !y = isTrue# (naturalGt# x y)
 
 -- | Lower test for Natural
 naturalLt# :: Natural -> Natural -> Bool#
-{-# NOINLINE naturalLt# #-}
 naturalLt# (NS x) (NS y) = x `ltWord#` y
 naturalLt# (NS _) (NB _) = 1#
 naturalLt# (NB _) (NS _) = 0#
@@ -211,7 +207,6 @@ naturalLt !x !y = isTrue# (naturalLt# x y)
 
 -- | Compare two Natural
 naturalCompare :: Natural -> Natural -> Ordering
-{-# NOINLINE naturalCompare #-}
 naturalCompare (NS x) (NS y) = cmpW# x y
 naturalCompare (NB x) (NB y) = bigNatCompare x y
 naturalCompare (NS _) (NB _) = LT