Allow SpecConstr to specialize for unboxed literals.wip/andreask/spec_lits

Fixes #22781

4 files changed, 94 insertions, 3 deletions

diff --git a/compiler/GHC/Core/Opt/SpecConstr.hs b/compiler/GHC/Core/Opt/SpecConstr.hs
index 81c2816334..67d85ce4e2 100644
--- a/compiler/GHC/Core/Opt/SpecConstr.hs
+++ b/compiler/GHC/Core/Opt/SpecConstr.hs
@@ -1530,10 +1530,15 @@ scExpr' env (Case scrut b ty alts)
           ; (usg, rhs') <- scExpr env2 rhs
           ; let (usg', b_occ:arg_occs) = lookupOccs usg (b':bs2)
                 scrut_occ = case con of
-                               DataAlt dc -- See Note [Do not specialise evals]
+                              DataAlt dc -- See Note [Do not specialise evals]
                                   | not (single_alt && all deadArgOcc arg_occs)
-                                  -> ScrutOcc (unitUFM dc arg_occs)
-                               _  -> UnkOcc
+                                  -> -- pprTrace "sc_alt1" (ppr b' $$ ppr con $$ ppr bs $$ ppr arg_occs) $
+                                     ScrutOcc (unitUFM dc arg_occs)
+                              LitAlt _
+                                  | not single_alt
+                                  -> ScrutOcc (emptyUFM)
+                              _   -> -- pprTrace "sc_alt1" (ppr b' $$ ppr con $$ ppr bs $$ ppr arg_occs) $
+                                     UnkOcc
           ; return (usg', b_occ `combineOcc` scrut_occ, Alt con bs2 rhs') }
 
 
@@ -2631,6 +2636,11 @@ argToPat in_scope val_env arg arg_occ
   -- Check for a constructor application
   -- NB: this *precedes* the Var case, so that we catch nullary constrs
 argToPat1 env in_scope val_env arg arg_occ _arg_str
+  | Just (ConVal (LitAlt lit) _args) <- isValue val_env arg
+  , mb_scrut_lit
+  = do {
+       ; return (True, Lit lit , []) }
+
   | Just (ConVal (DataAlt dc) args) <- isValue val_env arg
   , not (ignoreDataCon env dc)        -- See Note [NoSpecConstr]
   , Just arg_occs <- mb_scrut dc
@@ -2660,6 +2670,13 @@ argToPat1 env in_scope val_env arg arg_occ _arg_str
                             -> Just (repeat UnkOcc)
                             | otherwise
                             -> Nothing
+    mb_scrut_lit = case arg_occ of
+                ScrutOcc _  -> True
+                _other      | sc_force env || sc_keen (sc_opts env)
+                            -> True
+                            | otherwise
+                            -> False
+
     match_vals bangs (arg:args)
       | isTypeArg arg
       = NotMarkedStrict : match_vals bangs args
diff --git a/testsuite/tests/simplCore/should_compile/T22781.hs b/testsuite/tests/simplCore/should_compile/T22781.hs
new file mode 100644
index 0000000000..125c50cb79
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22781.hs
@@ -0,0 +1,16 @@
+{-# LANGUAGE MagicHash #-}
+
+module T22781 where
+
+import GHC.Exts
+
+bar = I# (go 0# 1#)
+  where
+    -- SpecConstr should generate a specialization for the call
+    --    (go 0# 1#) = $sgo
+    -- looking like
+    --    $sgo = / void -> 1#
+    go :: Int# -> Int# -> Int#
+    go 0# 1# = 1#
+    go _  0# = 3#
+    go n x = go n (x -# 1# )
diff --git a/testsuite/tests/simplCore/should_compile/T22781.stderr b/testsuite/tests/simplCore/should_compile/T22781.stderr
new file mode 100644
index 0000000000..4c813ccc3c
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22781.stderr
@@ -0,0 +1,56 @@
+
+==================== Specialise ====================
+Result size of Specialise
+  = {terms: 37, types: 10, coercions: 0, joins: 0/0}
+
+Rec {
+go
+  = \ ds ds ->
+      case ds of ds {
+        __DEFAULT ->
+          case ds of ds {
+            __DEFAULT -> go ds (-# ds 1#);
+            0# -> 3#
+          };
+        0# ->
+          case ds of ds {
+            __DEFAULT -> go 0# (-# ds 1#);
+            0# -> 3#;
+            1# -> 1#
+          }
+      }
+end Rec }
+
+bar = case go 0# 1# of ds { __DEFAULT -> I# ds }
+
+
+
+
+==================== SpecConstr ====================
+Result size of SpecConstr
+  = {terms: 40, types: 13, coercions: 0, joins: 0/0}
+
+Rec {
+$sgo = \ void -> 1#
+
+go
+  = \ ds ds ->
+      case ds of ds {
+        __DEFAULT ->
+          case ds of ds {
+            __DEFAULT -> go ds (-# ds 1#);
+            0# -> 3#
+          };
+        0# ->
+          case ds of ds {
+            __DEFAULT -> go 0# (-# ds 1#);
+            0# -> 3#;
+            1# -> 1#
+          }
+      }
+end Rec }
+
+bar = case go 0# 1# of ds { __DEFAULT -> I# ds }
+
+
+
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 12e2ecf042..e174ae1f64 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -474,3 +474,5 @@ test('T15205', normal, compile, ['-O -ddump-simpl -dno-typeable-binds -dsuppress
 test('T22761', normal, multimod_compile, ['T22761', '-O2 -v0'])
 
 test('RewriteHigherOrderPatterns', normal, compile, ['-O -ddump-rule-rewrites -dsuppress-all -dsuppress-uniques'])
+# SpecConstr firing on literals
+test('T22781', [grep_errmsg(r'.*go') ], compile, ['-O2 -ddump-spec -dsuppress-uniques -dsuppress-all -dno-typeable-binds'])