packages feed

ghc-tcplugin-api 0.3.0.0 → 0.3.1.0

raw patch · 6 files changed

+726/−311 lines, 6 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

changelog.md view
@@ -1,4 +1,10 @@ 
+# Version 0.3.1.0 (2021-08-09)
+
+Ensure that the coercions stored in `Reduction`s are always
+oriented left-to-right, by making the internal rewriting compatibility layer
+also use left-to-right coercions.
+
 # Version 0.3.0.0 (2021-08-04)
 
 Account for changes in rewriting in GHC 9.4:
ghc-tcplugin-api.cabal view
@@ -1,6 +1,6 @@ cabal-version:  3.0
 name:           ghc-tcplugin-api
-version:        0.3.0.0
+version:        0.3.1.0
 synopsis:       An API for type-checker plugins.
 license:        BSD-3-Clause
 build-type:     Simple
@@ -62,6 +62,7 @@   else
     other-modules:
       GHC.TcPlugin.API.Internal.Shim
+      GHC.TcPlugin.API.Internal.Shim.Reduction
 
   if impl(ghc < 9.5.0)
     cpp-options: -DHAS_DERIVEDS
src/GHC/TcPlugin/API.hs view
@@ -652,7 +652,9 @@          -> m ( Maybe Reduction )
 matchFam tycon args =
 #ifndef HAS_REWRITING
-  fmap mkReduction <$>
+  fmap ( \ (co,ty) -> mkReduction (mkSymCo co) ty ) <$>
+  -- GHC 9.0 and 9.2 use a different orientation
+  -- when rewriting type family applications.
 #endif
   ( liftTcPluginM $ GHC.matchFam tycon args )
 
@@ -798,11 +800,4 @@   -> TcType   -- ^ The type that the type family application reduces to
   -> Reduction
 mkTyFamAppReduction str role tc args ty =
-  -- The rewriter/flattener in GHC 9.2/GHC 9.0 uses coercions
-  -- in which the rewritten type is on the left.
-  -- Starting from GHC 9.4, the rewritten type is always on the right.
-#if HAS_REWRITING
   Reduction ( mkPluginUnivCo str role ( mkTyConApp tc args ) ty ) ty
-#else
-  Reduction ( mkPluginUnivCo str role ty ( mkTyConApp tc args ) ) ty
-#endif
src/GHC/TcPlugin/API/Internal.hs view
@@ -59,10 +59,6 @@   where
 
 -- base
-#ifndef HAS_REWRITING
-import Data.IORef
-  ( IORef, newIORef )
-#endif
 import Data.Kind
   ( Constraint, Type )
 import GHC.TypeLits
@@ -101,11 +97,7 @@     ( fsLit )
 import qualified GHC.Tc.Plugin
   as GHC
-    ( tcLookupDataCon, tcLookupTyCon
-#ifndef HAS_REWRITING
-    , tcPluginIO
-#endif
-    )
+    ( tcLookupDataCon, tcLookupTyCon )
 import qualified GHC.Tc.Types
   as GHC
     ( TcM, TcPlugin(..), TcPluginM
@@ -133,17 +125,12 @@ import qualified GHC.Types.Unique.FM
   as GHC
     ( UniqFM )
-#ifndef HAS_REWRITING
-import qualified GHC.Types.Unique.DFM
-  as GHC
-    ( emptyUDFM )
-#endif
 
 -- ghc-tcplugin-api
 #ifndef HAS_REWRITING
 import GHC.TcPlugin.API.Internal.Shim
   ( TcPluginSolveResult, TcPluginRewriteResult(..)
-  , RewrittenTyFamApps, RewriteEnv
+  , RewriteEnv
   , shimRewriter
   )
 #endif
@@ -169,7 +156,7 @@ 
 -- | For rewriting type family applications, a type-checking plugin provides
 -- a function of this type for each type family 'GHC.Core.TyCon.TyCon'.
--- 
+--
 -- The function is provided with the current set of Given constraints, together
 -- with the arguments to the type family.
 -- The type family application will always be fully saturated.
@@ -544,9 +531,6 @@     , showTypeTyCon  :: !GHC.TyCon
     , concatTyCon    :: !GHC.TyCon
     , vcatTyCon      :: !GHC.TyCon
-#ifndef HAS_REWRITING
-    , rewrittenTyFamsIORef :: !( IORef RewrittenTyFamApps )
-#endif
     }
 
 data TcPluginDefs s
@@ -562,9 +546,6 @@   showTypeTyCon   <- GHC.promoteDataCon <$> GHC.tcLookupDataCon GHC.TypeLits.typeErrorShowTypeDataConName
   concatTyCon     <- GHC.promoteDataCon <$> GHC.tcLookupDataCon GHC.TypeLits.typeErrorAppendDataConName
   vcatTyCon       <- GHC.promoteDataCon <$> GHC.tcLookupDataCon GHC.TypeLits.typeErrorVAppendDataConName
-#ifndef HAS_REWRITING
-  rewrittenTyFamsIORef <- GHC.tcPluginIO $ newIORef GHC.emptyUDFM
-#endif
   pure ( BuiltinDefs { .. } )
 
 interpretErrorMessage :: BuiltinDefs -> TcPluginErrorMessage -> GHC.PredType
