Fix a nasty continuation-duplication bug

For a long-time mkDupableCont has had a bug that allows it to duplicate
an arbitrary continuation, which it should not do, of course.

The bug was that in the Select case of mkDupableCont we were calling
prepareCaseCont, which did not duplicate the continuation if there is
but a single alternative.  This is quite right in the case of the call
in rebuildCase, but quite wrong in mkDupableCont.

The bug manifest as follows. In the expression
	f (case ... of { ..several alts.. })
(when f is strict), we should transform to
	f (...transformed arg...)
The application of f should not be pushed down (see notes with the
ArgOf case of mkDupableCont.  But that was not happening in an example
like this (see how the call to f is pushed inwards).

f (a `div` abs (b::Int))
	--->
    case b_afT of wild_aHa { GHC.Base.I# x_aHc ->
    let {
      $j_sIe :: GHC.Prim.Int# -> GHC.Base.Int
      []
      $j_sIe =
	\ (ds1_aHr [Nothing OneShot] :: GHC.Prim.Int#) ->
	  Foo7.f
	    (case ds1_aHr of ds2_aHq {
	       __DEFAULT ->
		 case a_afS of wild1_aHM { GHC.Base.I# x_aHO ->
		 GHC.Base.I# (GHC.Base.divInt# x_aHO ds2_aHq)
		 };
	       0 -> GHC.Err.divZeroError @ GHC.Base.Int
	     })
    } in
      case GHC.Prim.>=# x_aHc 0 of wild1_aHe [Dead Nothing] {
	GHC.Base.False ->
	  let {
	    ds1_aHr :: GHC.Prim.Int#
	    []
	    ds1_aHr = GHC.Prim.negateInt# x_aHc
	  } in  $j_sIe ds1_aHr;
	GHC.Base.True -> $j_sIe x_aHc
      }
    }

diff --git a/compiler/simplCore/Simplify.lhs b/compiler/simplCore/Simplify.lhs
index 6f2e887..dd2a22b 100644 (file)
--- a/compiler/simplCore/Simplify.lhs
+++ b/compiler/simplCore/Simplify.lhs
@@ -1786,7 +1786,8 @@ prepareCaseCont :: SimplEnv
                -> [InAlt] -> SimplCont
                -> SimplM (FloatsWith (SimplCont,SimplCont))    
                        -- Return a duplicatable continuation, a non-duplicable part 
-                       -- plus some extra bindings
+                       -- plus some extra bindings (that scope over the entire
+                       -- continunation)
 
        -- No need to make it duplicatable if there's only one alternative
 prepareCaseCont env [alt] cont = returnSmpl (emptyFloats env, (cont, mkBoringStop (contResultType cont)))
@@ -1842,56 +1843,56 @@ mkDupableCont env (ApplyTo _ arg se cont)
        --      ==>
        --              let a = ...arg... 
        --              in [...hole...] a
-    simplExpr (setInScope se env) arg                  `thenSmpl` \ arg' ->
-
-    mkDupableCont env cont                             `thenSmpl` \ (floats, (dup_cont, nondup_cont)) ->
-    addFloats env floats                               $ \ env ->
-
-    if exprIsDupable arg' then
-       returnSmpl (emptyFloats env, (ApplyTo OkToDup arg' (zapSubstEnv se) dup_cont, nondup_cont))
-    else
-    newId FSLIT("a") (exprType arg')                   `thenSmpl` \ arg_id ->
-
-    tick (CaseOfCase arg_id)                           `thenSmpl_`
-       -- Want to tick here so that we go round again,
-       -- and maybe copy or inline the code.
-       -- Not strictly CaseOfCase, but never mind
-
-    returnSmpl (unitFloat env arg_id arg', 
-               (ApplyTo OkToDup (Var arg_id) (zapSubstEnv se) dup_cont,
-                nondup_cont))
-       -- But what if the arg should be case-bound? 
-       -- This has been this way for a long time, so I'll leave it,
-       -- but I can't convince myself that it's right.
+    do { (floats, (dup_cont, nondup_cont)) <- mkDupableCont env cont
+       ; addFloats env floats $ \ env -> do
+       { arg1 <- simplExpr (setInScope se env) arg
+       ; (floats2, arg2) <- mkDupableArg env arg1
+       ; return (floats2, (ApplyTo OkToDup arg2 (zapSubstEnv se) dup_cont, nondup_cont)) }}
 
 mkDupableCont env (Select _ case_bndr alts se cont)
   =    -- e.g.         (case [...hole...] of { pi -> ei })
        --      ===>
        --              let ji = \xij -> ei 
        --              in case [...hole...] of { pi -> ji xij }
-    tick (CaseOfCase case_bndr)                                        `thenSmpl_`
-    let
-       alt_env = setInScope se env
-    in
-    prepareCaseCont alt_env alts cont                          `thenSmpl` \ (floats1, (dup_cont, nondup_cont)) ->
-    addFloats alt_env floats1                                  $ \ alt_env ->
-
-    simplBinder alt_env case_bndr                              `thenSmpl` \ (alt_env, case_bndr') ->
-       -- NB: simplBinder does not zap deadness occ-info, so
-       -- a dead case_bndr' will still advertise its deadness
-       -- This is really important because in
-       --      case e of b { (# a,b #) -> ... }
-       -- b is always dead, and indeed we are not allowed to bind b to (# a,b #),
-       -- which might happen if e was an explicit unboxed pair and b wasn't marked dead.
-       -- In the new alts we build, we have the new case binder, so it must retain
-       -- its deadness.
-
-    mkDupableAlts alt_env case_bndr' alts dup_cont     `thenSmpl` \ (floats2, alts') ->
-    addFloats alt_env floats2                          $ \ alt_env ->
-    returnSmpl (emptyFloats alt_env, 
-               (Select OkToDup case_bndr' alts' (zapSubstEnv se) 
-                       (mkBoringStop (contResultType dup_cont)),
-                nondup_cont))
+    do { tick (CaseOfCase case_bndr)
+       ; let alt_env = setInScope se env
+       ; (floats1, (dup_cont, nondup_cont)) <- mkDupableCont alt_env cont
+               -- NB: call mkDupableCont here, *not* prepareCaseCont
+               -- We must make a duplicable continuation, whereas prepareCaseCont
+               -- doesn't when there is a single case branch
+       ; addFloats alt_env floats1     $ \ alt_env -> do
+
+       { (alt_env, case_bndr') <- simplBinder alt_env case_bndr
+               -- NB: simplBinder does not zap deadness occ-info, so
+               -- a dead case_bndr' will still advertise its deadness
+               -- This is really important because in
+               --      case e of b { (# a,b #) -> ... }
+               -- b is always dead, and indeed we are not allowed to bind b to (# a,b #),
+               -- which might happen if e was an explicit unboxed pair and b wasn't marked dead.
+               -- In the new alts we build, we have the new case binder, so it must retain
+               -- its deadness.
+
+       ; (floats2, alts') <- mkDupableAlts alt_env case_bndr' alts dup_cont
+       ; return (floats2, (Select OkToDup case_bndr' alts' (zapSubstEnv se) 
+                                  (mkBoringStop (contResultType dup_cont)),
+                           nondup_cont))
+       }}
+
+mkDupableArg :: SimplEnv -> OutExpr -> SimplM (FloatsWith OutExpr)
+-- Let-bind the thing if necessary
+mkDupableArg env arg
+  | exprIsDupable arg 
+  = return (emptyFloats env, arg)
+  | otherwise     
+  = do { arg_id <- newId FSLIT("a") (exprType arg)
+       ; tick (CaseOfCase arg_id)
+               -- Want to tick here so that we go round again,
+               -- and maybe copy or inline the code.
+               -- Not strictly CaseOfCase, but never mind
+       ; return (unitFloat env arg_id arg, Var arg_id) }
+       -- What if the arg should be case-bound? 
+       -- This has been this way for a long time, so I'll leave it,
+       -- but I can't convince myself that it's right.
 
 mkDupableAlts :: SimplEnv -> OutId -> [InAlt] -> SimplCont
              -> SimplM (FloatsWith [InAlt])