CSE: Walk past join point lambdas (#15002)

As the CSE transformation traverses the syntax tree, it needs to go past
the lambdas of a join point, and only look for CSE opportunities inside,
as a join point’s lambdas must be preserved. Simple fix; comes with a
Note and a test case.

Thanks to Ryan Scott for an excellently minimized test case, and for
bisecting GHC.

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

1 files changed, 32 insertions, 2 deletions

diff --git a/compiler/simplCore/CSE.hs b/compiler/simplCore/CSE.hs
index 8f61128038..ee3a1eb538 100644
--- a/compiler/simplCore/CSE.hs
+++ b/compiler/simplCore/CSE.hs
@@ -17,7 +17,7 @@ import Var              ( Var )
 import VarEnv           ( elemInScopeSet, mkInScopeSet )
 import Id               ( Id, idType, idInlineActivation, isDeadBinder
                         , zapIdOccInfo, zapIdUsageInfo, idInlinePragma
-                        , isJoinId )
+                        , isJoinId, isJoinId_maybe )
 import CoreUtils        ( mkAltExpr, eqExpr
                         , exprIsTickedString
                         , stripTicksE, stripTicksT, mkTicks )
@@ -274,7 +274,28 @@ compiling ppHtml in Haddock.Backends.Xhtml).
 
 We could try and be careful by tracking which join points are still valid at
 each subexpression, but since join points aren't allocated or shared, there's
-less to gain by trying to CSE them.
+less to gain by trying to CSE them. (#13219)
+
+Note [Don’t tryForCSE the RHS of a Join Point]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Another way how CSE for joint points is tricky is
+
+  let join foo x = (x, 42)
+      join bar x = (x, 42)
+  in … jump foo 1 … jump bar 2 …
+
+naively, CSE would turn this into
+
+  let join foo x = (x, 42)
+      join bar = foo
+  in … jump foo 1 … jump bar 2 …
+
+but now bar is a join point that claims arity one, but its right-hand side
+is not a lambda, breaking the join-point invariant (this was #15002).
+
+Therefore, `cse_bind` will zoom past the lambdas of a join point (using
+`collectNBinders`) and resume searching for CSE opportunities only in the body
+of the join point.
 
 Note [CSE for recursive bindings]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -353,6 +374,13 @@ cse_bind toplevel env (in_id, in_rhs) out_id
       -- See Note [Take care with literal strings]
   = (env', (out_id, in_rhs))
 
+  | Just arity <- isJoinId_maybe in_id
+      -- See Note [Don’t tryForCSE the RHS of a Join Point]
+  = let (params, in_body) = collectNBinders arity in_rhs
+        (env', params') = addBinders env params
+        out_body = tryForCSE env' in_body
+    in (env, (out_id, mkLams params' out_body))
+
   | otherwise
   = (env', (out_id', out_rhs))
   where
@@ -392,6 +420,8 @@ addBinding env in_id out_id rhs'
                    Var {} -> True
                    _      -> False
 
+-- | Given a binder `let x = e`, this function
+-- determines whether we should add `e -> x` to the cs_map
 noCSE :: InId -> Bool
 noCSE id =  not (isAlwaysActive (idInlineActivation id)) &&
             not (noUserInlineSpec (inlinePragmaSpec (idInlinePragma id)))