Fix handing of CoVars in SetLevels: it wasn't renaming occurrences of case-bound coercion variabes

Conflicts:

	compiler/simplCore/SetLevels.lhs

1 files changed, 30 insertions, 31 deletions

diff --git a/compiler/simplCore/SetLevels.lhs b/compiler/simplCore/SetLevels.lhs
index c90d7a2b20..3f665fcbdb 100644
--- a/compiler/simplCore/SetLevels.lhs
+++ b/compiler/simplCore/SetLevels.lhs
@@ -59,8 +59,9 @@ import CoreMonad	( FloatOutSwitches(..) )
 import CoreUtils	( exprType, exprOkForSpeculation, mkPiTypes )
 import CoreArity	( exprBotStrictness_maybe )
 import CoreFVs		-- all of it
+import Coercion         ( isCoVar )
 import CoreSubst	( Subst, emptySubst, extendInScope, substBndr, substRecBndrs,
-			  extendIdSubst, cloneBndrs, cloneIdBndr, cloneRecIdBndrs )
+			  extendIdSubst, extendSubstWithVar, cloneBndr, cloneRecIdBndrs, substTy, substCo )
 import Id
 import IdInfo
 import Var
@@ -73,6 +74,7 @@ import Type		( isUnLiftedType, Type )
 import BasicTypes	( Arity )
 import UniqSupply
 import Util
+import MonadUtils
 import Outputable
 import FastString
 \end{code}
@@ -264,8 +266,8 @@ don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
 If there were another lambda in @r@'s rhs, it would get level-2 as well.
 
 \begin{code}
-lvlExpr _ _ (  _, AnnType ty) = return (Type ty)
-lvlExpr _ _ ( _, AnnCoercion co) = return (Coercion co)
+lvlExpr _ env (_, AnnType ty) = return (Type (substTy (le_subst env) ty))
+lvlExpr _ env (_, AnnCoercion co) = return (Coercion (substCo (le_subst env) co))
 lvlExpr _ env (_, AnnVar v)   = return (lookupVar env v)
 lvlExpr _ _   (_, AnnLit lit) = return (Lit lit)
 
@@ -312,7 +314,7 @@ lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
 
 lvlExpr ctxt_lvl env (_, AnnCast expr (_, co)) = do
     expr' <- lvlExpr ctxt_lvl env expr
