Increase an InScopeSet for a substitution
[ghc.git] / compiler / types / FamInstEnv.hs
index 9822236..62906dd 100644 (file)
@@ -12,11 +12,11 @@ module FamInstEnv (
 
         FamInstEnvs, FamInstEnv, emptyFamInstEnv, emptyFamInstEnvs,
         extendFamInstEnv, deleteFromFamInstEnv, extendFamInstEnvList,
-        identicalFamInstHead, famInstEnvElts, familyInstances, orphNamesOfFamInst,
+        identicalFamInstHead, famInstEnvElts, familyInstances,
 
         -- * CoAxioms
         mkCoAxBranch, mkBranchedCoAxiom, mkUnbranchedCoAxiom, mkSingleCoAxiom,
-        computeAxiomIncomps,
+        mkNewTypeCoAxiom,
 
         FamInstMatch(..),
         lookupFamInstEnv, lookupFamInstEnvConflicts, lookupFamInstEnvByTyCon,
@@ -25,13 +25,12 @@ module FamInstEnv (
 
         -- Injectivity
         InjectivityCheckResult(..),
-        lookupFamInstEnvInjectivityConflicts, unusedInjTvsInRHS, isTFHeaded,
-        bareTvInRHSViolated, injectiveBranches,
+        lookupFamInstEnvInjectivityConflicts, injectiveBranches,
 
         -- Normalisation
         topNormaliseType, topNormaliseType_maybe,
         normaliseType, normaliseTcApp,
-        reduceTyFamApp_maybe, chooseBranch,
+        reduceTyFamApp_maybe,
 
         -- Flattening
         flattenTys
@@ -39,17 +38,16 @@ module FamInstEnv (
 
 #include "HsVersions.h"
 
-import InstEnv
 import Unify
 import Type
-import TcType ( orphNamesOfTypes )
-import TypeRep
+import TyCoRep
 import TyCon
 import Coercion
 import CoAxiom
 import VarSet
 import VarEnv
 import Name
+import PrelNames ( eqPrimTyConKey )
 import UniqFM
 import Outputable
 import Maybes
@@ -59,8 +57,11 @@ import Util
 import Var
 import Pair
 import SrcLoc
-import NameSet
 import FastString
+import MonadUtils
+import Control.Monad
+import Data.Function ( on )
+import Data.List( mapAccumL )
 
 {-
 ************************************************************************
@@ -89,6 +90,12 @@ data FamInst  -- See Note [FamInsts and CoAxioms]
   = FamInst { fi_axiom  :: CoAxiom Unbranched -- The new coercion axiom
                                               -- introduced by this family
                                               -- instance
+                 -- INVARIANT: apart from freshening (see below)
+                 --    fi_tvs = cab_tvs of the (single) axiom branch
+                 --    fi_cvs = cab_cvs ...ditto...
+                 --    fi_tys = cab_lhs ...ditto...
+                 --    fi_rhs = cab_rhs ...ditto...
+
             , fi_flavor :: FamFlavor
 
             -- Everything below here is a redundant,
@@ -102,21 +109,62 @@ data FamInst  -- See Note [FamInsts and CoAxioms]
                 -- INVARIANT: fi_tcs = roughMatchTcs fi_tys
 
             -- Used for "proper matching"; ditto
-            , fi_tvs    :: [TyVar]      -- Template tyvars for full match
-                                 -- Like ClsInsts, these variables are always
-                                 -- fresh. See Note [Template tyvars are fresh]
-                                 -- in InstEnv
-                                 -- INVARIANT: fi_tvs = coAxiomTyVars fi_axiom
+            , fi_tvs :: [TyVar]      -- Template tyvars for full match
+            , fi_cvs :: [CoVar]      -- Template covars for full match
+                 -- Like ClsInsts, these variables are always fresh
+                 -- See Note [Template tyvars are fresh] in InstEnv
 
-            , fi_tys    :: [Type]       --   and its arg types
+            , fi_tys    :: [Type]       --   The LHS type patterns
+            -- May be eta-reduced; see Note [Eta reduction for data families]
 
-            , fi_rhs    :: Type         --   the RHS, with its freshened vars
+            , fi_rhs :: Type         --   the RHS, with its freshened vars
             }
 
 data FamFlavor
   = SynFamilyInst         -- A synonym family
   | DataFamilyInst TyCon  -- A data family, with its representation TyCon
 
+{- Note [Eta reduction for data families]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider this
+   data family T a b :: *
+   newtype instance T Int a = MkT (IO a) deriving( Monad )
+We'd like this to work.
+
+From the 'newtype instance' you might think we'd get:
+   newtype TInt a = MkT (IO a)
+   axiom ax1 a :: T Int a ~ TInt a   -- The newtype-instance part
+   axiom ax2 a :: TInt a ~ IO a      -- The newtype part
+
+But now what can we do?  We have this problem
+   Given:   d  :: Monad IO
+   Wanted:  d' :: Monad (T Int) = d |> ????
+What coercion can we use for the ???
+
+Solution: eta-reduce both axioms, thus:
+   axiom ax1 :: T Int ~ TInt
+   axiom ax2 :: TInt ~ IO
+Now
+   d' = d |> Monad (sym (ax2 ; ax1))
+
+This eta reduction happens for data instances as well as newtype
+instances. Here we want to eta-reduce the data family axiom.
+All this is done in TcInstDcls.tcDataFamInstDecl.
+
+See also Note [Newtype eta] in TyCon.
+
+Bottom line:
+  For a FamInst with fi_flavour = DataFamilyInst rep_tc,
+  - fi_tvs may be shorter than tyConTyVars of rep_tc
+  - fi_tys may be shorter than tyConArity of the family tycon
+       i.e. LHS is unsaturated
+  - fi_rhs will be (rep_tc fi_tvs)
+       i.e. RHS is un-saturated
+
+  But when fi_flavour = SynFamilyInst,
+  - fi_tys has the exact arity of the family tycon
+-}
+
 -- Obtain the axiom of a family instance
 famInstAxiom :: FamInst -> CoAxiom Unbranched
 famInstAxiom = fi_axiom
@@ -170,11 +218,13 @@ instance Outputable FamInst where
 --     See pprTyThing.pprFamInst for printing for the user
 pprFamInst :: FamInst -> SDoc
 pprFamInst famInst
-  = hang (pprFamInstHdr famInst)
-       2 (vcat [ ifPprDebug (ptext (sLit "Coercion axiom:") <+> ppr ax)
-               , ifPprDebug (ptext (sLit "RHS:") <+> ppr (famInstRHS famInst)) ])
+  = hang (pprFamInstHdr famInst) 2 (ifPprDebug debug_stuff)
   where
     ax = fi_axiom famInst
+    debug_stuff = vcat [ text "Coercion axiom:" <+> ppr ax
+                       , text "Tvs:" <+> ppr (fi_tvs famInst)
+                       , text "LHS:" <+> ppr (fi_tys famInst)
+                       , text "RHS:" <+> ppr (fi_rhs famInst) ]
 
 pprFamInstHdr :: FamInst -> SDoc
 pprFamInstHdr fi@(FamInst {fi_flavor = flavor})
@@ -184,7 +234,7 @@ pprFamInstHdr fi@(FamInst {fi_flavor = flavor})
     -- For *top level* type instances, say "type instance T Int = blah"
     pp_instance
       | isTyConAssoc fam_tc = empty
-      | otherwise           = ptext (sLit "instance")
+      | otherwise           = text "instance"
 
     (fam_tc, etad_lhs_tys) = famInstSplitLHS fi
     vanilla_pp_head = pprTypeApp fam_tc etad_lhs_tys
@@ -206,12 +256,12 @@ pprFamInstHdr fi@(FamInst {fi_flavor = flavor})
             = vanilla_pp_head
 
     pprTyConSort = case flavor of
-                     SynFamilyInst        -> ptext (sLit "type")
+                     SynFamilyInst        -> text "type"
                      DataFamilyInst tycon
-                       | isDataTyCon     tycon -> ptext (sLit "data")
-                       | isNewTyCon      tycon -> ptext (sLit "newtype")
-                       | isAbstractTyCon tycon -> ptext (sLit "data")
-                       | otherwise             -> ptext (sLit "WEIRD") <+> ppr tycon
+                       | isDataTyCon     tycon -> text "data"
+                       | isNewTyCon      tycon -> text "newtype"
+                       | isAbstractTyCon tycon -> text "data"
+                       | otherwise             -> text "WEIRD" <+> ppr tycon
 
 pprFamInsts :: [FamInst] -> SDoc
 pprFamInsts finsts = vcat (map pprFamInst finsts)
@@ -245,16 +295,17 @@ mkImportedFamInst fam mb_tcs axiom
       fi_fam    = fam,
       fi_tcs    = mb_tcs,
       fi_tvs    = tvs,
+      fi_cvs    = cvs,
       fi_tys    = tys,
       fi_rhs    = rhs,
       fi_axiom  = axiom,
       fi_flavor = flavor }
   where
      -- See Note [Lazy axiom match]
-     ~(CoAxiom { co_ax_branches =
-       ~(FirstBranch ~(CoAxBranch { cab_lhs = tys
-                                  , cab_tvs = tvs
-                                  , cab_rhs = rhs })) }) = axiom
+     ~(CoAxBranch { cab_lhs = tys
+                  , cab_tvs = tvs
+                  , cab_cvs = cvs
+                  , cab_rhs = rhs }) = coAxiomSingleBranch axiom
 
          -- Derive the flavor for an imported FamInst rather disgustingly
          -- Maybe we should store it in the IfaceFamInst?
@@ -306,9 +357,8 @@ Then we get a data type for each instance, and an axiom:
    axiom ax7   :: T Int ~ TInt   -- Eta-reduced
    axiom ax8 a :: T Bool [a] ~ TBoolList a
 
-These two axioms for T, one with one pattern, one with two.  The reason
-for this eta-reduction is decribed in TcInstDcls
-   Note [Eta reduction for data family axioms]
+These two axioms for T, one with one pattern, one with two;
+see Note [Eta reduction for data families]
 -}
 
 type FamInstEnv = UniqFM FamilyInstEnv  -- Maps a family to its instances
@@ -321,7 +371,7 @@ newtype FamilyInstEnv
   = FamIE [FamInst]     -- The instances for a particular family, in any order
 
 instance Outputable FamilyInstEnv where
-  ppr (FamIE fs) = ptext (sLit "FamIE") <+> vcat (map ppr fs)
+  ppr (FamIE fs) = text "FamIE" <+> vcat (map ppr fs)
 
 -- INVARIANTS:
 --  * The fs_tvs are distinct in each FamInst
@@ -344,21 +394,6 @@ familyInstances (pkg_fie, home_fie) fam
                 Just (FamIE insts) -> insts
                 Nothing                      -> []
 
--- | Collects the names of the concrete types and type constructors that
--- make up the LHS of a type family instance, including the family
--- name itself.
---
--- For instance, given `type family Foo a b`:
--- `type instance Foo (F (G (H a))) b = ...` would yield [Foo,F,G,H]
---
--- Used in the implementation of ":info" in GHCi.
-orphNamesOfFamInst :: FamInst -> NameSet
-orphNamesOfFamInst fam_inst
-  = orphNamesOfTypes (concat (brListMap cab_lhs (coAxiomBranches axiom)))
-    `extendNameSet` getName (coAxiomTyCon axiom)
-  where
-    axiom = fi_axiom fam_inst
-
 extendFamInstEnvList :: FamInstEnv -> [FamInst] -> FamInstEnv
 extendFamInstEnvList inst_env fis = foldl extendFamInstEnv inst_env fis
 
@@ -383,18 +418,16 @@ identicalFamInstHead :: FamInst -> FamInst -> Bool
 -- Used for overriding in GHCi
 identicalFamInstHead (FamInst { fi_axiom = ax1 }) (FamInst { fi_axiom = ax2 })
   =  coAxiomTyCon ax1 == coAxiomTyCon ax2
-  && brListLength brs1 == brListLength brs2
-  && and (brListZipWith identical_branch brs1 brs2)
+  && numBranches brs1 == numBranches brs2
+  && and ((zipWith identical_branch `on` fromBranches) brs1 brs2)
   where
     brs1 = coAxiomBranches ax1
     brs2 = coAxiomBranches ax2
 
     identical_branch br1 br2
-      =  isJust (tcMatchTys tvs1 lhs1 lhs2)
-      && isJust (tcMatchTys tvs2 lhs2 lhs1)
+      =  isJust (tcMatchTys lhs1 lhs2)
+      && isJust (tcMatchTys lhs2 lhs1)
       where
-        tvs1 = mkVarSet (coAxBranchTyVars br1)
-        tvs2 = mkVarSet (coAxBranchTyVars br2)
         lhs1 = coAxBranchLHS br1
         lhs2 = coAxBranchLHS br2
 
@@ -479,10 +512,11 @@ irrelevant (clause 1 of compatible) or benign (clause 2 of compatible).
 compatibleBranches :: CoAxBranch -> CoAxBranch -> Bool
 compatibleBranches (CoAxBranch { cab_lhs = lhs1, cab_rhs = rhs1 })
                    (CoAxBranch { cab_lhs = lhs2, cab_rhs = rhs2 })
-  = case tcUnifyTysFG instanceBindFun lhs1 lhs2 of
+  = case tcUnifyTysFG (const BindMe) lhs1 lhs2 of
       SurelyApart -> True
       Unifiable subst
-        | Type.substTy subst rhs1 `eqType` Type.substTy subst rhs2
+        | Type.substTyAddInScope subst rhs1 `eqType`
+          Type.substTyAddInScope subst rhs2
         -> True
       _ -> False
 
@@ -527,19 +561,19 @@ injectiveBranches injectivity
 -- takes a CoAxiom with unknown branch incompatibilities and computes
 -- the compatibilities
 -- See Note [Storing compatibility] in CoAxiom
-computeAxiomIncomps :: CoAxiom br -> CoAxiom br
-computeAxiomIncomps ax@(CoAxiom { co_ax_branches = branches })
-  = ax { co_ax_branches = go [] branches }
+computeAxiomIncomps :: [CoAxBranch] -> [CoAxBranch]
+computeAxiomIncomps branches
+  = snd (mapAccumL go [] branches)
   where
-    go :: [CoAxBranch] -> BranchList CoAxBranch br -> BranchList CoAxBranch br
-    go prev_branches (FirstBranch br)
-      = FirstBranch (br { cab_incomps = mk_incomps br prev_branches })
-    go prev_branches (NextBranch br tail)
-      = let br' = br { cab_incomps = mk_incomps br prev_branches } in
-        NextBranch br' (go (br' : prev_branches) tail)
+    go :: [CoAxBranch] -> CoAxBranch -> ([CoAxBranch], CoAxBranch)
+    go prev_brs cur_br
+       = (cur_br : prev_brs, new_br)
+       where
+         new_br = cur_br { cab_incomps = mk_incomps prev_brs cur_br }
 
