Fix in-scope set in simplifier

This patch fixes Trac #14408.  The problem was that the StaticEnv
field of an ApplyToVar or Select continuation didn't have enough
variables in scope.  The fix is simple, and I documented the
invariant in Note [StaticEnv invariant] in SimplUtils.

No change in behaviour: this just stops an ASSERT from tripping.

3 files changed, 56 insertions, 25 deletions

diff --git a/compiler/simplCore/SimplEnv.hs b/compiler/simplCore/SimplEnv.hs
index d2cb689a22..7f7977dd36 100644
--- a/compiler/simplCore/SimplEnv.hs
+++ b/compiler/simplCore/SimplEnv.hs
@@ -11,7 +11,7 @@ module SimplEnv (
         setMode, getMode, updMode, seDynFlags,
 
         -- * Environments
-        SimplEnv(..), StaticEnv, pprSimplEnv,   -- Temp not abstract
+        SimplEnv(..), pprSimplEnv,   -- Temp not abstract
         mkSimplEnv, extendIdSubst,
         SimplEnv.extendTvSubst, SimplEnv.extendCvSubst,
         zapSubstEnv, setSubstEnv,
@@ -101,8 +101,6 @@ data SimplEnv
       , seInScope   :: InScopeSet       -- OutVars only
     }
 
