Two improvements to optCoercion

* Fix a bug that meant that 
     (right (inst (forall tv.co) ty)) 
  wasn't getting optimised.  This showed up in the
  compiled code for ByteCodeItbls

* Add a substitution to optCoercion, so that it simultaneously
  substitutes and optimises.  Both call sites wanted this, and
  optCoercion itself can use it, so it seems a win all round.

compiler/coreSyn/CoreSubst.lhs

@@ -39,6 +39,7 @@ import OccurAnal( occurAnalyseExpr )
 
 import qualified Type
 import Type     ( Type, TvSubst(..), TvSubstEnv )
+import Coercion ( optCoercion )
 import VarSet
 import VarEnv
 import Id
@@ -290,7 +291,10 @@ substExpr subst expr
     go (Lit lit)       = Lit lit
     go (App fun arg)   = App (go fun) (go arg)
     go (Note note e)   = Note (go_note note) (go e)
-    go (Cast e co)     = Cast (go e) (substTy subst co)
+    go (Cast e co)     = Cast (go e) (optCoercion (getTvSubst subst) co)
+	-- Optimise coercions as we go; this is good, for example
+	-- in the RHS of rules, which are only substituted in
+
     go (Lam bndr body) = Lam bndr' (substExpr subst' body)
 		       where
 			 (subst', bndr') = substBndr subst bndr
@@ -463,8 +467,10 @@ substTyVarBndr (Subst in_scope id_env tv_env) tv
 
 -- | See 'Type.substTy'
 substTy :: Subst -> Type -> Type 
-substTy (Subst in_scope _id_env tv_env) ty
-  = Type.substTy (TvSubst in_scope tv_env) ty
+substTy subst ty = Type.substTy (getTvSubst subst) ty
+
+getTvSubst :: Subst -> TvSubst
+getTvSubst (Subst in_scope _id_env tv_env) = TvSubst in_scope tv_env
 \end{code}
 
 

compiler/simplCore/SimplEnv.lhs