-    mk_incomps :: CoAxBranch -> [CoAxBranch] -> [CoAxBranch]
-    mk_incomps br = filter (not . compatibleBranches br)
+    mk_incomps :: [CoAxBranch] -> CoAxBranch -> [CoAxBranch]
+    mk_incomps prev_brs cur_br
+       = filter (not . compatibleBranches cur_br) prev_brs
 
 {-
 ************************************************************************
@@ -547,6 +581,8 @@ computeAxiomIncomps ax@(CoAxiom { co_ax_branches = branches })
            Constructing axioms
     These functions are here because tidyType / tcUnifyTysFG
     are not available in CoAxiom
+
+    Also computeAxiomIncomps is too sophisticated for CoAxiom
 *                                                                      *
 ************************************************************************
 
@@ -559,32 +595,35 @@ Instead we must tidy those kind variables.  See Trac #7524.
 
 -- all axiom roles are Nominal, as this is only used with type families
 mkCoAxBranch :: [TyVar] -- original, possibly stale, tyvars
+             -> [CoVar] -- possibly stale covars
              -> [Type]  -- LHS patterns
              -> Type    -- RHS
+             -> [Role]
              -> SrcSpan
              -> CoAxBranch
-mkCoAxBranch tvs lhs rhs loc
+mkCoAxBranch tvs cvs lhs rhs roles loc
   = CoAxBranch { cab_tvs     = tvs1
+               , cab_cvs     = cvs1
                , cab_lhs     = tidyTypes env lhs
-               , cab_roles   = map (const Nominal) tvs1
+               , cab_roles   = roles
                , cab_rhs     = tidyType  env rhs
                , cab_loc     = loc
                , cab_incomps = placeHolderIncomps }
   where
-    (env, tvs1) = tidyTyVarBndrs emptyTidyEnv tvs
+    (env1, tvs1) = tidyTyCoVarBndrs emptyTidyEnv tvs
+    (env,  cvs1) = tidyTyCoVarBndrs env1         cvs
     -- See Note [Tidy axioms when we build them]
 
 -- all of the following code is here to avoid mutual dependencies with
 -- Coercion
 mkBranchedCoAxiom :: Name -> TyCon -> [CoAxBranch] -> CoAxiom Branched
 mkBranchedCoAxiom ax_name fam_tc branches
-  = computeAxiomIncomps $
-    CoAxiom { co_ax_unique   = nameUnique ax_name
+  = CoAxiom { co_ax_unique   = nameUnique ax_name
             , co_ax_name     = ax_name
             , co_ax_tc       = fam_tc
             , co_ax_role     = Nominal
             , co_ax_implicit = False
-            , co_ax_branches = toBranchList branches }
+            , co_ax_branches = manyBranches (computeAxiomIncomps branches) }
 
 mkUnbranchedCoAxiom :: Name -> TyCon -> CoAxBranch -> CoAxiom Unbranched
 mkUnbranchedCoAxiom ax_name fam_tc branch
@@ -593,23 +632,42 @@ mkUnbranchedCoAxiom ax_name fam_tc branch
             , co_ax_tc       = fam_tc
             , co_ax_role     = Nominal
             , co_ax_implicit = False
-            , co_ax_branches = FirstBranch (branch { cab_incomps = [] }) }
+            , co_ax_branches = unbranched (branch { cab_incomps = [] }) }
 
 mkSingleCoAxiom :: Role -> Name
