Fix exprIsConApp_maybe

In this commit
   commit 7833cf407d1f608bebb1d38bb99d3035d8d735e6
   Date:   Thu Jan 24 17:58:50 2019 +0100
      Look through newtype wrappers (Trac #16254)

we made exprIsConApp_maybe quite a bit cleverer.  But I had not paid
enough attention to keeping exactly the correct substitution and
in-scope set, which led to Trac #16348.

There were several buglets (like applying the substitution twice in
exprIsConApp_maybe, but the proximate source of the bug was that we were
calling addNewInScopeIds, which deleted things from the substitution as
well as adding them to the in-scope set.  That's usually right, but not
here!

This was quite tricky to track down.  But it is nicer now.

6 files changed, 81 insertions, 55 deletions

diff --git a/compiler/basicTypes/MkId.hs b/compiler/basicTypes/MkId.hs
index 98ff0b0c3d..1802cd769e 100644
--- a/compiler/basicTypes/MkId.hs
+++ b/compiler/basicTypes/MkId.hs
@@ -409,7 +409,7 @@ dictSelRule :: Int -> Arity -> RuleFun
 --
 dictSelRule val_index n_ty_args _ id_unf _ args
   | (dict_arg : _) <- drop n_ty_args args
-  , Just (floats, _, _, con_args) <- exprIsConApp_maybe id_unf dict_arg
+  , Just (_, floats, _, _, con_args) <- exprIsConApp_maybe id_unf dict_arg
   = Just (wrapFloats floats $ getNth con_args val_index)
   | otherwise
   = Nothing
diff --git a/compiler/coreSyn/CoreOpt.hs b/compiler/coreSyn/CoreOpt.hs
index a2ac7b5be9..80fb3a80cf 100644
--- a/compiler/coreSyn/CoreOpt.hs
+++ b/compiler/coreSyn/CoreOpt.hs
@@ -28,7 +28,7 @@ import CoreSyn
 import CoreSubst
 import CoreUtils
 import CoreFVs
-import MkCore ( FloatBind(..), mkCoreLet )
+import MkCore ( FloatBind(..) )
 import PprCore  ( pprCoreBindings, pprRules )
 import OccurAnal( occurAnalyseExpr, occurAnalysePgm )
 import Literal  ( Literal(LitString) )
@@ -232,7 +232,7 @@ simple_opt_expr env expr
     go (Case e b ty as)
        -- See Note [Getting the map/coerce RULE to work]
       | isDeadBinder b
-      , Just ([], con, _tys, es) <- exprIsConApp_maybe in_scope_env e'
+      , Just (_, [], con, _tys, es) <- exprIsConApp_maybe in_scope_env e'
         -- We don't need to be concerned about floats when looking for coerce.
       , Just (altcon, bs, rhs) <- findAlt (DataAlt con) as
       = case altcon of
@@ -889,42 +889,58 @@ data ConCont = CC [CoreExpr] Coercion
 -- are unfolded late, but we really want to trigger case-of-known-constructor as
 -- early as possible. See also Note [Activation for data constructor wrappers]
 -- in MkId.
-exprIsConApp_maybe :: InScopeEnv -> CoreExpr -> Maybe ([FloatBind], DataCon, [Type], [CoreExpr])
+--
+-- We also return the incoming InScopeSet, augmented with
+-- the binders from any [FloatBind] that we return
+exprIsConApp_maybe :: InScopeEnv -> CoreExpr
+                   -> Maybe (InScopeSet, [FloatBind], DataCon, [Type], [CoreExpr])
 exprIsConApp_maybe (in_scope, id_unf) expr
-  = do
-    (floats, con, ty, args) <- go (Left in_scope) [] expr (CC [] (mkRepReflCo (exprType expr)))
-    return $ (reverse floats, con, ty, args)
+  = go (Left in_scope) [] expr (CC [] (mkRepReflCo (exprType expr)))
   where
     go :: Either InScopeSet Subst
              -- Left in-scope  means "empty substitution"
              -- Right subst    means "apply this substitution to the CoreExpr"
+             -- NB: in the call (go subst floats expr cont)
+             --     the substitution applies to 'expr', but /not/ to 'floats' or 'cont'
        -> [FloatBind] -> CoreExpr -> ConCont
              -- Notice that the floats here are in reverse order
-       -> Maybe ([FloatBind], DataCon, [Type], [CoreExpr])
+       -> Maybe (InScopeSet, [FloatBind], DataCon, [Type], [CoreExpr])
     go subst floats (Tick t expr) cont
        | not (tickishIsCode t) = go subst floats expr cont
+
     go subst floats (Cast expr co1) (CC args co2)
        | Just (args', m_co1') <- pushCoArgs (subst_co subst co1) args
             -- See Note [Push coercions in exprIsConApp_maybe]
        = case m_co1' of
            MCo co1' -> go subst floats expr (CC args' (co1' `mkTransCo` co2))
            MRefl    -> go subst floats expr (CC args' co2)
+
     go subst floats (App fun arg) (CC args co)
-       = go subst floats fun (CC (subst_arg subst arg : args) co)
-    go subst floats (Lam var body) (CC (arg:args) co)
+       = go subst floats fun (CC (subst_expr subst arg : args) co)
+
+    go subst floats (Lam bndr body) (CC (arg:args) co)
        | exprIsTrivial arg          -- Don't duplicate stuff!
-       = go (extend subst var arg) floats body (CC args co)
-    go subst floats (Lam var body) (CC (arg:args) co)
-       = go subst floats (mkCoreLet (NonRec var arg) body) (CC args co)
-    go subst floats (Let bndr@(NonRec _ _) expr) cont
-       = let (subst', bndr') = subst_bind subst bndr in
-           go subst' (FloatLet bndr' : floats) expr cont
+       = go (extend subst bndr arg) floats body (CC args co)
+       | otherwise
+       = let (subst', bndr') = subst_bndr subst bndr
+             float           = FloatLet (NonRec bndr' arg)
+         in go subst' (float:floats) body (CC args co)
+
+    go subst floats (Let (NonRec bndr rhs) expr) cont
+       = let rhs'            = subst_expr subst rhs
+             (subst', bndr') = subst_bndr subst bndr
+             float           = FloatLet (NonRec bndr' rhs')
+         in go subst' (float:floats) expr cont
+
     go subst floats (Case scrut b _ [(con, vars, expr)]) cont
        = let
-          (subst', b') = subst_bndr subst b
+          scrut'           = subst_expr subst scrut
+          (subst', b')     = subst_bndr subst b
           (subst'', vars') = subst_bndrs subst' vars
+          float            = FloatCase scrut' b' con vars'
          in
-           go subst'' (FloatCase (subst_arg subst scrut) b' con vars' : floats) expr cont
+           go subst'' (float:floats) expr cont
+
     go (Right sub) floats (Var v) cont
        = go (Left (substInScope sub))
             floats
@@ -935,7 +951,8 @@ exprIsConApp_maybe (in_scope, id_unf) expr
 
         | Just con <- isDataConWorkId_maybe fun
         , count isValArg args == idArity fun
-        = pushFloats floats $ pushCoDataCon con args co
+        = succeedWith in_scope floats $
+          pushCoDataCon con args co
 
         -- See Note [Special case for newtype wrappers]
         | Just a <- isDataConWrapId_maybe fun
@@ -954,7 +971,7 @@ exprIsConApp_maybe (in_scope, id_unf) expr
         | DFunUnfolding { df_bndrs = bndrs, df_con = con, df_args = dfun_args } <- unfolding
         , bndrs `equalLength` args    -- See Note [DFun arity check]
         , let subst = mkOpenSubst in_scope (bndrs `zip` args)
-        = pushFloats floats $
+        = succeedWith in_scope floats $
           pushCoDataCon con (map (substExpr (text "exprIsConApp1") subst) dfun_args) co
 
         -- Look through unfoldings, but only arity-zero one;
@@ -972,42 +989,44 @@ exprIsConApp_maybe (in_scope, id_unf) expr
           (fun `hasKey` unpackCStringUtf8IdKey)
         , [arg]              <- args
         , Just (LitString str) <- exprIsLiteral_maybe (in_scope, id_unf) arg
-        = pushFloats floats $ dealWithStringLiteral fun str co
+        = succeedWith in_scope floats $
+          dealWithStringLiteral fun str co
         where
           unfolding = id_unf fun
 
     go _ _ _ _ = Nothing
 
-    pushFloats :: [FloatBind] -> Maybe (DataCon, [Type], [CoreExpr]) -> Maybe ([FloatBind], DataCon, [Type], [CoreExpr])
-    pushFloats floats x = do
-      (c, tys, args) <- x
-      return (floats, c, tys, args)
+    succeedWith :: InScopeSet -> [FloatBind]
+                -> Maybe (DataCon, [Type], [CoreExpr])
+                -> Maybe (InScopeSet, [FloatBind], DataCon, [Type], [CoreExpr])
+    succeedWith in_scope rev_floats x
+      = do { (con, tys, args) <- x
+           ; let floats = reverse rev_floats
+           ; return (in_scope, floats, con, tys, args) }
+
+    ----------------------------
+    -- Unconditionally substitute the argument of a newtype
+    dealWithNewtypeWrapper scope floats (Lam v body) (CC (arg:args) co)
+      = dealWithNewtypeWrapper (extend scope v arg) floats body (CC args co)
+    dealWithNewtypeWrapper scope floats expr args
+      = go scope floats expr args
 
-    dealWithNewtypeWrapper scope floats (Lam v body) (CC (arg:args) co) =
-      dealWithNewtypeWrapper (extend scope v arg) floats body (CC args co)
-    dealWithNewtypeWrapper scope floats expr args = go scope floats expr args
     ----------------------------
     -- Operations on the (Either InScopeSet CoreSubst)
     -- The Left case is wildly dominant
     subst_co (Left {}) co = co
     subst_co (Right s) co = CoreSubst.substCo s co
 
-    subst_arg (Left {}) e = e
-    subst_arg (Right s) e = substExpr (text "exprIsConApp2") s e
-
-    subst_bind (Left in_scope) bndr@(NonRec b _) =
-      (Left (extendInScopeSet in_scope b), bndr)
-    subst_bind (Left _) _ =
-      error "CoreOpt.exprIsConApp_maybe: recursive float."
-    subst_bind (Right subst) bndr =
-      let (subst', bndr') = substBind subst bndr in
-      (Right subst', bndr')
-
-    subst_bndr (Left in_scope) b =
-      (Left (extendInScopeSet in_scope b), b)
-    subst_bndr (Right subst) b =
-      let (subst', b') = substBndr subst b in
-      (Right subst', b')
+    subst_expr (Left {}) e = e
+    subst_expr (Right s) e = substExpr (text "exprIsConApp2") s e
+
+    subst_bndr msubst bndr
+      = (Right subst', bndr')
+      where
+        (subst', bndr') = substBndr subst bndr
+        subst = case msubst of
+                  Left in_scope -> mkEmptySubst in_scope
+                  Right subst   -> subst
 
     subst_bndrs subst bs = mapAccumL subst_bndr subst bs
 
diff --git a/compiler/prelude/PrelRules.hs b/compiler/prelude/PrelRules.hs
index a6d7bcc425..3a0b1f7b9f 100644
--- a/compiler/prelude/PrelRules.hs
+++ b/compiler/prelude/PrelRules.hs
@@ -1039,7 +1039,7 @@ dataToTagRule = a `mplus` b
       dflags <- getDynFlags
       [_, val_arg] <- getArgs
       in_scope <- getInScopeEnv
-      (floats, dc,_,_) <- liftMaybe $ exprIsConApp_maybe in_scope val_arg
+      (_,floats, dc,_,_) <- liftMaybe $ exprIsConApp_maybe in_scope val_arg
       ASSERT( not (isNewTyCon (dataConTyCon dc)) ) return ()
       return $ wrapFloats floats (mkIntVal dflags (toInteger (dataConTagZ dc)))
 
diff --git a/compiler/simplCore/Simplify.hs b/compiler/simplCore/Simplify.hs
index 2bb177d25b..2156dc55b8 100644
--- a/compiler/simplCore/Simplify.hs
+++ b/compiler/simplCore/Simplify.hs
@@ -2399,26 +2399,27 @@ rebuildCase env scrut case_bndr alts cont
   = do  { tick (KnownBranch case_bndr)
         ; case findAlt (LitAlt lit) alts of
             Nothing           -> missingAlt env case_bndr alts cont
-            Just (_, bs, rhs) -> simple_rhs [] scrut bs rhs }
+            Just (_, bs, rhs) -> simple_rhs env [] scrut bs rhs }
 
-  | Just (wfloats, con, ty_args, other_args) <- exprIsConApp_maybe (getUnfoldingInRuleMatch env) scrut
+  | Just (in_scope', wfloats, con, ty_args, other_args)
+      <- exprIsConApp_maybe (getUnfoldingInRuleMatch env) scrut
         -- Works when the scrutinee is a variable with a known unfolding
         -- as well as when it's an explicit constructor application
+  , let env0 = setInScopeSet env in_scope'
   = do  { tick (KnownBranch case_bndr)
         ; case findAlt (DataAlt con) alts of
-            Nothing  -> missingAlt env case_bndr alts cont
+            Nothing  -> missingAlt env0 case_bndr alts cont
             Just (DEFAULT, bs, rhs) -> let con_app = Var (dataConWorkId con)
                                                  `mkTyApps` ty_args
                                                  `mkApps`   other_args
-                                       in simple_rhs wfloats con_app bs rhs
-            Just (_, bs, rhs)       -> knownCon env scrut wfloats con ty_args other_args
+                                       in simple_rhs env0 wfloats con_app bs rhs
+            Just (_, bs, rhs)       -> knownCon env0 scrut wfloats con ty_args other_args
                                                 case_bndr bs rhs cont
         }
   where
-    simple_rhs wfloats scrut' bs rhs =
+    simple_rhs env wfloats scrut' bs rhs =
       ASSERT( null bs )
-      do { let env0 = addNewInScopeIds env (concatMap MkCore.floatBindings wfloats)
-         ; (floats1, env') <- simplNonRecX env0 case_bndr scrut'
+      do { (floats1, env') <- simplNonRecX env case_bndr scrut'
              -- scrut is a constructor application,
              -- hence satisfies let/app invariant
          ; (floats2, expr') <- simplExprF env' rhs cont
@@ -2863,8 +2864,7 @@ knownCon :: SimplEnv
          -> SimplM (SimplFloats, OutExpr)
 
 knownCon env scrut dc_floats dc dc_ty_args dc_args bndr bs rhs cont
-  = do  { let env0 = addNewInScopeIds env (concatMap MkCore.floatBindings dc_floats)
-        ; (floats1, env1)  <- bind_args env0 bs dc_args
+  = do  { (floats1, env1)  <- bind_args env bs dc_args
         ; (floats2, env2) <- bind_case_bndr env1
         ; (floats3, expr') <- simplExprF env2 rhs cont
         ; case dc_floats of
diff --git a/testsuite/tests/simplCore/should_compile/T16348.hs b/testsuite/tests/simplCore/should_compile/T16348.hs
new file mode 100644
index 0000000000..307ad64085
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T16348.hs
@@ -0,0 +1,6 @@
+module T16348 where
+
+data V2 a = V2 !a !a
+
+inv22 _ = case V2 (V2 1 2) (V2 3 4) of
+             V2 _ (V2 _ z) -> z
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 6e1979c5e6..1bb4694645 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -301,3 +301,4 @@ test('T15631',
      makefile_test, ['T15631'])
 test('T15673', normal, compile, ['-O'])
 test('T16288', normal, multimod_compile, ['T16288B', '-O -dcore-lint -v0'])
+test('T16348', normal, compile, ['-O'])