@@ -29,7 +29,7 @@ module SimplEnv (
 
 	simplNonRecBndr, simplRecBndrs, simplLamBndr, simplLamBndrs, 
  	simplBinder, simplBinders, addBndrRules,
-	substExpr, substTy, mkCoreSubst,
+	substExpr, substTy, getTvSubst, mkCoreSubst,
 
 	-- Floats
   	Floats, emptyFloats, isEmptyFloats, addNonRec, addFloats, extendFloats,
@@ -687,13 +687,16 @@ addBndrRules env in_id out_id
 %************************************************************************
 
 \begin{code}
+getTvSubst :: SimplEnv -> TvSubst
+getTvSubst (SimplEnv { seInScope = in_scope, seTvSubst = tv_env })
+  = mkTvSubst in_scope tv_env
+
 substTy :: SimplEnv -> Type -> Type 
-substTy (SimplEnv { seInScope = in_scope, seTvSubst = tv_env }) ty
-  = Type.substTy (TvSubst in_scope tv_env) ty
+substTy env ty = Type.substTy (getTvSubst env) ty
 
 substTyVarBndr :: SimplEnv -> TyVar -> (SimplEnv, TyVar)
-substTyVarBndr env@(SimplEnv { seInScope = in_scope, seTvSubst = tv_env }) tv
-  = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
+substTyVarBndr env tv
+  = case Type.substTyVarBndr (getTvSubst env) tv of
 	(TvSubst in_scope' tv_env', tv') 
 	   -> (env { seInScope = in_scope', seTvSubst = tv_env'}, tv')
 

compiler/simplCore/Simplify.lhs

@@ -874,7 +874,7 @@ simplType :: SimplEnv -> InType -> SimplM OutType
         -- Kept monadic just so we can do the seqType
 simplType env ty
   = -- pprTrace "simplType" (ppr ty $$ ppr (seTvSubst env)) $
-    seqType new_ty   `seq`   return new_ty
+    seqType new_ty `seq` return new_ty
   where
     new_ty = substTy env ty
 
@@ -883,8 +883,9 @@ simplCoercion :: SimplEnv -> InType -> SimplM OutType
 -- The InType isn't *necessarily* a coercion, but it might be
 -- (in a type application, say) and optCoercion is a no-op on types
 simplCoercion env co
-  = do { co' <- simplType env co
-       ; return (optCoercion co') }
+  = seqType new_co `seq` return new_co
+  where 
+    new_co = optCoercion (getTvSubst env) co
 \end{code}
 
 

compiler/types/Coercion.lhs

@@ -699,20 +699,24 @@ mkEqPredCoI _   (ACo co1) ty2 coi2      = ACo $ PredTy $ EqPred co1 (fromCoI coi
 %************************************************************************
 
 \begin{code}
+optCoercion :: TvSubst -> Coercion -> NormalCo
+-- ^ optCoercion applies a substitution to a coercion, 
+--   *and* optimises it to reduce its size
+optCoercion env co = opt_co env False co
+
 type NormalCo = Coercion
   -- Invariants: 
+  --  * The substitution has been fully applied
   --  * For trans coercions (co1 `trans` co2)
   --       co1 is not a trans, and neither co1 nor co2 is identity
   --  * If the coercion is the identity, it has no CoVars of CoTyCons in it (just types)
 
 type NormalNonIdCo = NormalCo  -- Extra invariant: not the identity
 
-optCoercion :: Coercion -> NormalCo
-optCoercion co = opt_co False co
-
-opt_co :: Bool	       -- True <=> return (sym co)
-       -> Coercion
-       -> NormalCo
+opt_co, opt_co' :: TvSubst
+       		-> Bool	       -- True <=> return (sym co)
+       		-> Coercion
+       		-> NormalCo	
 opt_co = opt_co'
 -- opt_co sym co = pprTrace "opt_co {" (ppr sym <+> ppr co) $
 --       	        co1 `seq` 
@@ -728,12 +732,14 @@ opt_co = opt_co'
 --            | otherwise = (s,t)
 --    (s2,t2) = coercionKind co1
 
-opt_co' sym (AppTy ty1 ty2) 	      = mkAppTy (opt_co sym ty1) (opt_co sym ty2)
-opt_co' sym (FunTy ty1 ty2) 	      = FunTy (opt_co sym ty1) (opt_co sym ty2)
-opt_co' sym (PredTy (ClassP cls tys)) = PredTy (ClassP cls (map (opt_co sym) tys))
-opt_co' sym (PredTy (IParam n ty))    = PredTy (IParam n (opt_co sym ty))
+opt_co' env sym (AppTy ty1 ty2) 	  = mkAppTy (opt_co env sym ty1) (opt_co env sym ty2)
+opt_co' env sym (FunTy ty1 ty2) 	  = FunTy (opt_co env sym ty1) (opt_co env sym ty2)
+opt_co' env sym (PredTy (ClassP cls tys)) = PredTy (ClassP cls (map (opt_co env sym) tys))
+opt_co' env sym (PredTy (IParam n ty))    = PredTy (IParam n (opt_co env sym ty))
+opt_co' _   _   co@(PredTy (EqPred {}))   = pprPanic "optCoercion" (ppr co)
 
-opt_co' sym co@(TyVarTy tv)
+opt_co' env sym co@(TyVarTy tv)
+  | Just ty <- lookupTyVar env tv = opt_co' (zapTvSubstEnv env) sym ty
   | not (isCoVar tv)     = co   -- Identity; does not mention a CoVar
   | ty1 `coreEqType` ty2 = ty1	-- Identity; ..ditto..
   | not sym              = co
@@ -741,41 +747,43 @@ opt_co' sym co@(TyVarTy tv)
   where
     (ty1,ty2) = coVarKind tv
 
-opt_co' sym (ForAllTy tv cor) 
-  | isCoVar tv = mkCoPredTy (opt_co sym co1) (opt_co sym co2) (opt_co sym cor)
-  | otherwise  = ForAllTy tv (opt_co sym cor)
+opt_co' env sym (ForAllTy tv cor) 
+  | isCoVar tv = mkCoPredTy (opt_co env sym co1) (opt_co env sym co2) (opt_co env sym cor)
+  | otherwise  = case substTyVarBndr env tv of
+                   (env', tv') -> ForAllTy tv' (opt_co env' sym cor)
   where
     (co1,co2) = coVarKind tv
 
-opt_co' sym (TyConApp tc cos)
+opt_co' env sym (TyConApp tc cos)
   | isCoercionTyCon tc
-  = foldl mkAppTy opt_co_tc 
-          (map (opt_co sym) (drop arity cos))
+  = foldl mkAppTy 
+	  (opt_co_tc_app env sym tc (take arity cos))
+          (map (opt_co env sym) (drop arity cos))
   | otherwise
-  = TyConApp tc (map (opt_co sym) cos)
+  = TyConApp tc (map (opt_co env sym) cos)
   where
     arity = tyConArity tc
-    opt_co_tc :: NormalCo
-    opt_co_tc = opt_co_tc_app sym tc (take arity cos)
 
 --------
-opt_co_tc_app :: Bool -> TyCon -> [Type] -> NormalCo
+opt_co_tc_app :: TvSubst -> Bool -> TyCon -> [Coercion] -> NormalCo
 -- Used for CoercionTyCons only
-opt_co_tc_app sym tc cos
+-- Arguments are *not* already simplified/substituted
+opt_co_tc_app env sym tc cos
   | tc `hasKey` symCoercionTyConKey
-  = opt_co (not sym) co1
+  = opt_co env (not sym) co1
 
   | tc `hasKey` transCoercionTyConKey
-  = if sym then opt_trans opt_co2 opt_co1
+  = if sym then opt_trans opt_co2 opt_co1   -- sym (g `o` h) = sym h `o` sym g
            else opt_trans opt_co1 opt_co2
 
   | tc `hasKey` leftCoercionTyConKey
-  , Just (co1, _) <- splitAppTy_maybe opt_co1
-  = co1
+  , Just (opt_co1_left, _) <- splitAppTy_maybe opt_co1
+  = opt_co1_left	-- sym (left g) = left (sym g) 
+			-- The opt_co has the sym pushed into it
 
   | tc `hasKey` rightCoercionTyConKey
-  , Just (_, co2) <- splitAppTy_maybe opt_co1
-  = co2
+  , Just (_, opt_co1_right) <- splitAppTy_maybe opt_co1
+  = opt_co1_right
 
   | tc `hasKey` csel1CoercionTyConKey
   , Just (s1,_,_) <- splitCoPredTy_maybe opt_co1
@@ -789,10 +797,16 @@ opt_co_tc_app sym tc cos
   , Just (_,_,r) <- splitCoPredTy_maybe opt_co1
   = r
 
-  | tc `hasKey` instCoercionTyConKey
-  , Just (tv, co'') <- splitForAllTy_maybe opt_co1
-  , let ty = co2
-  = substTyWith [tv] [ty] co''
+  | tc `hasKey` instCoercionTyConKey	-- See if the first arg 	
+					-- is already a forall
+  , Just (tv, co1_body) <- splitForAllTy_maybe co1
+  , let ty = substTy env co2
+  = opt_co (extendTvSubst env tv ty) sym co1_body
+
+  | tc `hasKey` instCoercionTyConKey	-- See if is *now* a forall
+  , Just (tv, opt_co1_body) <- splitForAllTy_maybe opt_co1
+  , let ty = substTy env co2
+  = substTyWith [tv] [ty] opt_co1_body	-- An inefficient one-variable substitution
 
   | otherwise	  -- Do not push sym inside top-level axioms
     		  -- e.g. if g is a top-level axiom
@@ -801,11 +815,15 @@ opt_co_tc_app sym tc cos
   = if sym then mkSymCoercion the_co 
            else the_co
   where
-    the_co = TyConApp tc cos
     (co1 : cos1) = cos
     (co2 : _)    = cos1
-    opt_co1 = opt_co sym co1
-    opt_co2 = opt_co sym co2
+
+	-- These opt_cos have the sym pushed into them
+    opt_co1 = opt_co env sym co1
+    opt_co2 = opt_co env sym co2
+
+	-- However the_co does *not* have sym pushed into it
+    the_co = TyConApp tc (map (opt_co env False) cos)
 
 -------------
 opt_trans :: NormalCo -> NormalCo -> NormalCo

compiler/types/Type.lhs

@@ -123,7 +123,8 @@ module Type (
 	emptyTvSubstEnv, emptyTvSubst,
 	
 	mkTvSubst, mkOpenTvSubst, zipOpenTvSubst, zipTopTvSubst, mkTopTvSubst, notElemTvSubst,
-	getTvSubstEnv, setTvSubstEnv, getTvInScope, extendTvInScope, extendTvInScopeList,
+	getTvSubstEnv, setTvSubstEnv, zapTvSubstEnv, getTvInScope, 
+        extendTvInScope, extendTvInScopeList,
  	extendTvSubst, extendTvSubstList, isInScope, composeTvSubst, zipTyEnv,
         isEmptyTvSubst,
 
@@ -1437,6 +1438,9 @@ notElemTvSubst tv (TvSubst _ env) = not (tv `elemVarEnv` env)
 setTvSubstEnv :: TvSubst -> TvSubstEnv -> TvSubst
 setTvSubstEnv (TvSubst in_scope _) env = TvSubst in_scope env
 
+zapTvSubstEnv :: TvSubst -> TvSubst
+zapTvSubstEnv (TvSubst in_scope _) = TvSubst in_scope emptyVarEnv
+
 extendTvInScope :: TvSubst -> Var -> TvSubst
 extendTvInScope (TvSubst in_scope env) var = TvSubst (extendInScopeSet in_scope var) env