-type StaticEnv = SimplEnv       -- Just the static part is relevant
-
 data SimplFloats
   = SimplFloats
       { -- Ordinary let bindings
diff --git a/compiler/simplCore/SimplUtils.hs b/compiler/simplCore/SimplUtils.hs
index 9420081d84..f2cf7a6606 100644
--- a/compiler/simplCore/SimplUtils.hs
+++ b/compiler/simplCore/SimplUtils.hs
@@ -17,7 +17,7 @@ module SimplUtils (
         simplEnvForGHCi, updModeForStableUnfoldings, updModeForRules,
 
         -- The continuation type
-        SimplCont(..), DupFlag(..),
+        SimplCont(..), DupFlag(..), StaticEnv,
         isSimplified, contIsStop,
         contIsDupable, contResultType, contHoleType,
         contIsTrivial, contArgs,
@@ -117,7 +117,7 @@ data SimplCont
   | ApplyToVal         -- (ApplyToVal arg K)[e] = K[ e arg ]
       { sc_dup  :: DupFlag      -- See Note [DupFlag invariants]
       , sc_arg  :: InExpr       -- The argument,
-      , sc_env  :: StaticEnv    --     and its static env
+      , sc_env  :: StaticEnv    -- see Note [StaticEnv invariant]
       , sc_cont :: SimplCont }
 
   | ApplyToTy          -- (ApplyToTy ty K)[e] = K[ e ty ]
@@ -130,7 +130,7 @@ data SimplCont
       { sc_dup  :: DupFlag        -- See Note [DupFlag invariants]
       , sc_bndr :: InId           -- case binder
       , sc_alts :: [InAlt]        -- Alternatives
-      , sc_env  :: StaticEnv      --   and their static environment
+      , sc_env  :: StaticEnv      -- See Note [StaticEnv invariant]
       , sc_cont :: SimplCont }
 
   -- The two strict forms have no DupFlag, because we never duplicate them
@@ -140,7 +140,7 @@ data SimplCont
       , sc_bndr  :: InId
       , sc_bndrs :: [InBndr]
       , sc_body  :: InExpr
-      , sc_env   :: StaticEnv
+      , sc_env   :: StaticEnv      -- See Note [StaticEnv invariant]
       , sc_cont  :: SimplCont }
 
   | StrictArg           -- (StrictArg (f e1 ..en) K)[e] = K[ f e1 .. en e ]
@@ -154,6 +154,8 @@ data SimplCont
         (Tickish Id)    -- Tick tickish <hole>
         SimplCont
 
+type StaticEnv = SimplEnv       -- Just the static part is relevant
+
 data DupFlag = NoDup       -- Unsimplified, might be big
              | Simplified  -- Simplified
              | OkToDup     -- Simplified and small
@@ -167,7 +169,25 @@ perhapsSubstTy dup env ty
   | isSimplified dup = ty
   | otherwise        = substTy env ty
 
-{-
+{- Note [StaticEnv invariant]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+We pair up an InExpr or InAlts with a StaticEnv, which establishes the
+lexical scope for that InExpr.  When we simplify that InExpr/InAlts, we
+use
+  - Its captured StaticEnv
+  - Overriding its InScopeSet with the larger one at the
+    simplification point.
+
+Why override the InScopeSet?  Example:
+      (let y = ey in f) ex
+By the time we simplify ex, 'y' will be in scope.
+
+However the InScopeSet in the StaticEnv is not irrelevant: it should
+include all the free vars of applying the substitution to the InExpr.
+Reason: contHoleType uses perhapsSubstTy to apply the substitution to
+the expression, and that (rightly) gives ASSERT failures if the InScopeSet
+isn't big enough.
+
 Note [DupFlag invariants]
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 In both (ApplyToVal dup _ env k)
diff --git a/compiler/simplCore/Simplify.hs b/compiler/simplCore/Simplify.hs
index 1e1b6ee27e..b24163695e 100644
--- a/compiler/simplCore/Simplify.hs
+++ b/compiler/simplCore/Simplify.hs
@@ -1467,7 +1467,7 @@ simplNonRecJoinPoint env bndr rhs body cont
        ; simplExprF (extendIdSubst env bndr (mkContEx env rhs)) body cont }
 
    | otherwise
-   = wrapJoinCont env cont $ \ cont ->
+   = wrapJoinCont env cont $ \ env cont ->
      do { -- We push join_cont into the join RHS and the body;
           -- and wrap wrap_cont around the whole thing
         ; let res_ty = contResultType cont
@@ -1483,7 +1483,7 @@ simplRecJoinPoint :: SimplEnv -> [(InId, InExpr)]
                   -> InExpr -> SimplCont
                   -> SimplM (SimplFloats, OutExpr)
 simplRecJoinPoint env pairs body cont
-  = wrapJoinCont env cont $ \ cont ->
+  = wrapJoinCont env cont $ \ env cont ->
     do { let bndrs = map fst pairs
              res_ty = contResultType cont
        ; env1 <- simplRecJoinBndrs env res_ty bndrs
@@ -1495,17 +1495,17 @@ simplRecJoinPoint env pairs body cont
 
 --------------------
 wrapJoinCont :: SimplEnv -> SimplCont
-             -> (SimplCont -> SimplM (SimplFloats, OutExpr))
+             -> (SimplEnv -> SimplCont -> SimplM (SimplFloats, OutExpr))
              -> SimplM (SimplFloats, OutExpr)
 -- Deal with making the continuation duplicable if necessary,
 -- and with the no-case-of-case situation.
 wrapJoinCont env cont thing_inside
   | contIsStop cont        -- Common case; no need for fancy footwork
-  = thing_inside cont
+  = thing_inside env cont
 
   | not (sm_case_case (getMode env))
     -- See Note [Join points wih -fno-case-of-case]
-  = do { (floats1, expr1) <- thing_inside (mkBoringStop (contHoleType cont))
+  = do { (floats1, expr1) <- thing_inside env (mkBoringStop (contHoleType cont))
        ; let (floats2, expr2) = wrapJoinFloatsX floats1 expr1
        ; (floats3, expr3) <- rebuild (env `setInScopeFromF` floats2) expr2 cont
        ; return (floats2 `addFloats` floats3, expr3) }
@@ -1513,7 +1513,7 @@ wrapJoinCont env cont thing_inside
   | otherwise
     -- Normal case; see Note [Join points and case-of-case]
   = do { (floats1, cont')  <- mkDupableCont env cont
-       ; (floats2, result) <- thing_inside cont'
+       ; (floats2, result) <- thing_inside (env `setInScopeFromF` floats1) cont'
        ; return (floats1 `addFloats` floats2, result) }
 
 
@@ -2463,7 +2463,10 @@ simplAlts :: SimplEnv
           -> SimplM OutExpr  -- Returns the complete simplified case expression
 
 simplAlts env0 scrut case_bndr alts cont'
-  = do  { (env1, case_bndr1) <- simplBinder env0 case_bndr
+  = do  { traceSmpl "simplAlts" (vcat [ ppr case_bndr
+                                      , text "cont':" <+> ppr cont'
+                                      , text "in_scope" <+> ppr (seInScope env0) ])
+        ; (env1, case_bndr1) <- simplBinder env0 case_bndr
         ; let case_bndr2 = case_bndr1 `setIdUnfolding` evaldUnfolding
               env2       = modifyInScope env1 case_bndr2
               -- See Note [Case binder evaluated-ness]
@@ -2855,7 +2858,8 @@ mkDupableCont env (StrictBind { sc_bndr = bndr, sc_bndrs = bndrs
        ; return ( floats2
                 , StrictBind { sc_bndr = bndr', sc_bndrs = []
                              , sc_body = body2
-                             , sc_env  = zapSubstEnv se
+                             , sc_env  = zapSubstEnv se `setInScopeFromF` floats2
+                                         -- See Note [StaticEnv invariant] in SimplUtils
                              , sc_dup  = OkToDup
                              , sc_cont = mkBoringStop res_ty } ) }
 
@@ -2888,12 +2892,18 @@ mkDupableCont env (ApplyToVal { sc_arg = arg, sc_dup = dup
         ; let env' = env `setInScopeFromF` floats1
         ; (_, se', arg') <- simplArg env' dup se arg
         ; (let_floats2, arg'') <- makeTrivial (getMode env) NotTopLevel (fsLit "karg") arg'
-        ; return ( floats1 `addLetFloats` let_floats2
-                 , ApplyToVal { sc_arg = arg'', sc_env = se'
+        ; let all_floats = floats1 `addLetFloats` let_floats2
+        ; return ( all_floats
+                 , ApplyToVal { sc_arg = arg''
+                              , sc_env = se' `setInScopeFromF` all_floats
+                                         -- Ensure that sc_env includes the free vars of
+                                         -- arg'' in its in-scope set, even if makeTrivial
+                                         -- has turned arg'' into a fresh variable
+                                         -- See Note [StaticEnv invariant] in SimplUtils
                               , sc_dup = OkToDup, sc_cont = cont' }) }
 
 mkDupableCont env (Select { sc_bndr = case_bndr, sc_alts = alts
-                            , sc_env = se, sc_cont = cont })
+                          , sc_env = se, sc_cont = cont })
   =     -- e.g.         (case [...hole...] of { pi -> ei })
         --      ===>
         --              let ji = \xij -> ei
@@ -2923,12 +2933,15 @@ mkDupableCont env (Select { sc_bndr = case_bndr, sc_alts = alts
         ; (join_floats, alts'') <- mapAccumLM (mkDupableAlt (seDynFlags env) case_bndr')
                                               emptyJoinFloats alts'
 
-        ; return (floats `addJoinFloats` join_floats,  -- Note [Duplicated env]
-                  Select { sc_dup  = OkToDup
-                         , sc_bndr = case_bndr'
-                         , sc_alts = alts''
-                         , sc_env  = zapSubstEnv env
-                         , sc_cont = mkBoringStop (contResultType cont) } ) }
+        ; let all_floats = floats `addJoinFloats` join_floats
+                           -- Note [Duplicated env]
+        ; return (all_floats
+                 , Select { sc_dup  = OkToDup
+                          , sc_bndr = case_bndr'
+                          , sc_alts = alts''
+                          , sc_env  = zapSubstEnv se `setInScopeFromF` all_floats
+                                      -- See Note [StaticEnv invariant] in SimplUtils
+                          , sc_cont = mkBoringStop (contResultType cont) } ) }
 
 mkDupableAlt :: DynFlags -> OutId
              -> JoinFloats -> OutAlt