-                -> [TyVar] -> TyCon -> [Type] -> Type
+                -> [TyVar] -> [CoVar] -> TyCon -> [Type] -> Type
                 -> CoAxiom Unbranched
 -- Make a single-branch CoAxiom, incluidng making the branch itself
 -- Used for both type family (Nominal) and data family (Representational)
 -- axioms, hence passing in the Role
-mkSingleCoAxiom role ax_name tvs fam_tc lhs_tys rhs_ty
+mkSingleCoAxiom role ax_name tvs cvs fam_tc lhs_tys rhs_ty
   = CoAxiom { co_ax_unique   = nameUnique ax_name
             , co_ax_name     = ax_name
             , co_ax_tc       = fam_tc
             , co_ax_role     = role
             , co_ax_implicit = False
-            , co_ax_branches = FirstBranch (branch { cab_incomps = [] }) }
+            , co_ax_branches = unbranched (branch { cab_incomps = [] }) }
   where
-    branch = mkCoAxBranch tvs lhs_tys rhs_ty (getSrcSpan ax_name)
+    branch = mkCoAxBranch tvs cvs lhs_tys rhs_ty
+                          (map (const Nominal) tvs)
+                          (getSrcSpan ax_name)
+
+-- | Create a coercion constructor (axiom) suitable for the given
+--   newtype 'TyCon'. The 'Name' should be that of a new coercion
+--   'CoAxiom', the 'TyVar's the arguments expected by the @newtype@ and
+--   the type the appropriate right hand side of the @newtype@, with
+--   the free variables a subset of those 'TyVar's.
+mkNewTypeCoAxiom :: Name -> TyCon -> [TyVar] -> [Role] -> Type -> CoAxiom Unbranched
+mkNewTypeCoAxiom name tycon tvs roles rhs_ty
+  = CoAxiom { co_ax_unique   = nameUnique name
+            , co_ax_name     = name
+            , co_ax_implicit = True  -- See Note [Implicit axioms] in TyCon
+            , co_ax_role     = Representational
+            , co_ax_tc       = tycon
+            , co_ax_branches = unbranched (branch { cab_incomps = [] }) }
+  where
+    branch = mkCoAxBranch tvs [] (mkTyVarTys tvs) rhs_ty
+                          roles (getSrcSpan name)
 
 {-
 ************************************************************************
@@ -640,13 +698,15 @@ we return the matching instance '(FamInst{.., fi_tycon = :R42T}, Int)'.
 -- and the list of types the axiom should be applied to
 data FamInstMatch = FamInstMatch { fim_instance :: FamInst
                                  , fim_tys      :: [Type]
+                                 , fim_cos      :: [Coercion]
                                  }
   -- See Note [Over-saturated matches]
 
 instance Outputable FamInstMatch where
   ppr (FamInstMatch { fim_instance = inst
-                    , fim_tys      = tys })
-    = ptext (sLit "match with") <+> parens (ppr inst) <+> ppr tys
+                    , fim_tys      = tys
+                    , fim_cos      = cos })
+    = text "match with" <+> parens (ppr inst) <+> ppr tys <+> ppr cos
 
 lookupFamInstEnvByTyCon :: FamInstEnvs -> TyCon -> [FamInst]
 lookupFamInstEnvByTyCon (pkg_ie, home_ie) fam_tc
@@ -665,7 +725,7 @@ lookupFamInstEnv
 lookupFamInstEnv
    = lookup_fam_inst_env match
    where
-     match _ tpl_tvs tpl_tys tys = tcMatchTys tpl_tvs tpl_tys tys
+     match _ _ tpl_tys tys = tcMatchTys tpl_tys tys
 
 lookupFamInstEnvConflicts
     :: FamInstEnvs
@@ -685,7 +745,7 @@ lookupFamInstEnvConflicts envs fam_inst@(FamInst { fi_axiom = new_axiom })
         -- In example above,   fam tys' = F [b]
 
     my_unify (FamInst { fi_axiom = old_axiom }) tpl_tvs tpl_tys _