src/GHC/TcPlugin/API/Internal/Shim.hs view
@@ -3,6 +3,7 @@ {-# LANGUAGE CPP             #-}
 {-# LANGUAGE DerivingVia     #-}
 {-# LANGUAGE LambdaCase      #-}
+{-# LANGUAGE NamedFieldPuns  #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE TupleSections   #-}
 {-# LANGUAGE ViewPatterns    #-}
@@ -15,9 +16,17 @@ in typechecking plugins on GHC 9.0 and 9.2.
 -}
 
-module GHC.TcPlugin.API.Internal.Shim where
+module GHC.TcPlugin.API.Internal.Shim
+  ( Reduction(..), mkReduction
+  , TcPluginSolveResult, TcPluginRewriteResult(..)
+  , RewriteEnv(..)
+  , shimRewriter
+  )
+  where
 
 -- base
+import Prelude
+  hiding ( Floating(cos), iterate )
 import Control.Monad
   ( forM, when )
 #if !MIN_VERSION_ghc(9,2,0)
@@ -39,32 +48,16 @@ 
 -- ghc
 import GHC.Core.Coercion
-  ( castCoercionKind1, coercionRKind
-  , mkReflCo, mkSymCo, mkFunCo, mkHomoForAllCos
-  , mkTransCo, mkAppCos, mkNomReflCo, mkSubCo
+  ( mkReflCo, mkSymCo
+  , mkAppCos, mkNomReflCo, mkSubCo
   , mkTyConAppCo, tyConRolesX
   , tyConRolesRepresentational
-  , simplifyArgsWorker
-#if !MIN_VERSION_ghc(9,2,0)
-  , coToMCo
-#endif
   )
-#if MIN_VERSION_ghc(9,2,0)
-import GHC.Core.Map.Type
-  ( LooseTypeMap )
-#else
-import GHC.Core.Map
-  ( LooseTypeMap )
-#endif
 import GHC.Core.Predicate
   ( EqRel(..), eqRelRole )
 import GHC.Core.TyCo.Rep
-  ( Type(..), Kind, Coercion(..)
-  , TyCoBinder(..)
-  , MCoercion(..), MCoercionN
-  , binderVars, mkForAllTys
-  , isNamedBinder
-  , mkTyVarTy
+  ( Type(..), Kind, Coercion(..), MCoercion(..), TyCoBinder(..)
+  , isNamedBinder, mkTyVarTy
   )
 import GHC.Core.TyCon
   ( TyCon(..), TyConBinder, TyConBndrVis(..)
@@ -78,19 +71,14 @@   )
 import GHC.Core.Type
   ( TyVar
-  , tcView
-  , mkCoercionTy, mkCastTy, mkAppTys
-  , mkTyConApp, mkScaled, coreView
-  , tymult, tyVarKind
+  , tcView , mkTyConApp, mkScaled, coreView , tymult, tyVarKind
   )
 #if MIN_VERSION_ghc(9,2,0)
 import GHC.Data.Maybe
   ( firstJustsM )
 #endif
-import GHC.Data.TrieMap
-  ( ListMap )
 import GHC.Tc.Plugin
-  ( tcPluginIO, newWanted, newDerived )
+  ( newWanted, newDerived )
 import GHC.Tc.Solver.Monad
   ( TcS
   , zonkCo, zonkTcType
@@ -114,7 +102,6 @@ import GHC.Tc.Types.Constraint
   ( Ct(..)
   , CtLoc, CtFlavour(..), CtFlavourRole, ShadowInfo(..)
-  , Xi
 #if MIN_VERSION_ghc(9,2,0)
   , CanEqLHS(..)
 #endif
@@ -124,10 +111,7 @@   )
 import GHC.Tc.Types.Evidence
   ( EvTerm(..), Role(..)
-  , evCast
-  , mkTcReflCo, mkTcTransCo, mkTcSymCo
-  , mkTcTyConAppCo
-  , tcDowngradeRole
+  , evCast, mkTcTransCo , mkTcTyConAppCo
   )
 import GHC.Tc.Utils.TcType
   ( TcTyCoVarSet
@@ -140,12 +124,6 @@   , tcTypeKind
   , tyCoVarsOfType
   )
-#if !MIN_VERSION_ghc(9,2,0)
-import GHC.Types.Unique
-  ( Unique )
-#endif
-import GHC.Types.Unique.DFM
-  ( UniqDFM )
 import GHC.Types.Unique.FM
   ( UniqFM, lookupUFM, isNullUFM )
 import GHC.Types.Var
@@ -164,34 +142,12 @@ import GHC.Utils.Monad
   ( zipWith3M )
 import GHC.Utils.Outputable
-  ( Outputable(..), SDoc
-  , (<+>), braces, empty, text, vcat
-  )
-
---------------------------------------------------------------------------------
+  ( Outputable(..), SDoc, empty )
 
--- | A reduction to the provided type, with a coercion witnessing the equality.
-data Reduction =
-  Reduction
-    { reductionCoercion    :: Coercion
-    , reductionReducedType :: !Type
-    }
-instance Outputable Reduction where
-  ppr redn =
-    braces $ vcat
-      [ text "reductionOriginalType:" <+> ppr (reductionOriginalType redn)
-      , text " reductionReducedType:" <+> ppr (reductionReducedType redn)
-      , text "    reductionCoercion:" <+> ppr (reductionCoercion redn)
-      ]
+-- ghc-tcplugin-api
+import GHC.TcPlugin.API.Internal.Shim.Reduction
 
-reductionOriginalType :: Reduction -> Type
-reductionOriginalType = coercionRKind . reductionCoercion
--- NB. This is coercionRKind, because in GHC 9.0 & 9.2, the flattener/rewriter
--- uses coercions where the original type is on the right and the rewritten type
--- is on the left.
--- This orientation changes in GHC 9.4, but this module is only present
--- when compiling with GHC 9.0 or GHC 9.2.
-{-# INLINE reductionOriginalType #-}
+--------------------------------------------------------------------------------
 
 data RewriteEnv
   = FE { fe_loc     :: !CtLoc
@@ -199,15 +155,6 @@        , fe_eq_rel  :: !EqRel
        }
 
-mkReduction :: ( Coercion, Type ) -> Reduction
-mkReduction ( co, ty ) = Reduction co ty
-
-runReduction1 :: Reduction -> ( Type, Coercion )
-runReduction1 ( Reduction co ty ) = ( ty, co )
-
-runReduction2 :: Reduction -> ( Coercion, Type )
-runReduction2 ( Reduction co ty ) = ( co, ty )
-
 type TcPluginSolveResult = TcPluginResult
 
 data TcPluginRewriteResult
@@ -260,7 +207,7 @@     rwEnv :: RewriteEnv
     rwEnv = FE ( ctLoc ct ) ( ctFlavour ct ) ( ctEqRel ct )
     shimRewriteEnv :: ShimRewriteEnv
-    shimRewriteEnv = ShimRewriteEnv rws rwEnv ct givens
+    shimRewriteEnv = ShimRewriteEnv rws rwEnv givens
   ( res, newCts ) <- runRewritePluginM shimRewriteEnv ( rewrite_one predTy )
   case res of
     Nothing -> pure ( Nothing, newCts )
@@ -271,7 +218,9 @@         Derived   -> newDerived ( ctLoc ct ) predTy'
       pure ( Just
               ( mkNonCanonical ctEv'
-              , ( evCast ( ctEvExpr ctEv' ) co, ct )
+              , ( evCast ( ctEvExpr ctEv' ) ( mkSymCo co )
+                , ct
+                )
               )
            , newCts
            )
@@ -292,110 +241,120 @@ -- The following is (mostly) copied from GHC 9.4's GHC.Tc.Solver.Rewrite module.
 
 rewrite_one :: Type -> RewriteM Reduction
-rewrite_one = \case
-  ( rewriterView -> Just ty' )
-    -> rewrite_one ty'
-  ty@( LitTy {} )
-    -> do
-      role <- getRole
-      pure $ Reduction ( mkReflCo role ty ) ty
-  TyVarTy tv
-    -> rewriteTyVar tv
-  AppTy ty1 ty2
-    -> rewrite_app_tys ty1 [ty2]
-  TyConApp tc tys
-    | isTypeFamilyTyCon tc
-    -> rewrite_fam_app tc tys
-    | otherwise
-    -> rewrite_ty_con_app tc tys
-  ty@( FunTy { ft_mult = mult, ft_arg = ty1, ft_res = ty2 } )
-    -> do
-      Reduction co1 xi1 <- rewrite_one ty1
-      Reduction co2 xi2 <- rewrite_one ty2
-      Reduction co3 xi3 <- setEqRel NomEq $ rewrite_one mult
-      role <- getRole
-      return $
-       Reduction
-         ( mkFunCo role co3 co1 co2 )
-         ( ty { ft_mult = xi3, ft_arg = xi1, ft_res = xi2 } )
+rewrite_one ty
+  | Just ty' <- rewriterView ty  -- See Note [Rewriting synonyms]
+  = rewrite_one ty'
 
-  ty@( ForAllTy {} )
-    -> do
-      let
-        (bndrs, rho) = tcSplitForAllTyVarBinders ty
-        tvs          = binderVars bndrs
-      Reduction co rho' <- rewrite_one rho
-      pure $ Reduction
-        ( mkHomoForAllCos tvs co )
-        ( mkForAllTys bndrs rho' )
+rewrite_one xi@(LitTy {})
+  = do { role <- getRole
+       ; return $ mkReflRedn role xi }
 
-  CastTy ty g
-    -> do 
-      Reduction co xi <- rewrite_one ty
-      (g', _) <- rewrite_co g
-      role <- getRole
-      pure $ Reduction
-        ( castCoercionKind1 co role xi ty g' )
-        ( mkCastTy xi g' )
+rewrite_one (TyVarTy tv)
+  = rewriteTyVar tv
 
-  CoercionTy co
-    -> do
-      ( co1, co2 ) <- rewrite_co co
-      pure $ Reduction co2 ( mkCoercionTy co1 )
+rewrite_one (AppTy ty1 ty2)
+  = rewrite_app_tys ty1 [ty2]
 
+rewrite_one (TyConApp tc tys)
+  | isTypeFamilyTyCon tc
+  = rewrite_fam_app tc tys
+
+  | otherwise
+  = rewrite_ty_con_app tc tys
+
+rewrite_one (FunTy { ft_af = vis, ft_mult = mult, ft_arg = ty1, ft_res = ty2 })
+  = do { arg_redn <- rewrite_one ty1
+       ; res_redn <- rewrite_one ty2
+       ; w_redn <- setEqRel NomEq $ rewrite_one mult
+       ; role <- getRole
+       ; return $ mkFunRedn role vis w_redn arg_redn res_redn }
+
+rewrite_one ty@(ForAllTy {})
+  = do { let (bndrs, rho) = tcSplitForAllTyVarBinders ty
+       ; redn <- rewrite_one rho
+       ; return $ mkHomoForAllRedn bndrs redn }
+
+rewrite_one (CastTy ty g)
+  = do { redn <- rewrite_one ty
+       ; g'   <- rewrite_co g
+       ; role <- getRole
+       ; return $ mkCastRedn1 role ty g' redn }
+
+rewrite_one (CoercionTy co)
+  = do { co' <- rewrite_co co
+       ; role <- getRole
+       ; return $ mkReflCoRedn role co' }
+
+rewrite_reduction :: Reduction -> RewriteM Reduction
+rewrite_reduction (Reduction co xi) = do
+  redn <- bumpDepth $ rewrite_one xi
+  pure $ co `mkTransRedn` redn
+
 rewrite_app_tys :: Type -> [Type] -> RewriteM Reduction
-rewrite_app_tys ( AppTy ty1 ty2 ) tys =
-  rewrite_app_tys ty1 ( ty2 : tys )
-rewrite_app_tys fun_ty arg_tys = do
-  Reduction fun_co fun_xi <- rewrite_one fun_ty
-  rewrite_app_ty_args fun_xi fun_co arg_tys
+rewrite_app_tys (AppTy ty1 ty2) tys = rewrite_app_tys ty1 (ty2:tys)
+rewrite_app_tys fun_ty arg_tys
+  = do { redn <- rewrite_one fun_ty
+       ; rewrite_app_ty_args redn arg_tys }
 
-rewrite_app_ty_args :: Xi -> Coercion -> [Type] -> RewriteM Reduction
-rewrite_app_ty_args fun_xi fun_co [] = pure $ Reduction fun_co fun_xi
-rewrite_app_ty_args fun_xi fun_co arg_tys = do
-  (xi, co, kind_co) <- case tcSplitTyConApp_maybe fun_xi of
-    Just (tc, xis) -> do
-      let tc_roles  = tyConRolesRepresentational tc
-          arg_roles = dropList xis tc_roles
-      (arg_xis, arg_cos, kind_co)
-        <- rewrite_vector (tcTypeKind fun_xi) arg_roles arg_tys
-      eq_rel <- getEqRel
-      let app_xi = mkTyConApp tc (xis ++ arg_xis)
-          app_co = case eq_rel of
-            NomEq  -> mkAppCos fun_co arg_cos
-            ReprEq -> mkTcTyConAppCo Representational tc
-                        (zipWith mkReflCo tc_roles xis ++ arg_cos)
-                      `mkTcTransCo`
-                      mkAppCos fun_co (map mkNomReflCo arg_tys)
-      return (app_xi, app_co, kind_co)
-    Nothing -> do
-      (arg_xis, arg_cos, kind_co)
-        <- rewrite_vector (tcTypeKind fun_xi) (repeat Nominal) arg_tys
-      let arg_xi = mkAppTys fun_xi arg_xis
-          arg_co = mkAppCos fun_co arg_cos
-      return (arg_xi, arg_co, kind_co)
-  role <- getRole
-  return (homogenise_result xi co role kind_co)
+rewrite_app_ty_args :: Reduction -> [Type] -> RewriteM Reduction
+rewrite_app_ty_args redn []
+  = return redn
+rewrite_app_ty_args fun_redn@(Reduction fun_co fun_xi) arg_tys
+  = do { het_redn <- case tcSplitTyConApp_maybe fun_xi of
+           Just (tc, xis) ->
+             do { let tc_roles  = tyConRolesRepresentational tc
+                      arg_roles = dropList xis tc_roles
+                ; ArgsReductions (Reductions arg_cos arg_xis) kind_co
+                    <- rewrite_vector (tcTypeKind fun_xi) arg_roles arg_tys
 
+                ; eq_rel <- getEqRel
+                ; let app_xi = mkTyConApp tc (xis ++ arg_xis)
+                      app_co = case eq_rel of
+                        NomEq  -> mkAppCos fun_co arg_cos
+                        ReprEq -> mkAppCos fun_co (map mkNomReflCo arg_tys)
+                                  `mkTcTransCo`
+                                  mkTcTyConAppCo Representational tc
+                                    (zipWith mkReflCo tc_roles xis ++ arg_cos)
+
+                ; return $
+                    mkHetReduction
+                      (mkReduction app_co app_xi )
+                      kind_co }
+           Nothing ->
+             do { ArgsReductions redns kind_co
+                    <- rewrite_vector (tcTypeKind fun_xi) (repeat Nominal) arg_tys
+                ; return $ mkHetReduction (mkAppRedns fun_redn redns) kind_co }
+
+       ; role <- getRole
+       ; return (homogeniseHetRedn role het_redn) }
+
 {-# INLINE rewrite_args_tc #-}
-rewrite_args_tc :: TyCon -> Maybe [Role] -> [Type] -> RewriteM ( [Xi], [Coercion] , MCoercionN)
+rewrite_args_tc :: TyCon -> Maybe [Role] -> [Type] -> RewriteM ArgsReductions
 rewrite_args_tc tc = rewrite_args all_bndrs any_named_bndrs inner_ki emptyVarSet
+  -- NB: TyCon kinds are always closed
   where
+  -- There are many bang patterns in here. It's been observed that they
+  -- greatly improve performance of an optimized build.
+  -- The T9872 test cases are good witnesses of this fact.
+
     (bndrs, named)
       = ty_con_binders_ty_binders' (tyConBinders tc)
+    -- it's possible that the result kind has arrows (for, e.g., a type family)
+    -- so we must split it
     (inner_bndrs, inner_ki, inner_named) = split_pi_tys' (tyConResKind tc)
     !all_bndrs                           = bndrs `chkAppend` inner_bndrs
     !any_named_bndrs                     = named || inner_named
+    -- NB: Those bangs there drop allocations in T9872{a,c,d} by 8%.
 
 rewrite_fam_app :: TyCon -> [Type] -> RewriteM Reduction
 rewrite_fam_app tc tys = do
   let (tys1, tys_rest) = splitAt (tyConArity tc) tys
-  Reduction co1 xi1 <- rewrite_exact_fam_app tc tys1
-  rewrite_app_ty_args xi1 co1 tys_rest
+  redn <- rewrite_exact_fam_app tc tys1
+  rewrite_app_ty_args redn tys_rest
 
 rewrite_exact_fam_app :: TyCon -> [Type] -> RewriteM Reduction
 rewrite_exact_fam_app tc tys = do
-  checkStackDepth (mkTyConApp tc tys)
+  checkStackDepth $ mkTyConApp tc tys
   rws <- getRewriters
   let
     mbRewriter :: Maybe Rewriter
@@ -405,45 +364,51 @@     Just redn -> finish False redn
     _ -> do
       eq_rel <- getEqRel
-      (xis, coercions, kind_co) <-
+      ArgsReductions (Reductions cos xis) kind_co <-
         if eq_rel == NomEq
         then rewrite_args_tc tc Nothing tys
         else setEqRel NomEq $
              rewrite_args_tc tc Nothing tys
-      let role    = eqRelRole eq_rel
-          args_co = mkTyConAppCo role tc coercions
-          homogenise :: Reduction -> Reduction
-          homogenise (Reduction co xi) =
-            homogenise_result xi (co `mkTcTransCo` args_co) role kind_co
-          giveUp :: Reduction
-          giveUp = homogenise $ Reduction (mkTcReflCo role reduced) reduced
-             where reduced = mkTyConApp tc xis
+      let
+        role    = eqRelRole eq_rel
+        args_co = mkTyConAppCo role tc cos
+        homogenise :: Reduction -> Reduction
+        homogenise redn
+          = homogeniseHetRedn role
+          $ mkHetReduction
+              (args_co `mkTransRedn` redn)
+              kind_co
+        give_up :: Reduction
+        give_up = homogenise $ mkReflRedn role (mkTyConApp tc xis)
+
       result2 <- liftTcS $ lookupFamAppInert tc xis
       flavour <- getFlavour
       case result2 of
         Just (co, xi, fr@(_, inert_eq_rel))
           | fr `eqCanRewriteFR` (flavour, eq_rel)
-          -> finish True (homogenise $ Reduction downgraded_co xi)
+          , let
+              redn :: Reduction
+              redn = Reduction (mkSymCo co) xi -- inerts use a different orientation in GHC 9.0 and 9.2
+          -> finish True (homogenise $ downgradeRedn role' inert_role redn)
           where
-            inert_role    = eqRelRole inert_eq_rel
-            role'         = eqRelRole eq_rel
-            downgraded_co = tcDowngradeRole role' inert_role (mkTcSymCo co)
+            inert_role      = eqRelRole inert_eq_rel
+            role'           = eqRelRole eq_rel
         _ -> do
           result3 <- try_to_reduce tc xis mbRewriter
           case result3 of
             Just redn -> finish True (homogenise redn)
-            _         -> return giveUp
+            _         -> return give_up
   where
     finish :: Bool -> Reduction -> RewriteM Reduction
     finish use_cache (Reduction co xi) = do
       Reduction fully_co fully <- bumpDepth $ rewrite_one xi
-      let final_redn = Reduction (fully_co `mkTcTransCo` co) fully
+      let final_redn@(Reduction final_co final_xi) = Reduction (fully_co `mkTcTransCo` co) fully
       eq_rel <- getEqRel
       flavour <- getFlavour
       when (use_cache && eq_rel == NomEq && flavour /= Derived) $
         liftTcS $
           extendFamAppCache tc tys
-            ( runReduction2 final_redn )
+            ( mkSymCo final_co, final_xi ) -- different orientation in GHC 9.0 and 9.2
 #if !MIN_VERSION_ghc(9,2,0)
             flavour
 #endif
@@ -497,21 +462,19 @@           -> addRewriting Nothing []
 
 rewrite_ty_con_app :: TyCon -> [Type] -> RewriteM Reduction
-rewrite_ty_con_app tc tys = do
-  role <- getRole
-  let m_roles | Nominal <- role = Nothing
-              | otherwise       = Just $ tyConRolesX role tc
-  (xis, coercions, kind_co) <- rewrite_args_tc tc m_roles tys
-  let tyconapp_xi = mkTyConApp tc xis
-      tyconapp_co = mkTyConAppCo role tc coercions
-  return (homogenise_result tyconapp_xi tyconapp_co role kind_co)
+rewrite_ty_con_app tc tys
+  = do { role <- getRole
+       ; let m_roles | Nominal <- role = Nothing
+                     | otherwise       = Just $ tyConRolesX role tc
+       ; ArgsReductions redns kind_co <- rewrite_args_tc tc m_roles tys
+       ; let tyconapp_redn
+                = mkHetReduction
+                    (mkTyConAppRedn role tc redns)
+                    kind_co
+       ; return $ homogeniseHetRedn role tyconapp_redn }
 
-rewrite_co :: Coercion -> RewriteM ( Coercion, Coercion )
-rewrite_co co = do
-  zonked_co <- liftTcS $ zonkCo co
-  env_role <- getRole
-  let co' = mkTcReflCo env_role ( mkCoercionTy zonked_co )
-  pure ( zonked_co, co' )
+rewrite_co :: Coercion -> RewriteM Coercion
+rewrite_co co = liftTcS $ zonkCo co
 
 rewriterView :: Type -> Maybe Type
 rewriterView ty@(TyConApp tc _)
@@ -526,18 +489,16 @@ rewriteTyVar tv = do
   mb_yes <- rewrite_tyvar1 tv
   case mb_yes of
-    RTRFollowed ty1 co1 -> do
-      Reduction co2 ty2 <- rewrite_one ty1
-      pure $ Reduction (co2 `mkTransCo` co1) ty2
+    RTRFollowed redn -> rewrite_reduction redn
     RTRNotFollowed -> do
       tv' <- liftTcS $ updateTyVarKindM zonkTcType tv
       role <- getRole
       let ty' = mkTyVarTy tv'
-      return $ Reduction (mkTcReflCo role ty') ty'
+      pure $ mkReflRedn role ty'
 
 data RewriteTvResult
   = RTRNotFollowed
-  | RTRFollowed Type Coercion
+  | RTRFollowed !Reduction
 
 rewrite_tyvar1 :: TcTyVar -> RewriteM RewriteTvResult
 rewrite_tyvar1 tv = do
@@ -545,7 +506,7 @@   case mb_ty of
     Just ty -> do
       role <- getRole
-      return (RTRFollowed ty (mkReflCo role ty))
+      return $ RTRFollowed $ mkReflRedn role ty
     Nothing -> do
       fr <- getFlavourRole
       rewrite_tyvar2 tv fr
@@ -565,49 +526,29 @@ #endif
       , let ct_fr = (ctEvFlavour ctev, ct_eq_rel)
       , ct_fr `eqCanRewriteFR` fr
-      -> do 
-        let rewrite_co1 = mkSymCo (ctEvCoercion ctev)
-            rewrite_co2 = case (ct_eq_rel, eq_rel) of
-              (ReprEq, _rel)  -> rewrite_co1
-              (NomEq, NomEq)  -> rewrite_co1
-              (NomEq, ReprEq) -> mkSubCo rewrite_co1
-        return (RTRFollowed rhs_ty rewrite_co2)
+      -> do
+          let rewriting_co1 = ctEvCoercion ctev
+              rewriting_co  = case (ct_eq_rel, eq_rel) of
+                (ReprEq, _rel)  -> rewriting_co1
+                (NomEq, NomEq)  -> rewriting_co1
+                (NomEq, ReprEq) -> mkSubCo rewriting_co1
+          return $ RTRFollowed $ mkReduction rewriting_co rhs_ty 
     _other -> return RTRNotFollowed
 
 rewrite_vector :: Kind
                -> [Role]
                -> [Type]
-               -> RewriteM ([Xi], [Coercion], MCoercionN)
-rewrite_vector ki roles tys = do
-  eq_rel <- getEqRel
-  case eq_rel of
-    NomEq ->
-      rewrite_args bndrs
-        any_named_bndrs
-        inner_ki
-        fvs
-        Nothing
-        tys
-    ReprEq ->
-      rewrite_args bndrs
-        any_named_bndrs
-        inner_ki
-        fvs
-        (Just roles)
-        tys
+               -> RewriteM ArgsReductions
+rewrite_vector ki roles tys
+  = do { eq_rel <- getEqRel
+       ; let mb_roles = case eq_rel of { NomEq -> Nothing; ReprEq -> Just roles }
+       ; rewrite_args bndrs any_named_bndrs inner_ki fvs mb_roles tys
+       }
   where
     (bndrs, inner_ki, any_named_bndrs) = split_pi_tys' ki
     fvs                                = tyCoVarsOfType ki
 {-# INLINE rewrite_vector #-}
 
-homogenise_result :: Xi
-                  -> Coercion
-                  -> Role
-                  -> MCoercionN
-                  -> Reduction
-homogenise_result xi co _ MRefl = Reduction co xi
-homogenise_result xi co r mco@(MCo kind_co)
-  = Reduction ((mkSymCo $ GRefl r xi mco) `mkTransCo` co) (xi `mkCastTy` kind_co)
 split_pi_tys' :: Type -> ([TyCoBinder], Type, Bool)
 split_pi_tys' ty = split ty ty
   where
@@ -631,10 +572,11 @@     {-# INLINE go #-}
 {-# INLINE ty_con_binders_ty_binders' #-}
 
+{-# INLINE rewrite_args #-}
 rewrite_args :: [TyCoBinder] -> Bool
              -> Kind -> TcTyCoVarSet
              -> Maybe [Role] -> [Type]
-             -> RewriteM ([Xi], [Coercion], MCoercionN)
+             -> RewriteM ArgsReductions
 rewrite_args orig_binders
              any_named_bndrs
              orig_inner_ki
@@ -647,41 +589,37 @@         where orig_roles = fromMaybe (repeat Nominal) orig_m_roles
 
 {-# INLINE rewrite_args_fast #-}
-rewrite_args_fast :: [Type]
-                  -> RewriteM ([Xi], [Coercion], MCoercionN)
+rewrite_args_fast :: [Type] -> RewriteM ArgsReductions
 rewrite_args_fast orig_tys
-  = fmap finish (iterateRewrite orig_tys)
+  = fmap finish (iterate orig_tys)
   where
 
-    iterateRewrite :: [Type] -> RewriteM ([Xi], [Coercion])
-    iterateRewrite (ty:tys) = do
-      Reduction co xi <- rewrite_one ty
-      (xis, coercions) <- iterateRewrite tys
-      pure (xi : xis, co : coercions)
-    iterateRewrite [] = pure ([], [])
+    iterate :: [Type] -> RewriteM Reductions
+    iterate (ty : tys) = do
+      Reduction  co  xi  <- rewrite_one ty
+      Reductions cos xis <- iterate tys
+      pure $ Reductions (co : cos) (xi : xis)
+    iterate [] = pure $ Reductions [] []
 
     {-# INLINE finish #-}
-    finish :: ([Xi], [Coercion]) -> ([Xi], [Coercion], MCoercionN)
-    finish (xis, coercions) = (xis, coercions, MRefl)
+    finish :: Reductions -> ArgsReductions
+    finish redns = ArgsReductions redns MRefl
 
 {-# INLINE rewrite_args_slow #-}
 rewrite_args_slow :: [TyCoBinder] -> Kind -> TcTyCoVarSet
                   -> [Role] -> [Type]
-                  -> RewriteM ([Xi], [Coercion], MCoercionN)
-rewrite_args_slow binders inner_ki fvs roles tys = do
-  rewritten_args <-
-    zipWith3M fl (map isNamedBinder binders ++ repeat True)
-      roles tys
-  pure
-#if !MIN_VERSION_ghc(9,2,0)
-    $ ( \ ( xs, cs, c ) -> ( xs, cs, coToMCo c ) )
-#endif
-    $ simplifyArgsWorker binders inner_ki fvs roles rewritten_args
+                  -> RewriteM ArgsReductions
+rewrite_args_slow binders inner_ki fvs roles tys
+  = do { rewritten_args <- zipWith3M fl (map isNamedBinder binders ++ repeat True)
+                                        roles tys
+       ; return $ simplifyArgsWorker binders inner_ki fvs roles rewritten_args }
   where
     {-# INLINE fl #-}
-    fl :: Bool -> Role -> Type -> RewriteM ( Type, Coercion )
-    fl True  r ty = noBogusCoercions $ runReduction1 <$> fl1 r ty
-    fl False r ty =                    runReduction1 <$> fl1 r ty
+    fl :: Bool   -- must we ensure to produce a real coercion here?
+                 -- see comment at top of function
+       -> Role -> Type -> RewriteM Reduction
+    fl True  r ty = noBogusCoercions $ fl1 r ty
+    fl False r ty =                    fl1 r ty
 
     {-# INLINE fl1 #-}
     fl1 :: Role -> Type -> RewriteM Reduction
@@ -695,7 +633,7 @@ 
     fl1 Phantom ty
       = do { ty' <- liftTcS $ zonkTcType ty
-           ; pure $ Reduction ( mkReflCo Phantom ty' ) ty' }
+           ; return $ mkReflRedn Phantom ty' }
 
 noBogusCoercions :: RewriteM a -> RewriteM a
 noBogusCoercions thing_inside
@@ -728,20 +666,10 @@    , reductionOccurred :: !ReduceQ
    }
 
-type RewrittenTyFamApps =
-  UniqDFM
-#if MIN_VERSION_ghc(9,2,0)
-    TyCon
-#else
-    Unique
-#endif
-    ( ListMap LooseTypeMap ( Maybe Reduction, [Ct] ) )
-
 data ShimRewriteEnv
   = ShimRewriteEnv
   { rewriters     :: !Rewriters
   , rewriteEnv    :: !RewriteEnv
-  , rewriteCt     :: !Ct
   , rewriteGivens :: ![ Ct ]
   }
 
@@ -763,19 +691,15 @@ runRewritePluginM env ( RewriteM { runRewriteM = run } ) = do
 
   evBindsVar <- getEvBindsTcPluginM
-  ( a, RewriteState newCts didReduce )
+  ( a, RewriteState { rewrittenCts, reductionOccurred } )
     <- unsafeTcPluginTcM
      $ runTcSWithEvBinds evBindsVar
      $ run env ( RewriteState [] NoReduction )
   let
-    mb_a = case didReduce of
+    mb_a = case reductionOccurred of
       NoReduction -> Nothing
       DidReduce   -> Just a
-  pure ( mb_a, newCts )
-
-setDidReduce :: RewriteM ()
-setDidReduce = RewriteM \ _env ( RewriteState cts _ ) ->
-  pure ( (), RewriteState cts DidReduce )
+  pure ( mb_a, rewrittenCts )
 
 addRewriting :: Maybe Reduction -> [Ct] -> RewriteM ( Maybe Reduction )
 addRewriting mbRedn newCts = RewriteM \ _ ( RewriteState cts s ) ->
@@ -788,18 +712,8 @@       = s
   in pure ( mbRedn , RewriteState ( cts <> newCts ) s' )
 
--- Silly workaround because wrapTcS is not exported in GHC 9.0
-liftIOTcS :: IO a -> TcS a
-liftIOTcS = runTcPluginTcS . tcPluginIO
-
 getRewriters :: RewriteM Rewriters
 getRewriters = RewriteM \ env s -> pure ( rewriters env, s )
-
-getGivens :: RewriteM [Ct]
-getGivens = RewriteM \ env s -> pure ( rewriteGivens env, s )
-
-getRewriteCt :: RewriteM Ct
-getRewriteCt = RewriteM \ env s -> pure ( rewriteCt env, s )
 
 getRewriteEnvField :: (RewriteEnv -> a) -> RewriteM a
 getRewriteEnvField accessor = RewriteM \ env s ->
+ src/GHC/TcPlugin/API/Internal/Shim/Reduction.hs view
@@ -0,0 +1,518 @@+{-# LANGUAGE BangPatterns #-}
+
+module GHC.TcPlugin.API.Internal.Shim.Reduction where
+
+-- base
+import Prelude
+  hiding (Floating(cos))
+
+-- ghc
+import GHC.Core.Class
+  ( Class(classTyCon) )
+import GHC.Core.Coercion
+  ( Coercion, CoercionN, MCoercion(..)
+  , Role(Nominal), LiftingContext
+  , castCoercionKind1, castCoercionKind2
+  , coercionLKind, coercionRKind, coercionKind
+  , coToMCo
+  , decomposePiCos, downgradeRole
+  , liftCoSubst, emptyLiftingContext, extendLiftingContextAndInScope, zapLiftingContext
+  , mkAppCo, mkAppCos
+  , mkCoherenceRightCo
+  , mkForAllCo, mkFunCo
+  , mkGReflLeftCo, mkGReflRightCo
+  , mkHomoForAllCos, mkProofIrrelCo
+  , mkReflCo, mkSubCo, mkSymCo, mkTransCo, mkTyConAppCo
+  )
+import GHC.Core.Predicate
+  ( mkClassPred )
+import GHC.Core.TyCo.Rep
+  ( TyCoBinder, mkFunTy )
+import GHC.Core.TyCon
+  ( TyCon )
+import GHC.Core.Type
+  ( AnonArgFlag, ArgFlag, Kind, Type, TyVar, TyVarBinder
+  , binderVars
+  , mkAppTy, mkAppTys, mkCastTy, mkCoercionTy, mkForAllTy, mkForAllTys
+  , mkTyConApp, mkPiTys
+  , noFreeVarsOfType
+  , splitPiTys, tyCoBinderType, tyCoBinderVar_maybe
+  )
+import GHC.Data.Pair
+  ( Pair(Pair) )
+import GHC.Types.Var
+  ( setTyVarKind )
+import GHC.Types.Var.Env
+  ( mkInScopeSet )
+import GHC.Types.Var.Set
+  ( TyCoVarSet )
+import GHC.Utils.Outputable
+  ( Outputable(ppr), (<+>)
+  , braces, text, vcat
+  )
+
+--------------------------------------------------------------------------------
+
+
+-- | A 'Reduction' is the result of an operation that rewrites a type @ty_in@.
+-- The 'Reduction' includes the rewritten type @ty_out@ and a 'Coercion' @co@
+-- such that @co :: ty_in ~ ty_out@, where the role of the coercion is determined
+-- by the context. That is, the LHS type of the coercion is the original type
+-- @ty_in@, while its RHS type is the rewritten type @ty_out@.
+--
+-- A Reduction is always homogeneous, unless it is wrapped inside a 'HetReduction',
+-- which separately stores the kind coercion.
+data Reduction =
+  Reduction
+    { reductionCoercion    :: Coercion
+    , reductionReducedType :: !Type
+    }
+
+-- | Stores a heterogeneous reduction.
+--
+-- The stored kind coercion must relate the kinds of the
+-- stored reduction. That is, in @HetReduction (Reduction co xi) kco@,
+-- we must have:
+--
+-- >  co :: ty ~ xi
+-- > kco :: typeKind ty ~ typeKind xi
+data HetReduction =
+  HetReduction
+    Reduction
+    MCoercion
+
+-- | Create a heterogeneous reduction.
+--
+-- Pre-condition: the provided kind coercion (second argument)
+-- relates the kinds of the stored reduction.
+-- That is, if the coercion stored in the 'Reduction' is of the form
+--
+-- > co :: ty ~ xi
+--
+-- Then the kind coercion supplied must be of the form:
+--
+-- > kco :: typeKind ty ~ typeKind xi
+mkHetReduction :: Reduction  -- ^ heterogeneous reduction
+               -> MCoercion  -- ^ kind coercion
+               -> HetReduction
+mkHetReduction redn mco = HetReduction redn mco
+{-# INLINE mkHetReduction #-}
+
+-- | Homogenise a heterogeneous reduction.
+--
+-- Given @HetReduction (Reduction co xi) kco@, with
+--
+-- >  co :: ty ~ xi
+-- > kco :: typeKind(ty) ~ typeKind(xi)
+--
+-- this returns the homogeneous reduction:
+--
+-- > hco :: ty ~ ( xi |> sym kco )
+homogeniseHetRedn :: Role -> HetReduction -> Reduction
+homogeniseHetRedn role (HetReduction redn kco)
+  = mkCoherenceRightMRedn role redn (mkSymMCo kco)
+{-# INLINE homogeniseHetRedn #-}
+
+-- | Create a 'Reduction' from a pair of a 'Coercion' and a 'Type.
+--
+-- Pre-condition: the RHS type of the coercion matches the provided type
+-- (perhaps up to zonking).
+--
+-- Use 'coercionRedn' when you only have the coercion.
+mkReduction :: Coercion -> Type -> Reduction
+mkReduction co ty = Reduction co ty
+{-# INLINE mkReduction #-}
+
+instance Outputable Reduction where
+  ppr redn =
+    braces $ vcat
+      [ text "reductionOriginalType:" <+> ppr (reductionOriginalType redn)
+      , text " reductionReducedType:" <+> ppr (reductionReducedType redn)
+      , text "    reductionCoercion:" <+> ppr (reductionCoercion redn)
+      ]
+
+-- | A 'Reduction' in which the 'Coercion' has 'Nominal' role.
+type ReductionN = Reduction
+
+-- | A 'Reduction' in which the 'Coercion' has 'Representational' role.
+type ReductionR = Reduction
+
+-- | Get the original, unreduced type corresponding to a 'Reduction'.
+--
+-- This is obtained by computing the LHS kind of the stored coercion,
+-- which may be slow.
+reductionOriginalType :: Reduction -> Type
+reductionOriginalType = coercionLKind . reductionCoercion
+{-# INLINE reductionOriginalType #-}
+
+-- | Turn a 'Coercion' into a 'Reduction'
+-- by inspecting the RHS type of the coercion.
+--
+-- Prefer using 'mkReduction' when you already know
+-- the RHS type of the coercion, to avoid computing it anew.
+coercionRedn :: Coercion -> Reduction
+coercionRedn co = Reduction co (coercionRKind co)
+{-# INLINE coercionRedn #-}
+
+-- | Downgrade the role of the coercion stored in the 'Reduction'.
+downgradeRedn :: Role -- ^ desired role
+              -> Role -- ^ current role
+              -> Reduction
+              -> Reduction
+downgradeRedn new_role old_role redn@(Reduction co _)
+  = redn { reductionCoercion = downgradeRole new_role old_role co }
+{-# INLINE downgradeRedn #-}
+
+-- | Downgrade the role of the coercion stored in the 'Reduction',
+-- from 'Nominal' to 'Representational'.
+mkSubRedn :: Reduction -> Reduction
+mkSubRedn redn@(Reduction co _) = redn { reductionCoercion = mkSubCo co }
+{-# INLINE mkSubRedn #-}
+
+-- | Compose a reduction with a coercion on the left.
+--
+-- Pre-condition: the provided coercion's RHS type must match the LHS type
+-- of the coercion that is stored in the reduction.
+mkTransRedn :: Coercion -> Reduction -> Reduction
+mkTransRedn co1 redn@(Reduction co2 _)
+  = redn { reductionCoercion = co1 `mkTransCo` co2 }
+{-# INLINE mkTransRedn #-}
+
+-- | The reflexive reduction.
+mkReflRedn :: Role -> Type -> Reduction
+mkReflRedn r ty = mkReduction (mkReflCo r ty) ty
+
+-- | Create a 'Reduction' from a kind cast, in which
+-- the casted type is the rewritten type.
+--
+-- Given @ty :: k1@, @mco :: k1 ~ k2@,
+-- produces the 'Reduction' @ty ~res_co~> (ty |> mco)@
+-- at the given 'Role'.
+mkGReflRightRedn :: Role -> Type -> CoercionN -> Reduction
+mkGReflRightRedn role ty co
+  = mkReduction
+      (mkGReflRightCo role ty co)
+      (mkCastTy ty co)
+{-# INLINE mkGReflRightRedn #-}
+
+-- | Create a 'Reduction' from a kind cast, in which
+-- the casted type is the rewritten type.
+--
+-- Given @ty :: k1@, @mco :: k1 ~ k2@,
+-- produces the 'Reduction' @ty ~res_co~> (ty |> mco)@
+-- at the given 'Role'.
+mkGReflRightMRedn :: Role -> Type -> MCoercion -> Reduction
+mkGReflRightMRedn role ty mco
+  = mkReduction
+      (mkGReflRightMCo role ty mco)
+      (mkCastTyMCo ty mco)
+{-# INLINE mkGReflRightMRedn #-}
+
+-- | Create a 'Reduction' from a kind cast, in which
+-- the casted type is the original (non-rewritten) type.
+--
+-- Given @ty :: k1@, @mco :: k1 ~ k2@,
+-- produces the 'Reduction' @(ty |> mco) ~res_co~> ty@
+-- at the given 'Role'.
+mkGReflLeftRedn :: Role -> Type -> CoercionN -> Reduction
+mkGReflLeftRedn role ty co
+  = mkReduction
+      (mkGReflLeftCo role ty co)
+      ty
+{-# INLINE mkGReflLeftRedn #-}
+
+-- | Create a 'Reduction' from a kind cast, in which
+-- the casted type is the original (non-rewritten) type.
+--
+-- Given @ty :: k1@, @mco :: k1 ~ k2@,
+-- produces the 'Reduction' @(ty |> mco) ~res_co~> ty@
+-- at the given 'Role'.
+mkGReflLeftMRedn :: Role -> Type -> MCoercion -> Reduction
+mkGReflLeftMRedn role ty mco
+  = mkReduction
+      (mkGReflLeftMCo role ty mco)
+      ty
+{-# INLINE mkGReflLeftMRedn #-}
+
+-- | Apply a cast to the result of a 'Reduction'.
+--
+-- Given a 'Reduction' @ty1 ~co1~> (ty2 :: k2)@ and a kind coercion @kco@
+-- with LHS kind @k2@, produce a new 'Reduction' @ty1 ~co2~> ( ty2 |> kco )@
+-- of the given 'Role' (which must match the role of the coercion stored
+-- in the 'Reduction' argument).
+mkCoherenceRightRedn :: Role -> Reduction -> CoercionN -> Reduction
+mkCoherenceRightRedn r (Reduction co1 ty2) kco
+  = mkReduction
+      (mkCoherenceRightCo r ty2 kco co1)
+      (mkCastTy ty2 kco)
+{-# INLINE mkCoherenceRightRedn #-}
+
+-- | Apply a cast to the result of a 'Reduction', using an 'MCoercionN'.
+--
+-- Given a 'Reduction' @ty1 ~co1~> (ty2 :: k2)@ and a kind coercion @mco@
+-- with LHS kind @k2@, produce a new 'Reduction' @ty1 ~co2~> ( ty2 |> mco )@
+-- of the given 'Role' (which must match the role of the coercion stored
+-- in the 'Reduction' argument).
+mkCoherenceRightMRedn :: Role -> Reduction -> MCoercion -> Reduction
+mkCoherenceRightMRedn r (Reduction co1 ty2) kco
+  = mkReduction
+      (mkCoherenceRightMCo r ty2 kco co1)
+      (mkCastTyMCo ty2 kco)
+{-# INLINE mkCoherenceRightMRedn #-}
+
+-- | Apply a cast to a 'Reduction', casting both the original and the reduced type.
+--
+-- Given @cast_co@ and 'Reduction' @ty ~co~> xi@, this function returns
+-- the 'Reduction' @(ty |> cast_co) ~return_co~> (xi |> cast_co)@
+-- of the given 'Role' (which must match the role of the coercion stored
+-- in the 'Reduction' argument).
+--
+-- Pre-condition: the 'Type' passed in is the same as the LHS type
+-- of the coercion stored in the 'Reduction'.
+mkCastRedn1 :: Role
+            -> Type      -- ^ original type
+            -> CoercionN -- ^ coercion to cast with
+            -> Reduction -- ^ rewritten type, with rewriting coercion
+            -> Reduction
+mkCastRedn1 r ty cast_co (Reduction co xi)
+  -- co :: ty ~r ty'
+  -- return_co :: (ty |> cast_co) ~r (ty' |> cast_co)
+  = mkReduction
+      (castCoercionKind1 co r ty xi cast_co)
+      (mkCastTy xi cast_co)
+{-# INLINE mkCastRedn1 #-}
+
+-- | Apply casts on both sides of a 'Reduction' (of the given 'Role').
+--
+-- Use 'mkCastRedn1' when you want to cast both the original and reduced types
+-- in a 'Reduction' using the same coercion.
+--
+-- Pre-condition: the 'Type' passed in is the same as the LHS type
+-- of the coercion stored in the 'Reduction'.
+mkCastRedn2 :: Role
+            -> Type      -- ^ original type
+            -> CoercionN -- ^ coercion to cast with on the left
+            -> Reduction -- ^ rewritten type, with rewriting coercion
+            -> CoercionN -- ^ coercion to cast with on the right
+            -> Reduction
+mkCastRedn2 r ty cast_co (Reduction nco nty) cast_co'
+  = mkReduction
+      (castCoercionKind2 nco r ty nty cast_co cast_co')
+      (mkCastTy nty cast_co')
+{-# INLINE mkCastRedn2 #-}
+
+-- | Apply one 'Reduction' to another.
+--
+-- Combines 'mkAppCo' and 'mkAppTy`.
+mkAppRedn :: Reduction -> Reduction -> Reduction
+mkAppRedn (Reduction co1 ty1) (Reduction co2 ty2)
+  = mkReduction (mkAppCo co1 co2) (mkAppTy ty1 ty2)
+{-# INLINE mkAppRedn #-}
+
+-- | Create a function 'Reduction'.
+--
+-- Combines 'mkFunCo' and 'mkFunTy'.
+mkFunRedn :: Role
+          -> AnonArgFlag
+          -> ReductionN -- ^ multiplicity reduction
+          -> Reduction  -- ^ argument reduction
+          -> Reduction  -- ^ result reduction
+          -> Reduction
+mkFunRedn r vis
+  (Reduction w_co w_ty)
+  (Reduction arg_co arg_ty)
+  (Reduction res_co res_ty)
+    = mkReduction
+        (mkFunCo r w_co arg_co res_co)
+        (mkFunTy vis w_ty arg_ty res_ty)
+{-# INLINE mkFunRedn #-}
+
+-- | Create a 'Reduction' associated to a Π type,
+-- from a kind 'Reduction' and a body 'Reduction'.
+--
+-- Combines 'mkForAllCo' and 'mkForAllTy'.
+mkForAllRedn :: ArgFlag
+             -> TyVar
+             -> ReductionN -- ^ kind reduction
+             -> Reduction  -- ^ body reduction
+             -> Reduction
+mkForAllRedn vis tv1 (Reduction h ki') (Reduction co ty)
+  = mkReduction
+      (mkForAllCo tv1 h co)
+      (mkForAllTy tv2 vis ty)
+  where
+    tv2 = setTyVarKind tv1 ki'
+{-# INLINE mkForAllRedn #-}
+
+-- | Create a 'Reduction' of a quantified type from a
+-- 'Reduction' of the body.
+--
+-- Combines 'mkHomoForAllCos' and 'mkForAllTys'.
+mkHomoForAllRedn :: [TyVarBinder] -> Reduction -> Reduction
+mkHomoForAllRedn bndrs (Reduction co ty)
+  = mkReduction
+      (mkHomoForAllCos (binderVars bndrs) co)
+      (mkForAllTys bndrs ty)
+{-# INLINE mkHomoForAllRedn #-}
+
+-- | Create a 'Reduction' from a coercion between coercions.
+--
+-- Combines 'mkProofIrrelCo' and 'mkCoercionTy'.
+mkProofIrrelRedn :: Role      -- ^ role of the created coercion, "r"
+                 -> CoercionN -- ^ co :: phi1 ~N phi2
+                 -> Coercion  -- ^ g1 :: phi1
+                 -> Coercion  -- ^ g2 :: phi2
+                 -> Reduction -- ^ res_co :: g1 ~r g2
+mkProofIrrelRedn role co g1 g2
+  = mkReduction
+      (mkProofIrrelCo role co g1 g2)
+      (mkCoercionTy g2)
+{-# INLINE mkProofIrrelRedn #-}
+
+-- | Create a reflexive 'Reduction' whose RHS is the given 'Coercion',
+-- with the specified 'Role'.
+mkReflCoRedn :: Role -> Coercion -> Reduction
+mkReflCoRedn role co
+  = mkReduction
+      (mkReflCo role co_ty)
+      co_ty
+  where
+    co_ty = mkCoercionTy co
+{-# INLINE mkReflCoRedn #-}
+
+-- | A collection of 'Reduction's where the coercions and the types are stored separately.
+--
+-- Use 'unzipRedns' to obtain 'Reductions' from a list of 'Reduction's.
+--
+-- This datatype is used in 'mkAppRedns', 'mkClassPredRedns' and 'mkTyConAppRedn',
+-- which expect separate types and coercions.
+--
+-- Invariant: the two stored lists are of the same length,
+-- and the RHS type of each coercion is the corresponding type.
+data Reductions = Reductions [Coercion] [Type]
+
+-- | Create 'Reductions' from individual lists of coercions and types.
+--
+-- The lists should be of the same length, and the RHS type of each coercion
+-- should match the specified type in the other list.
+mkReductions :: [Coercion] -> [Type] -> Reductions
+mkReductions cos tys = Reductions cos tys
+{-# INLINE mkReductions #-}
+
+-- | Combines 'mkAppCos' and 'mkAppTys'.
+mkAppRedns :: Reduction -> Reductions -> Reduction
+mkAppRedns (Reduction co ty) (Reductions cos tys)
+  = mkReduction (mkAppCos co cos) (mkAppTys ty tys)
+{-# INLINE mkAppRedns #-}
+
+-- | 'TyConAppCo' for 'Reduction's: combines 'mkTyConAppCo' and `mkTyConApp`.
+mkTyConAppRedn :: Role -> TyCon -> Reductions -> Reduction
+mkTyConAppRedn role tc (Reductions cos tys)
+  = mkReduction (mkTyConAppCo role tc cos) (mkTyConApp tc tys)
+{-# INLINE mkTyConAppRedn #-}
+
+-- | Reduce the arguments of a 'Class' 'TyCon'.
+mkClassPredRedn :: Class -> Reductions -> Reduction
+mkClassPredRedn cls (Reductions cos tys)
+  = mkReduction
+      (mkTyConAppCo Nominal (classTyCon cls) cos)
+      (mkClassPred cls tys)
+{-# INLINE mkClassPredRedn #-}
+
+-- | Obtain 'Reductions' from a list of 'Reduction's by unzipping.
+unzipRedns :: [Reduction] -> Reductions
+unzipRedns = foldr accRedn (Reductions [] [])
+  where
+    accRedn :: Reduction -> Reductions -> Reductions
+    accRedn (Reduction co xi) (Reductions cos xis)
+      = Reductions (co:cos) (xi:xis)
+{-# INLINE unzipRedns #-}
+
+--------------------------------------------------------------------------------
+
+data ArgsReductions =
+  ArgsReductions
+    {-# UNPACK #-} !Reductions
+    !MCoercion
+
+{-# INLINE simplifyArgsWorker #-}
+simplifyArgsWorker :: [TyCoBinder] -> Kind -> TyCoVarSet -> [Role] -> [Reduction] -> ArgsReductions
+simplifyArgsWorker orig_ki_binders orig_inner_ki orig_fvs
+                   orig_roles orig_simplified_args
+  = go orig_lc
+       orig_ki_binders orig_inner_ki
+       orig_roles orig_simplified_args
+  where
+    orig_lc = emptyLiftingContext $ mkInScopeSet $ orig_fvs
+
+    go :: LiftingContext -> [TyCoBinder] -> Kind -> [Role] -> [Reduction] -> ArgsReductions
+    go !lc binders inner_ki _ []
+      = ArgsReductions
+          (mkReductions [] [])
+          kind_co
+      where
+        final_kind = mkPiTys binders inner_ki
+        kind_co | noFreeVarsOfType final_kind = MRefl
+                | otherwise                   = MCo $ liftCoSubst Nominal lc final_kind
+
+    go lc (binder:binders) inner_ki (role:roles) (arg_redn:arg_redns)
+      =  let !kind_co = liftCoSubst Nominal lc (tyCoBinderType binder)
+             !(Reduction casted_co casted_xi)
+                      = mkCoherenceRightRedn role arg_redn kind_co
+         -- now, extend the lifting context with the new binding
+             !new_lc | Just tv <- tyCoBinderVar_maybe binder
+                     = extendLiftingContextAndInScope lc tv casted_co
+                     | otherwise
+                     = lc
+             !(ArgsReductions (Reductions cos xis) final_kind_co)
+               = go new_lc binders inner_ki roles arg_redns
+         in ArgsReductions
+              (Reductions (casted_co:cos) (casted_xi:xis))
+              final_kind_co
+
+    -- See Note [Last case in simplifyArgsWorker]
+    go lc [] inner_ki roles arg_redns
+      = let co1 = liftCoSubst Nominal lc inner_ki
+            co1_kind              = coercionKind co1
+            unrewritten_tys       = map reductionOriginalType arg_redns
+            (arg_cos, res_co)     = decomposePiCos co1 co1_kind unrewritten_tys
+            casted_args           = zipWith3 mkCoherenceRightRedn roles arg_redns arg_cos
+            zapped_lc             = zapLiftingContext lc
+            Pair rewritten_kind _ = co1_kind
+            (bndrs, new_inner)    = splitPiTys rewritten_kind
+
+            ArgsReductions redns_out res_co_out
+              = go zapped_lc bndrs new_inner roles casted_args
+        in
+          ArgsReductions redns_out (res_co `mkTransMCoR` res_co_out)
+
+    go _ _ _ _ _ = error "simplifyArgsWorker wandered into deeper water than usual"
+
+--------------------------------------------------------------------------------
+
+-- | Get the reverse of an 'MCoercion'
+mkSymMCo :: MCoercion -> MCoercion
+mkSymMCo MRefl    = MRefl
+mkSymMCo (MCo co) = MCo (mkSymCo co)
+
+mkGReflLeftMCo :: Role -> Type -> MCoercion -> Coercion
+mkGReflLeftMCo r ty MRefl    = mkReflCo r ty
+mkGReflLeftMCo r ty (MCo co) = mkGReflLeftCo r ty co
+
+mkGReflRightMCo :: Role -> Type -> MCoercion -> Coercion
+mkGReflRightMCo r ty MRefl    = mkReflCo r ty
+mkGReflRightMCo r ty (MCo co) = mkGReflRightCo r ty co
+
+-- | Cast a type by an 'MCoercion'
+mkCastTyMCo :: Type -> MCoercion -> Type
+mkCastTyMCo ty MRefl    = ty
+mkCastTyMCo ty (MCo co) = ty `mkCastTy` co
+
+-- | Like 'mkCoherenceRightCo', but with an 'MCoercion'
+mkCoherenceRightMCo :: Role -> Type -> MCoercion -> Coercion -> Coercion
+mkCoherenceRightMCo _ _  MRefl    co2 = co2
+mkCoherenceRightMCo r ty (MCo co) co2 = mkCoherenceRightCo r ty co co2
+
+mkTransMCoR :: Coercion -> MCoercion -> MCoercion
+mkTransMCoR co1 MRefl     = coToMCo co1
+mkTransMCoR co1 (MCo co2) = MCo (mkTransCo co1 co2)