Add constant-folding rule for Data.Bits.bit

This adds a constant-folding rule for `Integer`'s implementation of `bit` and
fixes the `T8832` testcase. Fixes #8832.

Reviewed By: simonpj, austin

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

GHC Trac Issues: #8832

5 files changed, 50 insertions, 22 deletions

diff --git a/compiler/prelude/PrelNames.hs b/compiler/prelude/PrelNames.hs
index a6eb834641..be6396cf21 100644
--- a/compiler/prelude/PrelNames.hs
+++ b/compiler/prelude/PrelNames.hs
@@ -308,7 +308,7 @@ basicKnownKeyNames
         decodeDoubleIntegerName,
         gcdIntegerName, lcmIntegerName,
         andIntegerName, orIntegerName, xorIntegerName, complementIntegerName,
-        shiftLIntegerName, shiftRIntegerName,
+        shiftLIntegerName, shiftRIntegerName, bitIntegerName,
 
         -- Float/Double
         rationalToFloatName,
@@ -939,7 +939,7 @@ integerTyConName, mkIntegerName, integerSDataConName,
     decodeDoubleIntegerName,
     gcdIntegerName, lcmIntegerName,
     andIntegerName, orIntegerName, xorIntegerName, complementIntegerName,
-    shiftLIntegerName, shiftRIntegerName :: Name
+    shiftLIntegerName, shiftRIntegerName, bitIntegerName :: Name
 integerTyConName      = tcQual  gHC_INTEGER_TYPE (fsLit "Integer")           integerTyConKey
 integerSDataConName   = dcQual gHC_INTEGER_TYPE (fsLit n)                   integerSDataConKey
   where n = case cIntegerLibraryType of
@@ -986,6 +986,7 @@ xorIntegerName        = varQual gHC_INTEGER_TYPE (fsLit "xorInteger")        xor
 complementIntegerName = varQual gHC_INTEGER_TYPE (fsLit "complementInteger") complementIntegerIdKey
 shiftLIntegerName     = varQual gHC_INTEGER_TYPE (fsLit "shiftLInteger")     shiftLIntegerIdKey
 shiftRIntegerName     = varQual gHC_INTEGER_TYPE (fsLit "shiftRInteger")     shiftRIntegerIdKey
+bitIntegerName        = varQual gHC_INTEGER_TYPE (fsLit "bitInteger")        bitIntegerIdKey
 
 -- GHC.Real types and classes
 rationalTyConName, ratioTyConName, ratioDataConName, realClassName,
@@ -1901,6 +1902,9 @@ typeSymbolTypeRepKey  = mkPreludeMiscIdUnique 507
 toDynIdKey :: Unique
 toDynIdKey = mkPreludeMiscIdUnique 508
 