-       = ASSERT2( tyVarsOfTypes tys `disjointVarSet` tpl_tvs,
+       = ASSERT2( tyCoVarsOfTypes tys `disjointVarSet` tpl_tvs,
                   (ppr fam <+> ppr tys) $$
                   (ppr tpl_tvs <+> ppr tpl_tys) )
                 -- Unification will break badly if the variables overlap
@@ -816,80 +876,11 @@ lookupFamInstEnvInjectivityConflicts injList (pkg_ie, home_ie)
 
       lookup_inj_fam_conflicts ie
           | isOpenFamilyTyCon fam, Just (FamIE insts) <- lookupUFM ie fam
-          = map (brFromUnbranchedSingleton . co_ax_branches . fi_axiom) $
+          = map (coAxiomSingleBranch . fi_axiom) $
             filter isInjConflict insts
           | otherwise = []
 
 
--- | Return a list of type variables that the function is injective in and that
--- do not appear on injective positions in the RHS of a family instance
--- declaration.
-unusedInjTvsInRHS :: [Bool] -> [Type] -> Type -> TyVarSet
--- INVARIANT: [Bool] list contains at least one True value
--- See Note [Verifying injectivity annotation]. This function implements fourth
--- check described there.
--- In theory, instead of implementing this whole check in this way, we could
--- attempt to unify equation with itself.  We would reject exactly the same
--- equations but this method gives us more precise error messages by returning
--- precise names of variables that are not mentioned in the RHS.
-unusedInjTvsInRHS injList lhs rhs =
-  injLHSVars `minusVarSet` injRhsVars
-    where
-      -- set of type and kind variables in which type family is injective
-      injLHSVars = tyVarsOfTypes (filterByList injList lhs)
-
-      -- set of type variables appearing in the RHS on an injective position.
-      -- For all returned variables we assume their associated kind variables
-      -- also appear in the RHS.
-      injRhsVars = closeOverKinds $ collectInjVars rhs
-
-      -- Collect all type variables that are either arguments to a type
-      -- constructor or to injective type families.
-      collectInjVars :: Type -> VarSet
-      collectInjVars ty | Just (ty1, ty2) <- splitAppTy_maybe ty
-        = collectInjVars ty1 `unionVarSet` collectInjVars ty2
-      collectInjVars (TyVarTy v)
-        = unitVarSet v
-      collectInjVars (TyConApp tc tys)
-        | isTypeFamilyTyCon tc = collectInjTFVars tys
-                                                 (familyTyConInjectivityInfo tc)
-        | otherwise            = mapUnionVarSet collectInjVars tys
-      collectInjVars (LitTy {})
-        = emptyVarSet
-      collectInjVars (FunTy arg res)
-        = collectInjVars arg `unionVarSet` collectInjVars res
-      collectInjVars (AppTy fun arg)
-        = collectInjVars fun `unionVarSet` collectInjVars arg
-      -- no forall types in the RHS of a type family
-      collectInjVars (ForAllTy _ _)    =
-          panic "unusedInjTvsInRHS.collectInjVars"
-
-      collectInjTFVars :: [Type] -> Injectivity -> VarSet
-      collectInjTFVars _ NotInjective
-          = emptyVarSet
-      collectInjTFVars tys (Injective injList)
-          = mapUnionVarSet collectInjVars (filterByList injList tys)
-
--- | Is type headed by a type family application?
-isTFHeaded :: Type -> Bool
--- See Note [Verifying injectivity annotation]. This function implements third
--- check described there.
-isTFHeaded ty | Just ty' <- tcView ty
-              = isTFHeaded ty'
-isTFHeaded ty | (TyConApp tc args) <- ty
-              , isTypeFamilyTyCon tc
-              = tyConArity tc == length args
-isTFHeaded _  = False
-
--- | If a RHS is a bare type variable return a set of LHS patterns that are not
--- bare type variables.
-bareTvInRHSViolated :: [Type] -> Type -> [Type]
--- See Note [Verifying injectivity annotation]. This function implements second
--- check described there.
-bareTvInRHSViolated pats rhs | isTyVarTy rhs
-   = filter (not . isTyVarTy) pats
-bareTvInRHSViolated _ _ = []
-
 --------------------------------------------------------------------------------
 --                    Type family overlap checking bits                       --
 --------------------------------------------------------------------------------
@@ -915,7 +906,7 @@ Note [Family instance overlap conflicts]
 type MatchFun =  FamInst                -- The FamInst template
               -> TyVarSet -> [Type]     --   fi_tvs, fi_tys of that FamInst
               -> [Type]                 -- Target to match against
-              -> Maybe TvSubst
+              -> Maybe TCvSubst
 
 lookup_fam_inst_env'          -- The worker, local to this module
     :: MatchFun
@@ -930,8 +921,8 @@ lookup_fam_inst_env' match_fun ie fam match_tys
   where
 
     find [] = []
-    find (item@(FamInst { fi_tcs = mb_tcs, fi_tvs = tpl_tvs,
-                          fi_tys = tpl_tys }) : rest)
+    find (item@(FamInst { fi_tcs = mb_tcs, fi_tvs = tpl_tvs, fi_cvs = tpl_cvs
+                        , fi_tys = tpl_tys }) : rest)
         -- Fast check for no match, uses the "rough match" fields
       | instanceCantMatch rough_tcs mb_tcs
       = find rest
@@ -939,7 +930,10 @@ lookup_fam_inst_env' match_fun ie fam match_tys
         -- Proper check
       | Just subst <- match_fun item (mkVarSet tpl_tvs) tpl_tys match_tys1
       = (FamInstMatch { fim_instance = item
-                      , fim_tys      = substTyVars subst tpl_tvs `chkAppend` match_tys2 })
+                      , fim_tys      = substTyVars subst tpl_tvs `chkAppend` match_tys2
+                      , fim_cos      = ASSERT( all (isJust . lookupCoVar subst) tpl_cvs )
+                                       substCoVars subst tpl_cvs
+                      })
         : find rest
 
         -- No match => try next
@@ -1001,8 +995,8 @@ We handle data families and type families separately here:
 
  * For data family instances, though, we need to re-split for each
    instance, because the breakdown might be different for each
-   instance.  Why?  Because of eta reduction; see Note [Eta reduction
-   for data family axioms] in TcInstDcls.
+   instance.  Why?  Because of eta reduction; see
+   Note [Eta reduction for data families].
 -}
 
 -- checks if one LHS is dominated by a list of other branches
@@ -1019,8 +1013,8 @@ isDominatedBy branch branches
   = or $ map match branches
     where
       lhs = coAxBranchLHS branch
