Algebraically simplify add/sub with carry/overflow

Previously, the `{add,sub}{Int,Word}C#` PrimOps weren't handled
in PrelRules (constant folding and algebraic simplification) at all.
This implements the necessary logic, so that using these primitives
isn't too punishing compared to their well-optimised, overflow-unaware
counterparts.

This is so that using these primitives in `enumFromThenTo @Int` can
be optimized by constant folding, reducing closure sizes.

Reviewers: bgamari, simonpj, hsyl20

Reviewed By: bgamari, simonpj

Subscribers: AndreasK, thomie, carter

GHC Trac Issues: #8763

Differential Revision: https://phabricator.haskell.org/D4605

6 files changed, 150 insertions, 21 deletions

diff --git a/compiler/basicTypes/Literal.hs b/compiler/basicTypes/Literal.hs
index f81f45224c..0392a98274 100644
--- a/compiler/basicTypes/Literal.hs
+++ b/compiler/basicTypes/Literal.hs
@@ -13,8 +13,8 @@ module Literal
           Literal(..)           -- Exported to ParseIface
 
         -- ** Creating Literals
-        , mkMachInt, mkMachIntWrap
-        , mkMachWord, mkMachWordWrap
+        , mkMachInt, mkMachIntWrap, mkMachIntWrapC
+        , mkMachWord, mkMachWordWrap, mkMachWordWrapC
         , mkMachInt64, mkMachInt64Wrap
         , mkMachWord64, mkMachWord64Wrap
         , mkMachFloat, mkMachDouble
@@ -247,30 +247,54 @@ mkMachInt :: DynFlags -> Integer -> Literal
 mkMachInt dflags x   = ASSERT2( inIntRange dflags x,  integer x )
                        MachInt x
 
+wrapInt :: DynFlags -> Integer -> Integer
+wrapInt dflags i
+ = case platformWordSize (targetPlatform dflags) of
+   4 -> toInteger (fromIntegral i :: Int32)
+   8 -> toInteger (fromIntegral i :: Int64)
+   w -> panic ("toIntRange: Unknown platformWordSize: " ++ show w)
+
 -- | Creates a 'Literal' of type @Int#@.
 --   If the argument is out of the (target-dependent) range, it is wrapped.
 --   See Note [Word/Int underflow/overflow]
 mkMachIntWrap :: DynFlags -> Integer -> Literal
-mkMachIntWrap dflags i
- = MachInt $ case platformWordSize (targetPlatform dflags) of
-   4 -> toInteger (fromIntegral i :: Int32)
-   8 -> toInteger (fromIntegral i :: Int64)
-   w -> panic ("toIntRange: Unknown platformWordSize: " ++ show w)
+mkMachIntWrap dflags i = MachInt (wrapInt dflags i)
+
+-- | Creates a 'Literal' of type @Int#@, as well as a 'Bool'ean flag indicating
+--   overflow. That is, if the argument is out of the (target-dependent) range
+--   the argument is wrapped and the overflow flag will be set.
+--   See Note [Word/Int underflow/overflow]
+mkMachIntWrapC :: DynFlags -> Integer -> (Literal, Bool)
+mkMachIntWrapC dflags i = (MachInt i', i /= i')
+  where
+    i' = wrapInt dflags i
 
 -- | Creates a 'Literal' of type @Word#@
 mkMachWord :: DynFlags -> Integer -> Literal
 mkMachWord dflags x   = ASSERT2( inWordRange dflags x, integer x )
                         MachWord x
 
+wrapWord :: DynFlags -> Integer -> Integer
+wrapWord dflags i
+ = case platformWordSize (targetPlatform dflags) of
+   4 -> toInteger (fromIntegral i :: Word32)
+   8 -> toInteger (fromIntegral i :: Word64)
+   w -> panic ("toWordRange: Unknown platformWordSize: " ++ show w)
+
 -- | Creates a 'Literal' of type @Word#@.
 --   If the argument is out of the (target-dependent) range, it is wrapped.
 --   See Note [Word/Int underflow/overflow]
 mkMachWordWrap :: DynFlags -> Integer -> Literal
