Add constant folding rule (#16402)

    narrowN (x .&. m)
    m .&. (2^N-1) = 2^N-1
    ==> narrowN x

e.g.  narrow16 (x .&. 0x12FFFF) ==> narrow16 x

4 files changed, 88 insertions, 6 deletions

diff --git a/compiler/prelude/PrelRules.hs b/compiler/prelude/PrelRules.hs
index bdcbecfe59..c6c27f8ffe 100644
--- a/compiler/prelude/PrelRules.hs
+++ b/compiler/prelude/PrelRules.hs
@@ -191,29 +191,35 @@ primOpRules nm Int2WordOp     = mkPrimOpRule nm 1 [ liftLitDynFlags int2WordLit
 primOpRules nm Narrow8IntOp   = mkPrimOpRule nm 1 [ liftLit narrow8IntLit
                                                   , subsumedByPrimOp Narrow8IntOp
                                                   , Narrow8IntOp `subsumesPrimOp` Narrow16IntOp
-                                                  , Narrow8IntOp `subsumesPrimOp` Narrow32IntOp ]
+                                                  , Narrow8IntOp `subsumesPrimOp` Narrow32IntOp
+                                                  , narrowSubsumesAnd AndIOp Narrow8IntOp 8 ]
 primOpRules nm Narrow16IntOp  = mkPrimOpRule nm 1 [ liftLit narrow16IntLit
                                                   , subsumedByPrimOp Narrow8IntOp
                                                   , subsumedByPrimOp Narrow16IntOp
-                                                  , Narrow16IntOp `subsumesPrimOp` Narrow32IntOp ]
+                                                  , Narrow16IntOp `subsumesPrimOp` Narrow32IntOp
+                                                  , narrowSubsumesAnd AndIOp Narrow16IntOp 16 ]
 primOpRules nm Narrow32IntOp  = mkPrimOpRule nm 1 [ liftLit narrow32IntLit
                                                   , subsumedByPrimOp Narrow8IntOp
                                                   , subsumedByPrimOp Narrow16IntOp
                                                   , subsumedByPrimOp Narrow32IntOp
-                                                  , removeOp32 ]
+                                                  , removeOp32
+                                                  , narrowSubsumesAnd AndIOp Narrow32IntOp 32 ]
 primOpRules nm Narrow8WordOp  = mkPrimOpRule nm 1 [ liftLit narrow8WordLit
                                                   , subsumedByPrimOp Narrow8WordOp
                                                   , Narrow8WordOp `subsumesPrimOp` Narrow16WordOp
-                                                  , Narrow8WordOp `subsumesPrimOp` Narrow32WordOp ]
+                                                  , Narrow8WordOp `subsumesPrimOp` Narrow32WordOp
+                                                  , narrowSubsumesAnd AndOp Narrow8WordOp 8 ]
 primOpRules nm Narrow16WordOp = mkPrimOpRule nm 1 [ liftLit narrow16WordLit
                                                   , subsumedByPrimOp Narrow8WordOp
                                                   , subsumedByPrimOp Narrow16WordOp
-                                                  , Narrow16WordOp `subsumesPrimOp` Narrow32WordOp ]
+                                                  , Narrow16WordOp `subsumesPrimOp` Narrow32WordOp
+                                                  , narrowSubsumesAnd AndOp Narrow16WordOp 16 ]
 primOpRules nm Narrow32WordOp = mkPrimOpRule nm 1 [ liftLit narrow32WordLit
                                                   , subsumedByPrimOp Narrow8WordOp
                                                   , subsumedByPrimOp Narrow16WordOp
                                                   , subsumedByPrimOp Narrow32WordOp
-                                                  , removeOp32 ]
+                                                  , removeOp32
+                                                  , narrowSubsumesAnd AndOp Narrow32WordOp 32 ]
 primOpRules nm OrdOp          = mkPrimOpRule nm 1 [ liftLit char2IntLit
                                                   , inversePrimOp ChrOp ]
 primOpRules nm ChrOp          = mkPrimOpRule nm 1 [ do [Lit lit] <- getArgs
@@ -649,6 +655,26 @@ subsumedByPrimOp primop = do
   matchPrimOpId primop primop_id
   return e
 
+-- | narrow subsumes bitwise `and` with full mask (cf #16402):
+--
+--       narrowN (x .&. m)
+--       m .&. (2^N-1) = 2^N-1
+--       ==> narrowN x
+--
+-- e.g.  narrow16 (x .&. 0xFFFF)
+--       ==> narrow16 x
+--
+narrowSubsumesAnd :: PrimOp -> PrimOp -> Int -> RuleM CoreExpr
+narrowSubsumesAnd and_primop narrw n = do
+  [Var primop_id `App` x `App` y] <- getArgs
+  matchPrimOpId and_primop primop_id
+  let mask = bit n -1
+      g v (Lit (LitNumber _ m _)) = do
+         guard (m .&. mask == mask)
+         return (Var (mkPrimOpId narrw) `App` v)
+      g _ _ = mzero
+  g x y <|> g y x
+
 idempotent :: RuleM CoreExpr
 idempotent = do [e1, e2] <- getArgs
                 guard $ cheapEqExpr e1 e2
diff --git a/testsuite/tests/numeric/should_compile/T16402.hs b/testsuite/tests/numeric/should_compile/T16402.hs
new file mode 100644
index 0000000000..c85fe0037b
--- /dev/null
+++ b/testsuite/tests/numeric/should_compile/T16402.hs
@@ -0,0 +1,19 @@
+{-# OPTIONS_GHC -ddump-simpl -dhex-word-literals -dsuppress-all -dsuppress-uniques -O2 #-}
+{-# LANGUAGE TypeApplications #-}
+module T16402 where
+
+import Data.Word
+import Data.Int
+import Data.Bits
+
+smallWord_foo :: Word64 -> Word64
+smallWord_foo x = fromIntegral @Word16 $ fromIntegral (x .&. 0xFFFF)
+
+smallWord_bar :: Word64 -> Word64
+smallWord_bar x = fromIntegral (fromIntegral x :: Word16)
+
+smallInt_foo :: Int64 -> Int64
+smallInt_foo x = fromIntegral @Int16 $ fromIntegral (x .&. 0x12FFFF)
+
+smallInt_bar :: Int64 -> Int64
+smallInt_bar x = fromIntegral (fromIntegral x :: Int16)
diff --git a/testsuite/tests/numeric/should_compile/T16402.stderr b/testsuite/tests/numeric/should_compile/T16402.stderr
new file mode 100644
index 0000000000..75db843376
--- /dev/null
+++ b/testsuite/tests/numeric/should_compile/T16402.stderr
@@ -0,0 +1,36 @@
+
+==================== Tidy Core ====================
+Result size of Tidy Core
+  = {terms: 34, types: 19, coercions: 0, joins: 0/0}
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+$trModule4 = "main"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+$trModule3 = TrNameS $trModule4
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+$trModule2 = "T16402"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+$trModule1 = TrNameS $trModule2
+
+-- RHS size: {terms: 3, types: 0, coercions: 0, joins: 0/0}
+$trModule = Module $trModule3 $trModule1
+
+-- RHS size: {terms: 7, types: 3, coercions: 0, joins: 0/0}
+smallWord_bar
+  = \ x -> case x of { W64# x# -> W64# (narrow16Word# x#) }
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+smallWord_foo = smallWord_bar
+
+-- RHS size: {terms: 7, types: 3, coercions: 0, joins: 0/0}
+smallInt_bar
+  = \ x -> case x of { I64# x# -> I64# (narrow16Int# x#) }
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+smallInt_foo = smallInt_bar
+
+
+
diff --git a/testsuite/tests/numeric/should_compile/all.T b/testsuite/tests/numeric/should_compile/all.T
index 625951f4d6..a7dc06cf44 100644
--- a/testsuite/tests/numeric/should_compile/all.T
+++ b/testsuite/tests/numeric/should_compile/all.T
@@ -9,3 +9,4 @@ test('T7881', normal, compile, [''])
 # desugaring, so we don't get the warning we expect.
 test('T8542', omit_ways(['hpc']), compile, [''])
 test('T10929', normal, compile, [''])
+test('T16402', [ grep_errmsg(r'and') ], compile, [''])