Don't eta-expand in stable unfoldings

See SimplUtils Note [No eta expansion in stable unfoldings],
and Trac #9509 for an excellend diagnosis by Nick Frisby

7 files changed, 64 insertions, 13 deletions

diff --git a/compiler/simplCore/SimplUtils.hs b/compiler/simplCore/SimplUtils.hs
index 6c4737507a..03adfe00fa 100644
--- a/compiler/simplCore/SimplUtils.hs
+++ b/compiler/simplCore/SimplUtils.hs
@@ -688,11 +688,12 @@ simplEnvForGHCi dflags
 updModeForStableUnfoldings :: Activation -> SimplifierMode -> SimplifierMode
 -- See Note [Simplifying inside stable unfoldings]
 updModeForStableUnfoldings inline_rule_act current_mode
-  = current_mode { sm_phase = phaseFromActivation inline_rule_act
-                 , sm_inline = True
+  = current_mode { sm_phase      = phaseFromActivation inline_rule_act
+                 , sm_inline     = True
                  , sm_eta_expand = False }
-                 -- For sm_rules, just inherit; sm_rules might be "off"
-                 -- because of -fno-enable-rewrite-rules
+                     -- sm_eta_expand: see Note [No eta expansion in stable unfoldings]
+       -- For sm_rules, just inherit; sm_rules might be "off"
+       -- because of -fno-enable-rewrite-rules
   where
     phaseFromActivation (ActiveAfter _ n) = Phase n
     phaseFromActivation _                 = InitialPhase
@@ -717,6 +718,25 @@ Ticks into the LHS, which makes matching trickier. Trac #10665, #10745.
 Doing this to either side confounds tools like HERMIT, which seek to reason
 about and apply the RULES as originally written. See Trac #10829.
 
+Note [No eta expansion in stable unfoldings]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+If we have a stable unfolding
+
+  f :: Ord a => a -> IO ()
+  -- Unfolding template
+  --    = /\a \(d:Ord a) (x:a). bla
+
+we do not want to eta-expand to
+
+  f :: Ord a => a -> IO ()
+  -- Unfolding template
+  --    = (/\a \(d:Ord a) (x:a) (eta:State#). bla eta) |> co
+
+because not specialisation of the overloading doesn't work properly
+(see Note [Specialisation shape] in Specialise), Trac #9509.
+
+So we disable eta-expansion in stable unfoldings.
+
 Note [Inlining in gentle mode]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Something is inlined if
@@ -1256,16 +1276,16 @@ won't inline because 'e' is too big.
 ************************************************************************
 -}
 
-mkLam :: [OutBndr] -> OutExpr -> SimplCont -> SimplM OutExpr
+mkLam :: SimplEnv -> [OutBndr] -> OutExpr -> SimplCont -> SimplM OutExpr
 -- mkLam tries three things
 --      a) eta reduction, if that gives a trivial expression
 --      b) eta expansion [only if there are some value lambdas]
 
-mkLam [] body _cont
+mkLam _env [] body _cont
   = return body