-      match (CoAxBranch { cab_tvs = tvs, cab_lhs = tys })
-        = isJust $ tcMatchTys (mkVarSet tvs) tys lhs
+      match (CoAxBranch { cab_lhs = tys })
+        = isJust $ tcMatchTys tys lhs
 
 {-
 ************************************************************************
@@ -1047,68 +1041,76 @@ reduceTyFamApp_maybe :: FamInstEnvs
 --
 -- The TyCon can be oversaturated.
 -- Works on both open and closed families
-
+--
+-- Always returns a *homogeneous* coercion -- type family reductions are always
+-- homogeneous
 reduceTyFamApp_maybe envs role tc tys
   | Phantom <- role
   = Nothing
 
   | case role of
-       Representational -> isOpenFamilyTyCon    tc
-       _                -> isOpenTypeFamilyTyCon tc
+      Representational -> isOpenFamilyTyCon     tc
+      _                -> isOpenTypeFamilyTyCon tc
        -- If we seek a representational coercion
        -- (e.g. the call in topNormaliseType_maybe) then we can
        -- unwrap data families as well as type-synonym families;
        -- otherwise only type-synonym families
-  , FamInstMatch { fim_instance = fam_inst
-                 , fim_tys =      inst_tys } : _ <- lookupFamInstEnv envs tc tys
+  , FamInstMatch { fim_instance = FamInst { fi_axiom = ax }
+                 , fim_tys      = inst_tys
+                 , fim_cos      = inst_cos } : _ <- lookupFamInstEnv envs tc tys
       -- NB: Allow multiple matches because of compatible overlap
-  = let ax     = famInstAxiom fam_inst
-        co     = mkUnbranchedAxInstCo role ax inst_tys
-        ty     = pSnd (coercionKind co)
+
+  = let co = mkUnbranchedAxInstCo role ax inst_tys inst_cos
+        ty = pSnd (coercionKind co)
     in Just (co, ty)
 
   | Just ax <- isClosedSynFamilyTyConWithAxiom_maybe tc
-  , Just (ind, inst_tys) <- chooseBranch ax tys
-  = let co     = mkAxInstCo role ax ind inst_tys
-        ty     = pSnd (coercionKind co)
+  , Just (ind, inst_tys, inst_cos) <- chooseBranch ax tys
+  = let co = mkAxInstCo role ax ind inst_tys inst_cos
+        ty = pSnd (coercionKind co)
     in Just (co, ty)
 
   | Just ax           <- isBuiltInSynFamTyCon_maybe tc
   , Just (coax,ts,ty) <- sfMatchFam ax tys
-  = let co = mkAxiomRuleCo coax ts []
+  = let co = mkAxiomRuleCo coax (zipWith mkReflCo (coaxrAsmpRoles coax) ts)
     in Just (co, ty)
 
   | otherwise
   = Nothing
 
 -- The axiom can be oversaturated. (Closed families only.)
-chooseBranch :: CoAxiom Branched -> [Type] -> Maybe (BranchIndex, [Type])
+chooseBranch :: CoAxiom Branched -> [Type]
+             -> Maybe (BranchIndex, [Type], [Coercion])  -- found match, with args
 chooseBranch axiom tys
   = do { let num_pats = coAxiomNumPats axiom
              (target_tys, extra_tys) = splitAt num_pats tys
              branches = coAxiomBranches axiom
-       ; (ind, inst_tys) <- findBranch (fromBranchList branches) target_tys
-       ; return (ind, inst_tys ++ extra_tys) }
+       ; (ind, inst_tys, inst_cos)
+           <- findBranch (fromBranches branches) target_tys
+       ; return ( ind, inst_tys `chkAppend` extra_tys, inst_cos ) }
 
 -- The axiom must *not* be oversaturated
 findBranch :: [CoAxBranch]             -- branches to check
            -> [Type]                   -- target types
-           -> Maybe (BranchIndex, [Type])
+           -> Maybe (BranchIndex, [Type], [Coercion])
+    -- coercions relate requested types to returned axiom LHS at role N
 findBranch branches target_tys
   = go 0 branches
   where