-mkMachWordWrap dflags i
- = MachWord $ case platformWordSize (targetPlatform dflags) of
-   4 -> toInteger (fromInteger i :: Word32)
-   8 -> toInteger (fromInteger i :: Word64)
-   w -> panic ("toWordRange: Unknown platformWordSize: " ++ show w)
+mkMachWordWrap dflags i = MachWord (wrapWord dflags i)
+
+-- | Creates a 'Literal' of type @Word#@, as well as a 'Bool'ean flag indicating
+--   carry. That is, if the argument is out of the (target-dependent) range
+--   the argument is wrapped and the carry flag will be set.
+--   See Note [Word/Int underflow/overflow]
+mkMachWordWrapC :: DynFlags -> Integer -> (Literal, Bool)
+mkMachWordWrapC dflags i = (MachWord i', i /= i')
+  where
+    i' = wrapWord dflags i
 
 -- | Creates a 'Literal' of type @Int64#@
 mkMachInt64 :: Integer -> Literal
diff --git a/compiler/prelude/PrelRules.hs b/compiler/prelude/PrelRules.hs
index 9fa0db6253..d0ad6c5dd1 100644
--- a/compiler/prelude/PrelRules.hs
+++ b/compiler/prelude/PrelRules.hs
@@ -94,6 +94,11 @@ primOpRules nm IntAddOp    = mkPrimOpRule nm 2 [ binaryLit (intOp2 (+))
 primOpRules nm IntSubOp    = mkPrimOpRule nm 2 [ binaryLit (intOp2 (-))
                                                , rightIdentityDynFlags zeroi
                                                , equalArgs >> retLit zeroi ]
+primOpRules nm IntAddCOp   = mkPrimOpRule nm 2 [ binaryLit (intOpC2 (+))
+                                               , identityCDynFlags zeroi ]
+primOpRules nm IntSubCOp   = mkPrimOpRule nm 2 [ binaryLit (intOpC2 (-))
+                                               , rightIdentityCDynFlags zeroi
+                                               , equalArgs >> retLitNoC zeroi ]
 primOpRules nm IntMulOp    = mkPrimOpRule nm 2 [ binaryLit (intOp2 (*))
                                                , zeroElem zeroi
                                                , identityDynFlags onei ]
@@ -135,6 +140,11 @@ primOpRules nm WordAddOp   = mkPrimOpRule nm 2 [ binaryLit (wordOp2 (+))
 primOpRules nm WordSubOp   = mkPrimOpRule nm 2 [ binaryLit (wordOp2 (-))
                                                , rightIdentityDynFlags zerow
                                                , equalArgs >> retLit zerow ]
+primOpRules nm WordAddCOp  = mkPrimOpRule nm 2 [ binaryLit (wordOpC2 (+))
+                                               , identityCDynFlags zerow ]
+primOpRules nm WordSubCOp  = mkPrimOpRule nm 2 [ binaryLit (wordOpC2 (-))
+                                               , rightIdentityCDynFlags zerow
+                                               , equalArgs >> retLitNoC zerow ]
 primOpRules nm WordMulOp   = mkPrimOpRule nm 2 [ binaryLit (wordOp2 (*))
                                                , identityDynFlags onew ]
 primOpRules nm WordQuotOp  = mkPrimOpRule nm 2 [ nonZeroLit 1 >> binaryLit (wordOp2 quot)
@@ -398,6 +408,13 @@ intOp2' op dflags (MachInt i1) (MachInt i2) =
   in  intResult dflags (fromInteger i1 `o` fromInteger i2)
 intOp2' _  _      _            _            = Nothing  -- Could find LitLit
 
+intOpC2 :: (Integral a, Integral b)
+        => (a -> b -> Integer)
+        -> DynFlags -> Literal -> Literal -> Maybe CoreExpr
+intOpC2 op dflags (MachInt i1) (MachInt i2) = do
+  intCResult dflags (fromInteger i1 `op` fromInteger i2)
+intOpC2 _  _      _            _            = Nothing  -- Could find LitLit
+
 shiftRightLogical :: DynFlags -> Integer -> Int -> Integer
 -- Shift right, putting zeros in rather than sign-propagating as Bits.shiftR would do
 -- Do this by converting to Word and back.  Obviously this won't work for big
@@ -412,6 +429,12 @@ retLit :: (DynFlags -> Literal) -> RuleM CoreExpr
 retLit l = do dflags <- getDynFlags
               return $ Lit $ l dflags
 
+retLitNoC :: (DynFlags -> Literal) -> RuleM CoreExpr
+retLitNoC l = do dflags <- getDynFlags
+                 let lit = l dflags
+                 let ty = literalType lit
+                 return $ mkCoreUbxTup [ty, ty] [Lit lit, Lit (zeroi dflags)]
+
 wordOp2 :: (Integral a, Integral b)
         => (a -> b -> Integer)
         -> DynFlags -> Literal -> Literal -> Maybe CoreExpr
@@ -419,6 +442,13 @@ wordOp2 op dflags (MachWord w1) (MachWord w2)
     = wordResult dflags (fromInteger w1 `op` fromInteger w2)
 wordOp2 _ _ _ _ = Nothing  -- Could find LitLit
 
+wordOpC2 :: (Integral a, Integral b)
+        => (a -> b -> Integer)
+        -> DynFlags -> Literal -> Literal -> Maybe CoreExpr
+wordOpC2 op dflags (MachWord w1) (MachWord w2) =
+  wordCResult dflags (fromInteger w1 `op` fromInteger w2)
+wordOpC2 _ _ _ _ = Nothing  -- Could find LitLit
+
 shiftRule :: (DynFlags -> Integer -> Int -> Integer) -> RuleM CoreExpr
                  -- Shifts take an Int; hence third arg of op is Int
 -- See Note [Guarding against silly shifts]
@@ -550,11 +580,31 @@ isMaxBound _      _              = False
 intResult :: DynFlags -> Integer -> Maybe CoreExpr
 intResult dflags result = Just (Lit (mkMachIntWrap dflags result))
 
+-- | Create an unboxed pair of an Int literal expression, ensuring the given
+-- Integer is in the target Int range and the corresponding overflow flag
+-- (@0#@/@1#@) if it wasn't.
+intCResult :: DynFlags -> Integer -> Maybe CoreExpr
+intCResult dflags result = Just (mkPair [Lit lit, Lit c])
+  where
+    mkPair = mkCoreUbxTup [intPrimTy, intPrimTy]
+    (lit, b) = mkMachIntWrapC dflags result
+    c = if b then onei dflags else zeroi dflags
+
 -- | Create a Word literal expression while ensuring the given Integer is in the
 -- target Word range
 wordResult :: DynFlags -> Integer -> Maybe CoreExpr
 wordResult dflags result = Just (Lit (mkMachWordWrap dflags result))
 
+-- | Create an unboxed pair of a Word literal expression, ensuring the given
+-- Integer is in the target Word range and the corresponding carry flag
+-- (@0#@/@1#@) if it wasn't.
+wordCResult :: DynFlags -> Integer -> Maybe CoreExpr
+wordCResult dflags result = Just (mkPair [Lit lit, Lit c])
+  where
+    mkPair = mkCoreUbxTup [wordPrimTy, intPrimTy]
+    (lit, b) = mkMachWordWrapC dflags result
+    c = if b then onei dflags else zeroi dflags
+
 inversePrimOp :: PrimOp -> RuleM CoreExpr
 inversePrimOp primop = do
   [Var primop_id `App` e] <- getArgs
@@ -738,6 +788,16 @@ leftIdentityDynFlags id_lit = do
   guard $ l1 == id_lit dflags
   return e2
 
+-- | Left identity rule for PrimOps like 'IntAddC' and 'WordAddC', where, in
+-- addition to the result, we have to indicate that no carry/overflow occured.
+leftIdentityCDynFlags :: (DynFlags -> Literal) -> RuleM CoreExpr
+leftIdentityCDynFlags id_lit = do
+  dflags <- getDynFlags
+  [Lit l1, e2] <- getArgs
+  guard $ l1 == id_lit dflags
+  let no_c = Lit (zeroi dflags)
+  return (mkCoreUbxTup [exprType e2, intPrimTy] [e2, no_c])
+
 rightIdentityDynFlags :: (DynFlags -> Literal) -> RuleM CoreExpr
 rightIdentityDynFlags id_lit = do
   dflags <- getDynFlags
@@ -745,8 +805,25 @@ rightIdentityDynFlags id_lit = do
   guard $ l2 == id_lit dflags
   return e1
 
+-- | Right identity rule for PrimOps like 'IntSubC' and 'WordSubC', where, in
+-- addition to the result, we have to indicate that no carry/overflow occured.
+rightIdentityCDynFlags :: (DynFlags -> Literal) -> RuleM CoreExpr
+rightIdentityCDynFlags id_lit = do
+  dflags <- getDynFlags
+  [e1, Lit l2] <- getArgs
+  guard $ l2 == id_lit dflags
+  let no_c = Lit (zeroi dflags)
+  return (mkCoreUbxTup [exprType e1, intPrimTy] [e1, no_c])
+
 identityDynFlags :: (DynFlags -> Literal) -> RuleM CoreExpr
-identityDynFlags lit = leftIdentityDynFlags lit `mplus` rightIdentityDynFlags lit
+identityDynFlags lit =
+  leftIdentityDynFlags lit `mplus` rightIdentityDynFlags lit
+
+-- | Identity rule for PrimOps like 'IntAddC' and 'WordAddC', where, in addition
+-- to the result, we have to indicate that no carry/overflow occured.
+identityCDynFlags :: (DynFlags -> Literal) -> RuleM CoreExpr
+identityCDynFlags lit =
+  leftIdentityCDynFlags lit `mplus` rightIdentityCDynFlags lit
 
 leftZero :: (DynFlags -> Literal) -> RuleM CoreExpr
 leftZero zero = do
diff --git a/testsuite/tests/numeric/should_run/T10962.hs b/testsuite/tests/numeric/should_run/T10962.hs
index 896c9e987f..435f3637d7 100644
--- a/testsuite/tests/numeric/should_run/T10962.hs
+++ b/testsuite/tests/numeric/should_run/T10962.hs
@@ -1,16 +1,32 @@
 {-# LANGUAGE MagicHash     #-}
 {-# LANGUAGE UnboxedTuples #-}
 
-module Main where
+module Main (main) where
 
 import GHC.Base
 
+unW# :: Word -> Word#
+unW# (W# w) = w
+
+type WordOpC = Word# -> Word# -> (# Word#, Int# #)
+
+check :: WordOpC -> Word# -> Word# -> IO ()
+check op a b = do
+  let (# w, c #) = op a b
+  print (W# w, I# c)
+
+checkSubInlNoInl :: WordOpC -> Word# -> Word# -> IO ()
+checkSubInlNoInl op a b = do
+  inline check op a b   -- constant folding
+  noinline check op a b -- lowering of PrimOp
+{-# INLINE checkSubInlNoInl #-}
+
 main :: IO ()
 main = do
   -- Overflow.
-  let (# w1, i1 #) = subWordC# 1## 3##
-  print (W# w1, I# i1)
+  checkSubInlNoInl subWordC# 1## 3##
+  checkSubInlNoInl addWordC# (unW# (inline maxBound)) 3##
 
   -- No overflow.
-  let (# w2, i2 #) = subWordC# 3## 1##
-  print (W# w2, I# i2)
+  checkSubInlNoInl subWordC# 5## 2##
+  checkSubInlNoInl addWordC# (unW# (inline maxBound-1)) 1##
diff --git a/testsuite/tests/numeric/should_run/T10962.stdout-ws-32 b/testsuite/tests/numeric/should_run/T10962.stdout-ws-32
index a1dec8410a..605265305d 100644
--- a/testsuite/tests/numeric/should_run/T10962.stdout-ws-32
+++ b/testsuite/tests/numeric/should_run/T10962.stdout-ws-32
@@ -1,2 +1,8 @@
 (4294967294,1)
-(2,0)
+(4294967294,1)
+(2,1)
+(2,1)
+(3,0)
+(3,0)
+(4294967295,0)
+(4294967295,0)
diff --git a/testsuite/tests/numeric/should_run/T10962.stdout-ws-64 b/testsuite/tests/numeric/should_run/T10962.stdout-ws-64
index 853bf94a61..d36f660eb1 100644
--- a/testsuite/tests/numeric/should_run/T10962.stdout-ws-64
+++ b/testsuite/tests/numeric/should_run/T10962.stdout-ws-64
@@ -1,2 +1,8 @@
 (18446744073709551614,1)
-(2,0)
+(18446744073709551614,1)
+(2,1)
+(2,1)
+(3,0)
+(3,0)
+(18446744073709551615,0)
+(18446744073709551615,0)
diff --git a/testsuite/tests/numeric/should_run/all.T b/testsuite/tests/numeric/should_run/all.T
index 37fff44bde..691fc26f7a 100644
--- a/testsuite/tests/numeric/should_run/all.T
+++ b/testsuite/tests/numeric/should_run/all.T
@@ -62,6 +62,6 @@ test('CarryOverflow', omit_ways(['ghci']), compile_and_run, [''])
 test('T9407', normal, compile_and_run, [''])
 test('T9810', normal, compile_and_run, [''])
 test('T10011', normal, compile_and_run, [''])
-test('T10962', omit_ways(['ghci']), compile_and_run, [''])
+test('T10962', omit_ways(['ghci']), compile_and_run, ['-O2'])
 test('T11702', extra_ways(['optasm']), compile_and_run, [''])
 test('T12136', normal, compile_and_run, [''])