-    return (Cast expr' co)
+    return (Cast expr' (substCo (le_subst env) co))
 
 -- We don't split adjacent lambdas.  That is, given
 --	\x y -> (x+1,y)
@@ -436,8 +438,8 @@ lvlMFE ::  Bool			-- True <=> strict context [body of case or let]
 -- lvlMFE is just like lvlExpr, except that it might let-bind
 -- the expression, so that it can itself be floated.
 
-lvlMFE _ _ _ (_, AnnType ty)
-  = return (Type ty)
+lvlMFE _ _ env (_, AnnType ty)
+  = return (Type (substTy (le_subst env) ty))
 
 -- No point in floating out an expression wrapped in a coercion or note
 -- If we do we'll transform  lvl = e |> co 
@@ -449,7 +451,7 @@ lvlMFE strict_ctxt ctxt_lvl env (_, AnnNote n e)
 
 lvlMFE strict_ctxt ctxt_lvl env (_, AnnCast e (_, co))
   = do	{ e' <- lvlMFE strict_ctxt ctxt_lvl env e
-	; return (Cast e' co) }
+	; return (Cast e' (substCo (le_subst env) co)) }
 
 -- Note [Case MFEs]
 lvlMFE True ctxt_lvl env e@(_, AnnCase {})
@@ -646,6 +648,8 @@ lvlBind :: Level		-- Context level; might be Top even for bindings
 lvlBind ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
   | isTyVar bndr    -- Don't do anything for TyVar binders
 	            --   (simplifier gets rid of them pronto)
+  || isCoVar bndr   -- Difficult to fix up CoVar occurrences (see extendPolyLvlEnv)
+                    -- so we will ignore this case for now
   || not (profitableFloat ctxt_lvl dest_lvl)
   || (isTopLvl dest_lvl && isUnLiftedType (idType bndr))
 	  -- We can't float an unlifted binding to top level, so we don't 
@@ -846,6 +850,9 @@ data LevelEnv
        , le_lvl_env  :: VarEnv Level	-- Domain is *post-cloned* TyVars and Ids
        , le_subst    :: Subst 		-- Domain is pre-cloned Ids; tracks the in-scope set
 					-- 	so that substitution is capture-avoiding
+                                        -- The Id -> CoreExpr in the Subst is ignored
+                                        -- (since we want to substitute in LevelledExpr
+                                        -- instead) but we do use the Co/TyVar substs
        , le_env      :: IdEnv ([Var], LevelledExpr)	-- Domain is pre-cloned Ids
     }
 	-- We clone let-bound variables so that they are still
@@ -914,17 +921,18 @@ extendCaseBndrLvlEnv :: LevelEnv -> Expr LevelledBndr
                      -> LevelledBndr -> LevelEnv
 extendCaseBndrLvlEnv le@(LE { le_subst = subst, le_env = id_env }) 
                      (Var scrut_var) (TB case_bndr _)
-  = le { le_subst   = extendIdSubst subst case_bndr (Var scrut_var)
-       , le_env     = extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var) }
+  = le { le_subst   = extendSubstWithVar subst case_bndr scrut_var
+       , le_env     = extendVarEnv id_env case_bndr ([scrut_var], ASSERT(not (isCoVar scrut_var)) Var scrut_var) }
      
 extendCaseBndrLvlEnv env _scrut case_bndr
   = extendLvlEnv env [case_bndr]
 
-extendPolyLvlEnv :: Level -> LevelEnv -> [Var] -> [(Var, Var)] -> LevelEnv
+extendPolyLvlEnv :: Level -> LevelEnv -> [Var] -> [(Var {- :: t -}, Var {- :: mkPiTypes abs_vars t -})] -> LevelEnv
 extendPolyLvlEnv dest_lvl 
                  le@(LE { le_lvl_env = lvl_env, le_subst = subst, le_env = id_env }) 
                  abs_vars bndr_pairs
-  = le { le_lvl_env = foldl add_lvl   lvl_env bndr_pairs
+   = ASSERT( all (not . isCoVar . fst) bndr_pairs ) -- What would we add to the CoSubst in this case. No easy answer, so avoid floating 
+    le { le_lvl_env = foldl add_lvl   lvl_env bndr_pairs
        , le_subst   = foldl add_subst subst   bndr_pairs
        , le_env     = foldl add_id    id_env  bndr_pairs }
   where
@@ -939,8 +947,10 @@ extendCloneLvlEnv lvl le@(LE { le_lvl_env = lvl_env, le_env = id_env })
        , le_subst   = new_subst
        , le_env     = foldl add_id  id_env  bndr_pairs }
   where
-     add_lvl env (_, v') = extendVarEnv env v' lvl
-     add_id  env (v, v') = extendVarEnv env v ([v'], Var v')
+     add_lvl env (_, v_cloned) = extendVarEnv env v_cloned lvl
+     add_id  env (v, v_cloned) = if isTyVar v
+                                 then delVarEnv    env v
+                                 else extendVarEnv env v ([v_cloned], ASSERT(not (isCoVar v_cloned)) Var v_cloned)
 
 maxFvLevel :: (Var -> Bool) -> LevelEnv -> VarSet -> Level
 maxFvLevel max_me (LE { le_lvl_env = lvl_env, le_env = id_env }) var_set
@@ -1084,30 +1094,19 @@ substLetBndrsRec
                           , le_subst = subst'
                           , le_env = delVarEnvList id_env bndrs }
 
-cloneVar :: LevelEnv -> Id -> Level -> LvlM (LevelEnv, Id)
-cloneVar env v dest_lvl
-  = ASSERT( isId v )
-    do { us <- getUniqueSupplyM
-       ; let (subst', v1) = cloneIdBndr (le_subst env) us v
-      	     v2     	  = zapDemandIdInfo v1
+cloneVar :: LevelEnv -> Var -> Level -> LvlM (LevelEnv, Var)
+cloneVar env v dest_lvl -- Works for Ids, TyVars and CoVars
+  = do { u <- getUniqueM
+       ; let (subst', v1) = cloneBndr (le_subst env) u v
+      	     v2     	  = if isId v1 then zapDemandIdInfo v1 else v1
       	     env'	  = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
        ; return (env', v2) }
 
 cloneVars :: LevelEnv -> [Var] -> Level -> LvlM (LevelEnv, [Var])
-cloneVars env vs dest_lvl    -- Works for tyvars etc too; typically case alts
-  = do { us <- getUniqueSupplyM
-       ; let (subst', vs1) = cloneBndrs (le_subst env) us vs
-      	     vs2     	   = map zap_demand vs1 
-      	     env'	   = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
-       ; return (env', vs2) }
-  where
-    zap_demand :: Var -> Var -- Note [Zapping the demand info]
-    zap_demand  v | not (isId v) = v
-                  | otherwise    = zapDemandIdInfo v
-
+cloneVars env vs dest_lvl = mapAccumLM (\env v -> cloneVar env v dest_lvl) env vs
 
 cloneRecVars :: LevelEnv -> [Id] -> Level -> LvlM (LevelEnv, [Id])
-cloneRecVars env vs dest_lvl
+cloneRecVars env vs dest_lvl -- Works for CoVars too (since cloneRecIdBndrs does)
   = ASSERT( all isId vs ) do
     us <- getUniqueSupplyM
     let