-mkLam bndrs body cont
-  = do  { dflags <- getDynFlags
-        ; mkLam' dflags bndrs body }
+mkLam env bndrs body cont
+  = do { dflags <- getDynFlags
+       ; mkLam' dflags bndrs body }
   where
     mkLam' :: DynFlags -> [OutBndr] -> OutExpr -> SimplM OutExpr
     mkLam' dflags bndrs (Cast body co)
@@ -1293,7 +1313,7 @@ mkLam bndrs body cont
            ; return etad_lam }
 
       | not (contIsRhs cont)   -- See Note [Eta-expanding lambdas]
-      , gopt Opt_DoLambdaEtaExpansion dflags
+      , sm_eta_expand (getMode env)
       , any isRuntimeVar bndrs
       , let body_arity = exprEtaExpandArity dflags body
       , body_arity > 0
@@ -1325,6 +1345,9 @@ better eta-expander (in the form of tryEtaExpandRhs), so we don't
 bother to try expansion in mkLam in that case; hence the contIsRhs
 guard.
 
+NB: We check the SimplEnv (sm_eta_expand), not DynFlags.
+    See Note [No eta expansion in stable unfoldings]
+
 Note [Casts and lambdas]
 ~~~~~~~~~~~~~~~~~~~~~~~~
 Consider
diff --git a/compiler/simplCore/Simplify.hs b/compiler/simplCore/Simplify.hs
index d52aacdde5..4f65b2b379 100644
--- a/compiler/simplCore/Simplify.hs
+++ b/compiler/simplCore/Simplify.hs
@@ -353,7 +353,7 @@ simplLazyBind env top_lvl is_rec bndr bndr1 rhs rhs_se
         ; (env', rhs')
             <-  if not (doFloatFromRhs top_lvl is_rec False body2 body_env2)
                 then                            -- No floating, revert to body1
-                     do { rhs' <- mkLam tvs' (wrapFloats body_env1 body1) rhs_cont
+                     do { rhs' <- mkLam env tvs' (wrapFloats body_env1 body1) rhs_cont
                         ; return (env, rhs') }
 
                 else if null tvs then           -- Simple floating
@@ -363,7 +363,7 @@ simplLazyBind env top_lvl is_rec bndr bndr1 rhs rhs_se
                 else                            -- Do type-abstraction first
                      do { tick LetFloatFromLet
                         ; (poly_binds, body3) <- abstractFloats tvs' body_env2 body2
-                        ; rhs' <- mkLam tvs' body3 rhs_cont
+                        ; rhs' <- mkLam env tvs' body3 rhs_cont
                         ; env' <- foldlM (addPolyBind top_lvl) env poly_binds
                         ; return (env', rhs') }
 
@@ -1272,7 +1272,7 @@ simplLam env bndrs body (TickIt tickish cont)
 simplLam env bndrs body cont
   = do  { (env', bndrs') <- simplLamBndrs env bndrs
         ; body' <- simplExpr env' body
-        ; new_lam <- mkLam bndrs' body' cont
+        ; new_lam <- mkLam env bndrs' body' cont
         ; rebuild env' new_lam cont }
 
 simplLamBndrs :: SimplEnv -> [InBndr] -> SimplM (SimplEnv, [OutBndr])
diff --git a/testsuite/tests/simplCore/should_compile/Makefile b/testsuite/tests/simplCore/should_compile/Makefile
index 288e3f96e5..3b1c2a52b5 100644
--- a/testsuite/tests/simplCore/should_compile/Makefile
+++ b/testsuite/tests/simplCore/should_compile/Makefile
@@ -7,6 +7,14 @@ T3990:
 	'$(TEST_HC)' $(TEST_HC_OPTS) -O -c -ddump-simpl T3990.hs | grep 'test_case'
         # Grep output should show an unpacked constructor
 
+T9509:
+	$(RM) -f T9509*.o T9509*.hi
+	'$(TEST_HC)' $(TEST_HC_OPTS) -O -c T9509a.hs
+	'$(TEST_HC)' $(TEST_HC_OPTS) -O -c T9509.hs -funfolding-use-threshold=20 \
+              -ddump-rule-rewrites | grep SPEC
+        # Grep output should show a SPEC rule firing
+        # The unfolding use threshold is to prevent foo inlining before it is specialised
+
 T8832:
 	$(RM) -f T8832.o T8832.hi
 	'$(TEST_HC)' $(TEST_HC_OPTS) $(T8832_WORDSIZE_OPTS) -O -c -ddump-simpl T8832.hs | grep '^[a-zA-Z0-9]\+ ='
diff --git a/testsuite/tests/simplCore/should_compile/T9509.hs b/testsuite/tests/simplCore/should_compile/T9509.hs
new file mode 100644
index 0000000000..86d2ce101f
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T9509.hs
@@ -0,0 +1,5 @@
+module T9509 (main) where
+
+import T9509a
+
+main = foo (5 :: Int) >>= print
diff --git a/testsuite/tests/simplCore/should_compile/T9509.stdout b/testsuite/tests/simplCore/should_compile/T9509.stdout
new file mode 100644
index 0000000000..0acd484558
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T9509.stdout
@@ -0,0 +1 @@
+    Rule: SPEC/T9509 foo @ Int
diff --git a/testsuite/tests/simplCore/should_compile/T9509a.hs b/testsuite/tests/simplCore/should_compile/T9509a.hs
new file mode 100644
index 0000000000..bd6511eedb
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T9509a.hs
@@ -0,0 +1,10 @@
+module T9509a (foo) where
+
+import Data.IORef
+
+foo :: Ord a => a -> IO a
+{-# INLINABLE foo #-}
+foo x = newIORef x >>= readIORef >>= \y ->
+        case compare x y of
+           LT ->  return x ;
+           _  -> return y
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 19d806f21c..459aa47eb7 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -246,3 +246,7 @@ test('T12212', normal, compile, ['-O'])
 test('noinline01', only_ways(['optasm']), compile, ['-ddump-stg -dsuppress-uniques -O'])
 test('par01', only_ways(['optasm']), compile, ['-ddump-prep -dsuppress-uniques -O2'])
 test('T12776', normal, compile, ['-O2'])
+test('T9509',
+     normal,
+     run_command,
+     ['$MAKE -s --no-print-directory T9509'])