-    go ind (branch@(CoAxBranch { cab_tvs = tpl_tvs, cab_lhs = tpl_lhs
+    go ind (branch@(CoAxBranch { cab_tvs = tpl_tvs, cab_cvs = tpl_cvs
+                               , cab_lhs = tpl_lhs
                                , cab_incomps = incomps }) : rest)
       = let in_scope = mkInScopeSet (unionVarSets $
-                            map (tyVarsOfTypes . coAxBranchLHS) incomps)
+                            map (tyCoVarsOfTypes . coAxBranchLHS) incomps)
             -- See Note [Flattening] below
             flattened_target = flattenTys in_scope target_tys
-        in case tcMatchTys (mkVarSet tpl_tvs) tpl_lhs target_tys of
+        in case tcMatchTys tpl_lhs target_tys of
         Just subst -- matching worked. now, check for apartness.
           |  apartnessCheck flattened_target branch
           -> -- matching worked & we're apart from all incompatible branches.
              -- success
-             Just (ind, substTyVars subst tpl_tvs)
+             ASSERT( all (isJust . lookupCoVar subst) tpl_cvs )
+             Just (ind, substTyVars subst tpl_tvs, substCoVars subst tpl_cvs)
 
         -- failure. keep looking
         _ -> go (ind+1) rest
@@ -1129,7 +1131,7 @@ apartnessCheck :: [Type]     -- ^ /flattened/ target arguments. Make sure
                -> Bool       -- ^ True <=> equation can fire
 apartnessCheck flattened_target (CoAxBranch { cab_incomps = incomps })
   = all (isSurelyApart
-         . tcUnifyTysFG instanceBindFun flattened_target
+         . tcUnifyTysFG (const BindMe) flattened_target
          . coAxBranchLHS) incomps
   where
     isSurelyApart SurelyApart = True
@@ -1141,6 +1143,46 @@ apartnessCheck flattened_target (CoAxBranch { cab_incomps = incomps })
                 Looking up a family instance
 *                                                                      *
 ************************************************************************
+
+Note [Normalising types]
+~~~~~~~~~~~~~~~~~~~~~~~~
+The topNormaliseType function removes all occurrences of type families
+and newtypes from the top-level structure of a type. normaliseTcApp does
+the type family lookup and is fairly straightforward. normaliseType is
+a little more involved.
+
+The complication comes from the fact that a type family might be used in the
+kind of a variable bound in a forall. We wish to remove this type family
+application, but that means coming up with a fresh variable (with the new
+kind). Thus, we need a substitution to be built up as we recur through the
+type. However, an ordinary TCvSubst just won't do: when we hit a type variable
+whose kind has changed during normalisation, we need both the new type
+variable *and* the coercion. We could conjure up a new VarEnv with just this
+property, but a usable substitution environment already exists:
+LiftingContexts from the liftCoSubst family of functions, defined in Coercion.
+A LiftingContext maps a type variable to a coercion and a coercion variable to
+a pair of coercions. Let's ignore coercion variables for now. Because the
+coercion a type variable maps to contains the destination type (via
+coercionKind), we don't need to store that destination type separately. Thus,
+a LiftingContext has what we need: a map from type variables to (Coercion,
+Type) pairs.
+
+We also benefit because we can piggyback on the liftCoSubstVarBndr function to
+deal with binders. However, I had to modify that function to work with this
+application. Thus, we now have liftCoSubstVarBndrCallback, which takes
+a function used to process the kind of the binder. We don't wish
+to lift the kind, but instead normalise it. So, we pass in a callback function
+that processes the kind of the binder.
+
+After that brilliant explanation of all this, I'm sure you've forgotten the
+dangling reference to coercion variables. What do we do with those? Nothing at
+all. The point of normalising types is to remove type family applications, but
+there's no sense in removing these from coercions. We would just get back a
+new coercion witnessing the equality between the same types as the original
+coercion. Because coercions are irrelevant anyway, there is no point in doing
+this. So, whenever we encounter a coercion, we just say that it won't change.
+That's what the CoercionTy case is doing within normalise_type.
+
 -}
 
 topNormaliseType :: FamInstEnvs -> Type -> Type
@@ -1172,70 +1214,178 @@ topNormaliseType_maybe env ty
 
     tyFamStepper rec_nts tc tys  -- Try to step a type/data familiy
       = let (args_co, ntys) = normaliseTcArgs env Representational tc tys in
+          -- NB: It's OK to use normaliseTcArgs here instead of
+          -- normalise_tc_args (which takes the LiftingContext described
+          -- in Note [Normalising types]) because the reduceTyFamApp below
+          -- works only at top level. We'll never recur in this function
+          -- after reducing the kind of a bound tyvar.
+
         case reduceTyFamApp_maybe env Representational tc ntys of
           Just (co, rhs) -> NS_Step rec_nts rhs (args_co `mkTransCo` co)
-          Nothing        -> NS_Done
+          _              -> NS_Done
 
 ---------------
 normaliseTcApp :: FamInstEnvs -> Role -> TyCon -> [Type] -> (Coercion, Type)
 -- See comments on normaliseType for the arguments of this function
 normaliseTcApp env role tc tys
-  | isTypeSynonymTyCon tc
-  , Just (tenv, rhs, ntys') <- expandSynTyCon_maybe tc ntys
-  , (co2, ninst_rhs) <- normaliseType env role (Type.substTy (mkTopTvSubst tenv) rhs)
-  = if isReflCo co2 then (args_co,                 mkTyConApp tc ntys)
-                    else (args_co `mkTransCo` co2, mkAppTys ninst_rhs ntys')
-
-  | Just (first_co, ty') <- reduceTyFamApp_maybe env role tc ntys
-  , (rest_co,nty) <- normaliseType env role ty'
-  = (args_co `mkTransCo` first_co `mkTransCo` rest_co, nty)
-
-  | otherwise   -- No unique matching family instance exists;
+  = initNormM env role (tyCoVarsOfTypes tys) $
+    normalise_tc_app tc tys
+
+-- See Note [Normalising types] about the LiftingContext
+normalise_tc_app :: TyCon -> [Type] -> NormM (Coercion, Type)
+normalise_tc_app tc tys
+  = do { (args_co, ntys) <- normalise_tc_args tc tys
+       ; case expandSynTyCon_maybe tc ntys of
+         { Just (tenv, rhs, ntys') ->
+           do { (co2, ninst_rhs)
+                  <- normalise_type (substTy (mkTvSubstPrs tenv) rhs)
+              ; return $
+                if isReflCo co2
+                then (args_co,                 mkTyConApp tc ntys)
+                else (args_co `mkTransCo` co2, mkAppTys ninst_rhs ntys') }
+         ; Nothing ->
+    do { env <- getEnv
+       ; role <- getRole
+       ; case reduceTyFamApp_maybe env role tc ntys of
+           Just (first_co, ty')
+             -> do { (rest_co,nty) <- normalise_type ty'
+                   ; return ( args_co `mkTransCo` first_co `mkTransCo` rest_co
+                            , nty ) }
+           _ -> -- No unique matching family instance exists;
                 -- we do not do anything
-  = (args_co, mkTyConApp tc ntys)
-
-  where
-    (args_co, ntys) = normaliseTcArgs env role tc tys
-
+                return (args_co, mkTyConApp tc ntys) }}}
 
 ---------------
-normaliseTcArgs :: FamInstEnvs            -- environment with family instances
-                 -> Role                   -- desired role of output coercion
-                 -> TyCon -> [Type]        -- tc tys
-                 -> (Coercion, [Type])     -- (co, new_tys), where
-                                           -- co :: tc tys ~ tc new_tys
+-- | Normalise arguments to a tycon
+normaliseTcArgs :: FamInstEnvs          -- ^ env't with family instances
+                -> Role                 -- ^ desired role of output coercion
+                -> TyCon                -- ^ tc
+                -> [Type]               -- ^ tys
+                -> (Coercion, [Type])   -- ^ co :: tc tys ~ tc new_tys
 normaliseTcArgs env role tc tys
-  = (mkTyConAppCo role tc cois, ntys)
+  = initNormM env role (tyCoVarsOfTypes tys) $
+    normalise_tc_args tc tys
+
+normalise_tc_args :: TyCon -> [Type]             -- tc tys
+                  -> NormM (Coercion, [Type])    -- (co, new_tys), where
+                                                 -- co :: tc tys ~ tc new_tys
+normalise_tc_args tc tys
+  = do { role <- getRole
+       ; (cois, ntys) <- zipWithAndUnzipM normalise_type_role
+                                          tys (tyConRolesX role tc)
+       ; return (mkTyConAppCo role tc cois, ntys) }
   where
-    (cois, ntys) = zipWithAndUnzip (normaliseType env) (tyConRolesX role tc) tys
+    normalise_type_role ty r = withRole r $ normalise_type ty
 
 ---------------
-normaliseType :: FamInstEnvs            -- environment with family instances
-               -> Role                   -- desired role of output coercion
-               -> Type                   -- old type
-               -> (Coercion, Type)       -- (coercion,new type), where
-                                        -- co :: old-type ~ new_type
+normaliseType :: FamInstEnvs
+              -> Role  -- desired role of coercion
+              -> Type -> (Coercion, Type)
+normaliseType env role ty
+  = initNormM env role (tyCoVarsOfType ty) $ normalise_type ty
+
+normalise_type :: Type                     -- old type
+               -> NormM (Coercion, Type)   -- (coercion,new type), where
+                                         -- co :: old-type ~ new_type
 -- Normalise the input type, by eliminating *all* type-function redexes
 -- but *not* newtypes (which are visible to the programmer)
 -- Returns with Refl if nothing happens
+-- Does nothing to newtypes
+-- The returned coercion *must* be *homogeneous*
+-- See Note [Normalising types]
 -- Try to not to disturb type synonyms if possible
 
-normaliseType env role (TyConApp tc tys)
-  = normaliseTcApp env role tc tys
-normaliseType _env role ty@(LitTy {}) = (mkReflCo role ty, ty)
-normaliseType env role (AppTy ty1 ty2)
-  = let (coi1,nty1) = normaliseType env role    ty1
-        (coi2,nty2) = normaliseType env Nominal ty2
-    in  (mkAppCo coi1 coi2, mkAppTy nty1 nty2)
-normaliseType env role (FunTy ty1 ty2)
-  = let (coi1,nty1) = normaliseType env role ty1
-        (coi2,nty2) = normaliseType env role ty2
-    in  (mkFunCo role coi1 coi2, mkFunTy nty1 nty2)
-normaliseType env role (ForAllTy tyvar ty1)
-  = let (coi,nty1) = normaliseType env role ty1
-    in  (mkForAllCo tyvar coi, ForAllTy tyvar nty1)
-normaliseType _  role ty@(TyVarTy _)
-  = (mkReflCo role ty,ty)
+normalise_type
+  = go
+  where
+    go (TyConApp tc tys) = normalise_tc_app tc tys
+    go ty@(LitTy {})     = do { r <- getRole
+                              ; return (mkReflCo r ty, ty) }
+    go (AppTy ty1 ty2)
+      = do { (co,  nty1) <- go ty1
+           ; (arg, nty2) <- withRole Nominal $ go ty2
+           ; return (mkAppCo co arg, mkAppTy nty1 nty2) }
+    go (ForAllTy (Anon ty1) ty2)
+      = do { (co1, nty1) <- go ty1
+           ; (co2, nty2) <- go ty2
+           ; r <- getRole
+           ; return (mkFunCo r co1 co2, mkFunTy nty1 nty2) }
+    go (ForAllTy (Named tyvar vis) ty)
+      = do { (lc', tv', h, ki') <- normalise_tyvar_bndr tyvar
+           ; (co, nty)          <- withLC lc' $ normalise_type ty
+           ; let tv2 = setTyVarKind tv' ki'
+           ; return (mkForAllCo tv' h co, mkNamedForAllTy tv2 vis nty) }
+    go (TyVarTy tv)    = normalise_tyvar tv
+    go (CastTy ty co)
+      = do { (nco, nty) <- go ty
+           ; lc <- getLC
+           ; let co' = substRightCo lc co
+           ; return (castCoercionKind nco co co', mkCastTy nty co') }
+    go (CoercionTy co)
+      = do { lc <- getLC
+           ; r <- getRole
+           ; let right_co = substRightCo lc co
+           ; return ( mkProofIrrelCo r
+                         (liftCoSubst Nominal lc (coercionType co))
+                         co right_co
+                    , mkCoercionTy right_co ) }
+
+normalise_tyvar :: TyVar -> NormM (Coercion, Type)
+normalise_tyvar tv
+  = ASSERT( isTyVar tv )
+    do { lc <- getLC
+       ; r  <- getRole
+       ; return $ case liftCoSubstTyVar lc r tv of
+           Just co -> (co, pSnd $ coercionKind co)
+           Nothing -> (mkReflCo r ty, ty) }
+  where ty = mkTyVarTy tv
+
+normalise_tyvar_bndr :: TyVar -> NormM (LiftingContext, TyVar, Coercion, Kind)
+normalise_tyvar_bndr tv
+  = do { lc1 <- getLC
+       ; env <- getEnv
+       ; let callback lc ki = runNormM (normalise_type ki) env lc Nominal
+       ; return $ liftCoSubstVarBndrCallback callback lc1 tv }
+
+-- | a monad for the normalisation functions, reading 'FamInstEnvs',
+-- a 'LiftingContext', and a 'Role'.
+newtype NormM a = NormM { runNormM ::
+                            FamInstEnvs -> LiftingContext -> Role -> a }
+
+initNormM :: FamInstEnvs -> Role
+          -> TyCoVarSet   -- the in-scope variables
+          -> NormM a -> a
+initNormM env role vars (NormM thing_inside)
+  = thing_inside env lc role
+  where
+    in_scope = mkInScopeSet vars
+    lc       = emptyLiftingContext in_scope
+
+getRole :: NormM Role
+getRole = NormM (\ _ _ r -> r)
+
+getLC :: NormM LiftingContext
+getLC = NormM (\ _ lc _ -> lc)
+
+getEnv :: NormM FamInstEnvs
+getEnv = NormM (\ env _ _ -> env)
+
+withRole :: Role -> NormM a -> NormM a
+withRole r thing = NormM $ \ envs lc _old_r -> runNormM thing envs lc r
+
+withLC :: LiftingContext -> NormM a -> NormM a
+withLC lc thing = NormM $ \ envs _old_lc r -> runNormM thing envs lc r
+
+instance Monad NormM where
+  ma >>= fmb = NormM $ \env lc r ->
+               let a = runNormM ma env lc r in
+               runNormM (fmb a) env lc r
+
+instance Functor NormM where
+  fmap = liftM
+instance Applicative NormM where
+  pure x = NormM $ \ _ _ _ -> x
+  (<*>)  = ap
 
 {-
 ************************************************************************
@@ -1277,11 +1427,17 @@ is!  Flattening as done below ensures this.
 flattenTys is defined here because of module dependencies.
 -}
 
-type FlattenMap = TypeMap TyVar
+data FlattenEnv = FlattenEnv { fe_type_map :: TypeMap TyVar
+                             , fe_subst    :: TCvSubst }
+
+emptyFlattenEnv :: InScopeSet -> FlattenEnv
+emptyFlattenEnv in_scope
+  = FlattenEnv { fe_type_map = emptyTypeMap
+               , fe_subst    = mkEmptyTCvSubst in_scope }
 
 -- See Note [Flattening]
 flattenTys :: InScopeSet -> [Type] -> [Type]
-flattenTys in_scope tys = snd $ coreFlattenTys all_in_scope emptyTypeMap tys
+flattenTys in_scope tys = snd $ coreFlattenTys env tys
   where
     -- when we hit a type function, we replace it with a fresh variable
     -- but, we need to make sure that this fresh variable isn't mentioned
@@ -1289,75 +1445,157 @@ flattenTys in_scope tys = snd $ coreFlattenTys all_in_scope emptyTypeMap tys
     -- a forall. That way, we can ensure consistency both within and outside
     -- of that forall.
     all_in_scope = in_scope `extendInScopeSetSet` allTyVarsInTys tys
+    env          = emptyFlattenEnv all_in_scope
 
-coreFlattenTys :: InScopeSet -> FlattenMap -> [Type] -> (FlattenMap, [Type])
-coreFlattenTys in_scope = go []
+coreFlattenTys :: FlattenEnv -> [Type] -> (FlattenEnv, [Type])
+coreFlattenTys = go []
   where
-    go rtys m []         = (m, reverse rtys)
-    go rtys m (ty : tys)
-      = let (m', ty') = coreFlattenTy in_scope m ty in
-        go (ty' : rtys) m' tys
+    go rtys env []         = (env, reverse rtys)
+    go rtys env (ty : tys)
+      = let (env', ty') = coreFlattenTy env ty in
+        go (ty' : rtys) env' tys
 
-coreFlattenTy :: InScopeSet -> FlattenMap -> Type -> (FlattenMap, Type)
-coreFlattenTy in_scope = go
+coreFlattenTy :: FlattenEnv -> Type -> (FlattenEnv, Type)
+coreFlattenTy = go
   where
-    go m ty | Just ty' <- coreView ty = go m ty'
+    go env ty | Just ty' <- coreView ty = go env ty'
 
-    go m ty@(TyVarTy {}) = (m, ty)
-    go m (AppTy ty1 ty2) = let (m1, ty1') = go m  ty1
-                               (m2, ty2') = go m1 ty2 in
-                           (m2, AppTy ty1' ty2')
-    go m (TyConApp tc tys)
+    go env (TyVarTy tv)    = (env, substTyVar (fe_subst env) tv)
+    go env (AppTy ty1 ty2) = let (env1, ty1') = go env  ty1
+                                 (env2, ty2') = go env1 ty2 in
+                             (env2, AppTy ty1' ty2')
+    go env (TyConApp tc tys)
          -- NB: Don't just check if isFamilyTyCon: this catches *data* families,
          -- which are generative and thus can be preserved during flattening
       | not (isGenerativeTyCon tc Nominal)
-      = let (m', tv) = coreFlattenTyFamApp in_scope m tc tys in
-        (m', mkTyVarTy tv)
+      = let (env', tv) = coreFlattenTyFamApp env tc tys in
+        (env', mkTyVarTy tv)
 
       | otherwise
-      = let (m', tys') = coreFlattenTys in_scope m tys in
-        (m', mkTyConApp tc tys')
+      = let (env', tys') = coreFlattenTys env tys in
+        (env', mkTyConApp tc tys')
+
+    go env (ForAllTy (Anon ty1) ty2) = let (env1, ty1') = go env  ty1
+                                           (env2, ty2') = go env1 ty2 in
+                                       (env2, mkFunTy ty1' ty2')
+
+    go env (ForAllTy (Named tv vis) ty)
+      = let (env1, tv') = coreFlattenVarBndr env tv
+            (env2, ty') = go env1 ty in
+        (env2, mkNamedForAllTy tv' vis ty')
+
+    go env ty@(LitTy {}) = (env, ty)
 
-    go m (FunTy ty1 ty2) = let (m1, ty1') = go m  ty1
-                               (m2, ty2') = go m1 ty2 in
-                           (m2, FunTy ty1' ty2')
+    go env (CastTy ty co) = let (env1, ty') = go env ty
+                                (env2, co') = coreFlattenCo env1 co in
+                            (env2, CastTy ty' co')
 
-      -- Note to RAE: this will have to be changed with kind families
-    go m (ForAllTy tv ty) = let (m', ty') = go m ty in
-                            (m', ForAllTy tv ty')
+    go env (CoercionTy co) = let (env', co') = coreFlattenCo env co in
+                             (env', CoercionTy co')
 
-    go m ty@(LitTy {}) = (m, ty)
+-- when flattening, we don't care about the contents of coercions.
+-- so, just return a fresh variable of the right (flattened) type
+coreFlattenCo :: FlattenEnv -> Coercion -> (FlattenEnv, Coercion)
+coreFlattenCo env co
+  = (env2, mkCoVarCo covar)
+  where
+    (env1, kind') = coreFlattenTy env (coercionType co)
+    fresh_name    = mkFlattenFreshCoName
+    subst1        = fe_subst env1
+    in_scope      = getTCvInScope subst1
+    covar         = uniqAway in_scope (mkCoVar fresh_name kind')
+    env2          = env1 { fe_subst = subst1 `extendTCvInScope` covar }
+
+coreFlattenVarBndr :: FlattenEnv -> TyVar -> (FlattenEnv, TyVar)
+coreFlattenVarBndr env tv
+  | kind' `eqType` kind
+  = ( env { fe_subst = extendTvSubst old_subst tv (mkTyVarTy tv) }
+             -- override any previous binding for tv
+    , tv)
 
-coreFlattenTyFamApp :: InScopeSet -> FlattenMap
+  | otherwise
+  = let new_tv    = uniqAway (getTCvInScope old_subst) (setTyVarKind tv kind')
+        new_subst = extendTvSubstWithClone old_subst tv new_tv
+    in
+    (env' { fe_subst = new_subst }, new_tv)
+  where
+    kind          = tyVarKind tv
+    (env', kind') = coreFlattenTy env kind
+    old_subst     = fe_subst env
+
+coreFlattenTyFamApp :: FlattenEnv
                     -> TyCon         -- type family tycon
                     -> [Type]        -- args
-                    -> (FlattenMap, TyVar)
-coreFlattenTyFamApp in_scope m fam_tc fam_args
-  = case lookupTypeMap m fam_ty of
-      Just tv -> (m, tv)
+                    -> (FlattenEnv, TyVar)
+coreFlattenTyFamApp env fam_tc fam_args
+  = case lookupTypeMap type_map fam_ty of
+      Just tv -> (env, tv)
               -- we need fresh variables here, but this is called far from
               -- any good source of uniques. So, we just use the fam_tc's unique
               -- and trust uniqAway to avoid clashes. Recall that the in_scope set
               -- contains *all* tyvars, even locally bound ones elsewhere in the
               -- overall type, so this really is fresh.
-      Nothing -> let tyvar_name   = mkSysTvName (getUnique fam_tc) (fsLit "fl")
-                     tv = uniqAway in_scope $ mkTyVar tyvar_name (typeKind fam_ty)
-                     m' = extendTypeMap m fam_ty tv in
-                 (m', tv)
-  where fam_ty = TyConApp fam_tc fam_args
-
+      Nothing -> let tyvar_name = mkFlattenFreshTyName fam_tc
+                     tv = uniqAway (getTCvInScope subst) $
+                          mkTyVar tyvar_name (typeKind fam_ty)
+                     env' = env { fe_type_map = extendTypeMap type_map fam_ty tv
+                                , fe_subst = extendTCvInScope subst tv }
+                 in (env', tv)
+  where fam_ty   = mkTyConApp fam_tc fam_args
+        FlattenEnv { fe_type_map = type_map
+                   , fe_subst = subst } = env
+
+-- | Get the set of all type variables mentioned anywhere in the list
+-- of types. These variables are not necessarily free.
 allTyVarsInTys :: [Type] -> VarSet
 allTyVarsInTys []       = emptyVarSet
 allTyVarsInTys (ty:tys) = allTyVarsInTy ty `unionVarSet` allTyVarsInTys tys
 
+-- | Get the set of all type variables mentioned anywhere in a type.
 allTyVarsInTy :: Type -> VarSet
 allTyVarsInTy = go
   where
     go (TyVarTy tv)      = unitVarSet tv
     go (AppTy ty1 ty2)   = (go ty1) `unionVarSet` (go ty2)
     go (TyConApp _ tys)  = allTyVarsInTys tys
-    go (FunTy ty1 ty2)   = (go ty1) `unionVarSet` (go ty2)
-    go (ForAllTy tv ty)  = (go (tyVarKind tv)) `unionVarSet`
-                           unitVarSet tv `unionVarSet`
-                           (go ty) -- don't remove tv
+    go (ForAllTy bndr ty) =
+      caseBinder bndr (\tv -> unitVarSet tv) (const emptyVarSet)
+      `unionVarSet` go (binderType bndr) `unionVarSet` go ty
+        -- don't remove the tv from the set!
     go (LitTy {})        = emptyVarSet
+    go (CastTy ty co)    = go ty `unionVarSet` go_co co
+    go (CoercionTy co)   = go_co co
+
+    go_co (Refl _ ty)           = go ty
+    go_co (TyConAppCo _ _ args) = go_cos args
+    go_co (AppCo co arg)        = go_co co `unionVarSet` go_co arg
+    go_co (ForAllCo tv h co)
+      = unionVarSets [unitVarSet tv, go_co co, go_co h]
+    go_co (CoVarCo cv)          = unitVarSet cv
+    go_co (AxiomInstCo _ _ cos) = go_cos cos
+    go_co (UnivCo p _ t1 t2)    = go_prov p `unionVarSet` go t1 `unionVarSet` go t2
+    go_co (SymCo co)            = go_co co
+    go_co (TransCo c1 c2)       = go_co c1 `unionVarSet` go_co c2
+    go_co (NthCo _ co)          = go_co co
+    go_co (LRCo _ co)           = go_co co
+    go_co (InstCo co arg)       = go_co co `unionVarSet` go_co arg
+    go_co (CoherenceCo c1 c2)   = go_co c1 `unionVarSet` go_co c2
+    go_co (KindCo co)           = go_co co
+    go_co (SubCo co)            = go_co co
+    go_co (AxiomRuleCo _ cs)    = go_cos cs
+
+    go_cos = foldr (unionVarSet . go_co) emptyVarSet
+
+    go_prov UnsafeCoerceProv    = emptyVarSet
+    go_prov (PhantomProv co)    = go_co co
+    go_prov (ProofIrrelProv co) = go_co co
+    go_prov (PluginProv _)      = emptyVarSet
+    go_prov (HoleProv _)        = emptyVarSet
+
+mkFlattenFreshTyName :: Uniquable a => a -> Name
+mkFlattenFreshTyName unq
+  = mkSysTvName (getUnique unq) (fsLit "flt")
+
+mkFlattenFreshCoName :: Name
+mkFlattenFreshCoName
+  = mkSystemVarName (deriveUnique eqPrimTyConKey 71) (fsLit "flc")