+bitIntegerIdKey :: Unique
+bitIntegerIdKey       = mkPreludeMiscIdUnique 509
+
 {-
 ************************************************************************
 *                                                                      *
diff --git a/compiler/prelude/PrelRules.hs b/compiler/prelude/PrelRules.hs
index 1ab8543afc..d44c224479 100644
--- a/compiler/prelude/PrelRules.hs
+++ b/compiler/prelude/PrelRules.hs
@@ -1003,6 +1003,7 @@ builtinIntegerRules =
   rule_unop           "complementInteger"   complementIntegerName   complement,
   rule_Int_binop      "shiftLInteger"       shiftLIntegerName       shiftL,
   rule_Int_binop      "shiftRInteger"       shiftRIntegerName       shiftR,
+  rule_bitInteger     "bitInteger"          bitIntegerName,
   -- See Note [Integer division constant folding] in libraries/base/GHC/Real.hs
   rule_divop_one      "quotInteger"         quotIntegerName         quot,
   rule_divop_one      "remInteger"          remIntegerName          rem,
@@ -1039,6 +1040,9 @@ builtinIntegerRules =
           rule_unop str name op
            = BuiltinRule { ru_name = fsLit str, ru_fn = name, ru_nargs = 1,
                            ru_try = match_Integer_unop op }
+          rule_bitInteger str name
+           = BuiltinRule { ru_name = fsLit str, ru_fn = name, ru_nargs = 1,
+                           ru_try = match_IntToInteger_unop (bit . fromIntegral) }
           rule_binop str name op
            = BuiltinRule { ru_name = fsLit str, ru_fn = name, ru_nargs = 2,
                            ru_try = match_Integer_binop op }
@@ -1155,14 +1159,7 @@ match_magicDict _ = Nothing
 -- Similarly Int64, Word64
 
 match_IntToInteger :: RuleFun
-match_IntToInteger _ id_unf fn [xl]
-  | Just (MachInt x) <- exprIsLiteral_maybe id_unf xl
-  = case idType fn of
-    FunTy _ integerTy ->
-        Just (Lit (LitInteger x integerTy))
-    _ ->
-        panic "match_IntToInteger: Id has the wrong type"
-match_IntToInteger _ _ _ _ = Nothing
+match_IntToInteger = match_IntToInteger_unop id
 
 match_WordToInteger :: RuleFun
 match_WordToInteger _ id_unf id [xl]
@@ -1209,6 +1206,32 @@ match_Integer_unop unop _ id_unf _ [xl]
   = Just (Lit (LitInteger (unop x) i))
 match_Integer_unop _ _ _ _ _ = Nothing
 
+{- Note [Rewriting bitInteger]
+
+For most types the bitInteger operation can be implemented in terms of shifts.
+The integer-gmp package, however, can do substantially better than this if
+allowed to provide its own implementation. However, in so doing it previously lost
+constant-folding (see Trac #8832). The bitInteger rule above provides constant folding
+specifically for this function.
+
+There is, however, a bit of trickiness here when it comes to ranges. While the
+AST encodes all integers (even MachInts) as Integers, `bit` expects the bit
+index to be given as an Int. Hence we coerce to an Int in the rule definition.
+This will behave a bit funny for constants larger than the word size, but the user
+should expect some funniness given that they will have at very least ignored a
+warning in this case.
+-}
+
+match_IntToInteger_unop :: (Integer -> Integer) -> RuleFun
+match_IntToInteger_unop unop _ id_unf fn [xl]
+  | Just (MachInt x) <- exprIsLiteral_maybe id_unf xl
+  = case idType fn of
+    FunTy _ integerTy ->
+        Just (Lit (LitInteger (unop x) integerTy))
+    _ ->
+        panic "match_IntToInteger_unop: Id has the wrong type"
+match_IntToInteger_unop _ _ _ _ _ = Nothing
+
 match_Integer_binop :: (Integer -> Integer -> Integer) -> RuleFun
 match_Integer_binop binop _ id_unf _ [xl,yl]
   | Just (LitInteger x i) <- exprIsLiteral_maybe id_unf xl
diff --git a/testsuite/tests/simplCore/should_compile/Makefile b/testsuite/tests/simplCore/should_compile/Makefile
index 7f43dafdc8..8c6ec45796 100644
--- a/testsuite/tests/simplCore/should_compile/Makefile
+++ b/testsuite/tests/simplCore/should_compile/Makefile
@@ -4,7 +4,7 @@ include $(TOP)/mk/test.mk
 
 T8832:
 	$(RM) -f T8832.o T8832.hi
-	'$(TEST_HC)' $(TEST_HC_OPTS) $(T8832_WORDSIZE_OPTS) -O -c -ddump-simpl T8832.hs | grep '#'
+	'$(TEST_HC)' $(TEST_HC_OPTS) $(T8832_WORDSIZE_OPTS) -O -c -ddump-simpl T8832.hs | grep '^[a-zA-Z0-9]\+ ='
 
 T8274:
 	$(RM) -f T8274.o T8274.hi
diff --git a/testsuite/tests/simplCore/should_compile/T8832.stdout b/testsuite/tests/simplCore/should_compile/T8832.stdout
index 9c10451669..a351735cd0 100644
--- a/testsuite/tests/simplCore/should_compile/T8832.stdout
+++ b/testsuite/tests/simplCore/should_compile/T8832.stdout
@@ -1,10 +1,11 @@
-i = GHC.Types.I# 0#
-i8 = GHC.Int.I8# 0#
-i16 = GHC.Int.I16# 0#
-i32 = GHC.Int.I32# 0#
-i64 = GHC.Int.I64# 0#
-w = GHC.Types.W# 0##
-w8 = GHC.Word.W8# 0##
-w16 = GHC.Word.W16# 0##
-w32 = GHC.Word.W32# 0##
-w64 = GHC.Word.W64# 0##
+i = I# 0#
+i8 = I8# 0#
+i16 = I16# 0#
+i32 = I32# 0#
+i64 = I64# 0#
+w = W# 0##
+w8 = W8# 0##
+w16 = W16# 0##
+w32 = W32# 0##
+w64 = W64# 0##
+z = 0
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index b337c9c170..c99b8f2bab 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -202,7 +202,7 @@ test('T5996',
      ['$MAKE -s --no-print-directory T5996'])
 test('T8537', normal, compile, [''])
 test('T8832',
-     expect_broken(8832),
+     normal,
      run_command,
      ['$MAKE -s --no-print-directory T8832 T8832_WORDSIZE_OPTS=' +
       ('-DT8832_WORDSIZE_64' if wordsize(64) else '')])