Hide concat order for HsWrapper so Semigroup is legitwip/hswrapper-safe-elim

See comments, we use an OrdList of concatenations.

9 files changed, 196 insertions, 173 deletions

diff --git a/compiler/GHC/HsToCore/Binds.hs b/compiler/GHC/HsToCore/Binds.hs
index 7bd9506ca1..9d87bab419 100644
--- a/compiler/GHC/HsToCore/Binds.hs
+++ b/compiler/GHC/HsToCore/Binds.hs
@@ -11,9 +11,10 @@ lower levels it is preserved with @let@/@letrec@s).
 -}
 
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE ViewPatterns #-}
-{-# LANGUAGE FlexibleContexts #-}
 
 {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
 
@@ -357,7 +358,7 @@ dsAbsBinds dflags tyvars dicts exports
          return (ABE { abe_ext   = noExtField
                      , abe_poly  = global
                      , abe_mono  = local
-                     , abe_wrap  = WpHole
+                     , abe_wrap  = idHsWrapper
                      , abe_prags = SpecPrags [] })
 
 -- | This is where we apply INLINE and INLINABLE pragmas. All we need to
@@ -1110,31 +1111,34 @@ So for now, we ban them altogether as requested by #13290. See also #7398.
 -}
 
 dsHsWrapper :: HsWrapper -> DsM (CoreExpr -> CoreExpr)
-dsHsWrapper WpHole            = return $ \e -> e
-dsHsWrapper (WpTyApp ty)      = return $ \e -> App e (Type ty)
-dsHsWrapper (WpEvLam ev)      = return $ Lam ev
-dsHsWrapper (WpTyLam tv)      = return $ Lam tv
-dsHsWrapper (WpLet ev_binds)  = do { bs <- dsTcEvBinds ev_binds
-                                   ; return (mkCoreLets bs) }
-dsHsWrapper (WpCompose c1 c2) = do { w1 <- dsHsWrapper c1
-                                   ; w2 <- dsHsWrapper c2
-                                   ; return (w1 . w2) }
- -- See comments on WpFun in GHC.Tc.Types.Evidence for an explanation of what
- -- the specification of this clause is
-dsHsWrapper (WpFun c1 c2 t1 doc)
-                              = do { x  <- newSysLocalDsNoLP t1
-                                   ; w1 <- dsHsWrapper c1
-                                   ; w2 <- dsHsWrapper c2
-                                   ; let app f a = mkCoreAppDs (text "dsHsWrapper") f a
-                                         arg     = w1 (Var x)
-                                   ; (_, ok) <- askNoErrsDs $ dsNoLevPolyExpr arg doc
-                                   ; if ok
-                                     then return (\e -> (Lam x (w2 (app e arg))))
-                                     else return id }  -- this return is irrelevant
-dsHsWrapper (WpCast co)       = ASSERT(coercionRole co == Representational)
-                                return $ \e -> mkCastDs e co
-dsHsWrapper (WpEvApp tm)      = do { core_tm <- dsEvTerm tm
-                                   ; return (\e -> App e core_tm) }
+dsHsWrapper (HsWrapper steps) = fmap (foldr (.) id) $
+  traverse dsHsWrapperStep steps
+
+dsHsWrapperStep :: HsWrapperStep -> DsM (CoreExpr -> CoreExpr)
+dsHsWrapperStep = \case
+  WpTyApp ty      -> return $ \e -> App e (Type ty)
+  WpEvLam ev      -> return $ Lam ev
+  WpTyLam tv      -> return $ Lam tv
+  WpLet ev_binds  -> do
+    bs <- dsTcEvBinds ev_binds
+    return $ mkCoreLets bs
+  -- See comments on WpFun in GHC.Tc.Types.Evidence for an explanation of what
+  -- the specification of this clause is
+  WpFun c1 c2 t1 doc -> do
+    x  <- newSysLocalDsNoLP t1
+    w1 <- dsHsWrapper c1
+    w2 <- dsHsWrapper c2
+    let app f a = mkCoreAppDs (text "dsHsWrapper") f a
+        arg     = w1 (Var x)
+    (_, ok) <- askNoErrsDs $ dsNoLevPolyExpr arg doc
+    if ok
+      then return (\e -> (Lam x (w2 (app e arg))))
+      else return id  -- this return is irrelevant
+  WpCast co       -> ASSERT(coercionRole co == Representational)
+                     return $ \e -> mkCastDs e co
+  WpEvApp tm      -> do
+    core_tm <- dsEvTerm tm
+    return $ \e -> App e core_tm
 
 --------------------------------------
 dsTcEvBinds_s :: [TcEvBinds] -> DsM [CoreBind]
diff --git a/compiler/GHC/HsToCore/Match.hs b/compiler/GHC/HsToCore/Match.hs
index 55f2709cf9..cc002856d1 100644
--- a/compiler/GHC/HsToCore/Match.hs
+++ b/compiler/GHC/HsToCore/Match.hs
@@ -59,6 +59,7 @@ import GHC.Utils.Misc
 import GHC.Types.Name
 import GHC.Utils.Outputable
 import GHC.Types.Basic ( isGenerated, il_value, fl_value )
+import GHC.Data.OrdList ( strictlyZipWith )
 import GHC.Data.FastString
 import GHC.Types.Unique
 import GHC.Types.Unique.DFM
@@ -1100,15 +1101,16 @@ viewLExprEq (e1,_) (e2,_) = lexp e1 e2
     -- Coarser notions of equality are possible
     -- (e.g., reassociating compositions,
     --        equating different ways of writing a coercion)
-    wrap WpHole WpHole = True
-    wrap (WpCompose w1 w2) (WpCompose w1' w2') = wrap w1 w1' && wrap w2 w2'
-    wrap (WpFun w1 w2 _ _) (WpFun w1' w2' _ _) = wrap w1 w1' && wrap w2 w2'
-    wrap (WpCast co)       (WpCast co')        = co `eqCoercion` co'
-    wrap (WpEvApp et1)     (WpEvApp et2)       = et1 `ev_term` et2
-    wrap (WpTyApp t)       (WpTyApp t')        = eqType t t'
+    wrap (HsWrapper x) (HsWrapper y) = and $ strictlyZipWith wrapOne x y
+
+    wrapOne :: HsWrapperStep -> HsWrapperStep -> Bool
+    wrapOne (WpFun w1 w2 _ _) (WpFun w1' w2' _ _) = wrap w1 w1' && wrap w2 w2'
+    wrapOne (WpCast co)       (WpCast co')        = co `eqCoercion` co'
+    wrapOne (WpEvApp et1)     (WpEvApp et2)       = et1 `ev_term` et2
+    wrapOne (WpTyApp t)       (WpTyApp t')        = eqType t t'
     -- Enhancement: could implement equality for more wrappers
     --   if it seems useful (lams and lets)
-    wrap _ _ = False
+    wrapOne _ _ = False
 
     ---------
     ev_term :: EvTerm -> EvTerm -> Bool
diff --git a/compiler/GHC/HsToCore/PmCheck.hs b/compiler/GHC/HsToCore/PmCheck.hs
index f9de7c8282..6a2a879f8f 100644
--- a/compiler/GHC/HsToCore/PmCheck.hs
+++ b/compiler/GHC/HsToCore/PmCheck.hs
@@ -6,10 +6,10 @@ Pattern Matching Coverage Checking.
 
 {-# LANGUAGE CPP            #-}
 {-# LANGUAGE GADTs          #-}
+{-# LANGUAGE LambdaCase     #-}
+{-# LANGUAGE MultiWayIf     #-}
 {-# LANGUAGE TupleSections  #-}
 {-# LANGUAGE ViewPatterns   #-}
-{-# LANGUAGE MultiWayIf     #-}
-{-# LANGUAGE LambdaCase     #-}
 
 module GHC.HsToCore.PmCheck (
         -- Checking and printing
@@ -45,7 +45,9 @@ import GHC.Core.DataCon
 import GHC.Core.TyCon
 import GHC.Types.Var (EvVar)
 import GHC.Core.Coercion
-import GHC.Tc.Types.Evidence (HsWrapper(..), isIdHsWrapper)
+import GHC.Tc.Types.Evidence
+  ( HsWrapperStep(..), isIdHsWrapper
+  , stepFromWrapper )
 import GHC.Tc.Utils.TcType (evVarPred)
 import {-# SOURCE #-} GHC.HsToCore.Expr (dsExpr, dsLExpr, dsSyntaxExpr)
 import {-# SOURCE #-} GHC.HsToCore.Binds (dsHsWrapper)
@@ -445,8 +447,11 @@ translatePat fam_insts x pat = case pat of
   -- Generally the translation is
   -- pat |> co   ===>   let y = x |> co, pat <- y  where y is a match var of pat
   XPat (CoPat wrapper p _ty)
-    | isIdHsWrapper wrapper                   -> translatePat fam_insts x p
-    | WpCast co <-  wrapper, isReflexiveCo co -> translatePat fam_insts x p
+    | isIdHsWrapper wrapper
+      -> translatePat fam_insts x p
+    | Just (WpCast co) <- stepFromWrapper wrapper
+    , isReflexiveCo co
+      -> translatePat fam_insts x p
     | otherwise -> do
         (y, grds) <- translatePatV fam_insts p
         wrap_rhs_y <- dsHsWrapper wrapper
diff --git a/compiler/GHC/Tc/Errors/Hole.hs b/compiler/GHC/Tc/Errors/Hole.hs
index 7cbd8dbc0b..614eb488ed 100644
--- a/compiler/GHC/Tc/Errors/Hole.hs
+++ b/compiler/GHC/Tc/Errors/Hole.hs
@@ -793,9 +793,8 @@ tcFilterHoleFits limit typed_hole ht@(hole_ty, _) candidates =
 
 
     unfoldWrapper :: HsWrapper -> [Type]
-    unfoldWrapper = reverse . unfWrp'
+    unfoldWrapper = reverse . foldMap unfWrp' . unHsWrapper
       where unfWrp' (WpTyApp ty) = [ty]
-            unfWrp' (WpCompose w1 w2) = unfWrp' w1 ++ unfWrp' w2
             unfWrp' _ = []
 
 
diff --git a/compiler/GHC/Tc/Gen/Expr.hs b/compiler/GHC/Tc/Gen/Expr.hs
index 3a89daac0b..4369e53f26 100644
--- a/compiler/GHC/Tc/Gen/Expr.hs
+++ b/compiler/GHC/Tc/Gen/Expr.hs
@@ -1986,7 +1986,7 @@ tcTagToEnum expr fun args app_res_ty res_ty
          check_enumeration ty' rep_tc
        ; let val_arg = dropWhile (not . isHsValArg) args
              rep_ty  = mkTyConApp rep_tc rep_args
-             fun'    = mkHsWrap (WpTyApp rep_ty) fun
+             fun'    = mkHsWrap (wrapperFromStep $ WpTyApp rep_ty) fun
              expr'   = applyHsArgs fun' val_arg
              df_wrap = mkWpCastR (mkTcSymCo coi)
        ; return (mkHsWrap df_wrap expr') }}}}}
diff --git a/compiler/GHC/Tc/Gen/Match.hs b/compiler/GHC/Tc/Gen/Match.hs
index 857470b155..caabfd8956 100644
--- a/compiler/GHC/Tc/Gen/Match.hs
+++ b/compiler/GHC/Tc/Gen/Match.hs
@@ -518,7 +518,7 @@ tcLcStmt m_tc ctxt (TransStmt { trS_form = form, trS_stmts = stmts
                              poly_arg_ty `mkVisFunTy` poly_res_ty
 
        ; using' <- tcCheckExpr using using_poly_ty
-       ; let final_using = fmap (mkHsWrap (WpTyApp tup_ty)) using'
+       ; let final_using = fmap (mkHsWrap (wrapperFromStep $ WpTyApp tup_ty)) using'
 
              -- 'stmts' returns a result of type (m1_ty tuple_ty),
              -- typically something like [(Int,Bool,Int)]
@@ -704,7 +704,7 @@ tcMcStmt ctxt (TransStmt { trS_stmts = stmts, trS_bndrs = bindersMap
        -- using :: ((a,b,c)->t) -> m1 (a,b,c) -> m2 (n (a,b,c))
 
        ; using' <- tcCheckExpr using using_poly_ty
-       ; let final_using = fmap (mkHsWrap (WpTyApp tup_ty)) using'
+       ; let final_using = fmap (mkHsWrap (wrapperFromStep $ WpTyApp tup_ty)) using'
 
        --------------- Building the bindersMap ----------------
        ; let mk_n_bndr :: Name -> TcId -> TcId
diff --git a/compiler/GHC/Tc/Types/Evidence.hs b/compiler/GHC/Tc/Types/Evidence.hs
index 843a12cbcf..1dea058941 100644
--- a/compiler/GHC/Tc/Types/Evidence.hs
+++ b/compiler/GHC/Tc/Types/Evidence.hs
@@ -1,12 +1,18 @@
 -- (c) The University of Glasgow 2006
 
-{-# LANGUAGE CPP, DeriveDataTypeable #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedLists #-}
 
 module GHC.Tc.Types.Evidence (
 
-  -- * HsWrapper
-  HsWrapper(..),
+  -- * HsWrappery
+  HsWrapper(..), HsWrapperStep(..),
+  wrapperFromStep, stepFromWrapper,
   (<.>), mkWpTyApps, mkWpEvApps, mkWpEvVarApps, mkWpTyLams,
   mkWpLams, mkWpLet, mkWpCastN, mkWpCastR, collectHsWrapBinders,
   mkWpFun, idHsWrapper, isIdHsWrapper,
@@ -77,6 +83,7 @@ import GHC.Core.FVs   ( exprSomeFreeVars )
 
 import GHC.Utils.Misc
 import GHC.Data.Bag
+import GHC.Data.OrdList
 import qualified Data.Data as Data
 import GHC.Utils.Outputable
 import GHC.Types.SrcLoc
@@ -188,16 +195,24 @@ maybeTcSubCo ReprEq = mkTcSubCo
 ************************************************************************
 -}
 
-data HsWrapper
-  = WpHole                      -- The identity coercion
-
-  | WpCompose HsWrapper HsWrapper
-       -- (wrap1 `WpCompose` wrap2)[e] = wrap1[ wrap2[ e ]]
-       --
-       -- Hence  (\a. []) `WpCompose` (\b. []) = (\a b. [])
-       -- But    ([] a)   `WpCompose` ([] b)   = ([] b a)
-
-  | WpFun HsWrapper HsWrapper TcType SDoc
+-- | A defunctionalized transformation on Core expressions.
+--
+-- On one hand, we want concatenation to be fast. On the other, we don't want
+-- the exact concatenation order to be observable as it would violate
+-- associativity. A great way to have our cake and it two is to use OrdList,
+-- which has O(1) concatenation but is abstract.
+--
+-- (wrap1 <.> wrap2)[e] = wrap1[ wrap2[ e ]]
+--
+-- Hence  (\a. []) <.> (\b. []) = (\a b. [])
+-- But    ([] a)   <.> ([] b)   = ([] b a)
+newtype HsWrapper = HsWrapper
+  { unHsWrapper :: OrdList HsWrapperStep
+  } deriving (Data.Data)
+    deriving newtype (Semigroup, Monoid)
+
+data HsWrapperStep
+  = WpFun HsWrapper HsWrapper TcType SDoc
        -- (WpFun wrap1 wrap2 t1)[e] = \(x:t1). wrap2[ e wrap1[x] ]
        -- So note that if  wrap1 :: exp_arg <= act_arg
        --                  wrap2 :: act_res <= exp_res
@@ -227,9 +242,7 @@ data HsWrapper
 
 -- Cannot derive Data instance because SDoc is not Data (it stores a function).
 -- So we do it manually:
-instance Data.Data HsWrapper where
-  gfoldl _ z WpHole             = z WpHole
-  gfoldl k z (WpCompose a1 a2)  = z WpCompose `k` a1 `k` a2
+instance Data.Data HsWrapperStep where
   gfoldl k z (WpFun a1 a2 a3 _) = z wpFunEmpty `k` a1 `k` a2 `k` a3
   gfoldl k z (WpCast a1)        = z WpCast `k` a1
   gfoldl k z (WpEvLam a1)       = z WpEvLam `k` a1
@@ -239,18 +252,14 @@ instance Data.Data HsWrapper where
   gfoldl k z (WpLet a1)         = z WpLet `k` a1
 
   gunfold k z c = case Data.constrIndex c of
-                    1 -> z WpHole
-                    2 -> k (k (z WpCompose))
-                    3 -> k (k (k (z wpFunEmpty)))
-                    4 -> k (z WpCast)
-                    5 -> k (z WpEvLam)
-                    6 -> k (z WpEvApp)
-                    7 -> k (z WpTyLam)
-                    8 -> k (z WpTyApp)
+                    1 -> k (k (k (z wpFunEmpty)))
+                    2 -> k (z WpCast)
+                    3 -> k (z WpEvLam)
+                    4 -> k (z WpEvApp)
+                    5 -> k (z WpTyLam)
+                    6 -> k (z WpTyApp)
                     _ -> k (z WpLet)
 
-  toConstr WpHole          = wpHole_constr
-  toConstr (WpCompose _ _) = wpCompose_constr
   toConstr (WpFun _ _ _ _) = wpFun_constr
   toConstr (WpCast _)      = wpCast_constr
   toConstr (WpEvLam _)     = wpEvLam_constr
@@ -261,36 +270,21 @@ instance Data.Data HsWrapper where
 
   dataTypeOf _ = hsWrapper_dataType
 
--- | The Semigroup instance is a bit fishy, since @WpCompose@, as a data
--- constructor, is "syntactic" and not associative. Concretely, if @a@, @b@,
--- and @c@ aren't @WpHole@:
---
--- > (a <> b) <> c ?= a <> (b <> c)
---
--- ==>
---
--- > (a `WpCompose` b) `WpCompose` c /= @ a `WpCompose` (b `WpCompose` c)
---
--- However these two associations are are "semantically equal" in the sense
--- that they produce equal functions when passed to
--- @GHC.HsToCore.Binds.dsHsWrapper@.
-instance S.Semigroup HsWrapper where
-  (<>) = (<.>)
+wrapperFromStep :: HsWrapperStep -> HsWrapper
+wrapperFromStep = HsWrapper . unitOL
 
-instance Monoid HsWrapper where
-  mempty = WpHole
+stepFromWrapper :: HsWrapper -> Maybe HsWrapperStep
+stepFromWrapper (HsWrapper steps) = viewSingle steps
 
 hsWrapper_dataType :: Data.DataType
 hsWrapper_dataType
   = Data.mkDataType "HsWrapper"
-      [ wpHole_constr, wpCompose_constr, wpFun_constr, wpCast_constr
+      [ wpFun_constr, wpCast_constr
       , wpEvLam_constr, wpEvApp_constr, wpTyLam_constr, wpTyApp_constr
       , wpLet_constr]
 
-wpHole_constr, wpCompose_constr, wpFun_constr, wpCast_constr, wpEvLam_constr,
+wpFun_constr, wpCast_constr, wpEvLam_constr,
   wpEvApp_constr, wpTyLam_constr, wpTyApp_constr, wpLet_constr :: Data.Constr
-wpHole_constr    = mkHsWrapperConstr "WpHole"
-wpCompose_constr = mkHsWrapperConstr "WpCompose"
 wpFun_constr     = mkHsWrapperConstr "WpFun"
 wpCast_constr    = mkHsWrapperConstr "WpCast"
 wpEvLam_constr   = mkHsWrapperConstr "WpEvLam"
@@ -302,13 +296,11 @@ wpLet_constr     = mkHsWrapperConstr "WpLet"
 mkHsWrapperConstr :: String -> Data.Constr
 mkHsWrapperConstr name = Data.mkConstr hsWrapper_dataType name [] Data.Prefix
 
-wpFunEmpty :: HsWrapper -> HsWrapper -> TcType -> HsWrapper
+wpFunEmpty :: HsWrapper -> HsWrapper -> TcType -> HsWrapperStep
 wpFunEmpty c1 c2 t1 = WpFun c1 c2 t1 empty
 
 (<.>) :: HsWrapper -> HsWrapper -> HsWrapper
-WpHole <.> c = c
-c <.> WpHole = c
-c1 <.> c2    = c1 `WpCompose` c2
+(<.>) = (S.<>)
 
 mkWpFun :: HsWrapper -> HsWrapper
         -> TcType    -- the "from" type of the first wrapper
@@ -316,23 +308,25 @@ mkWpFun :: HsWrapper -> HsWrapper
                      -- second wrapper is the identity)
         -> SDoc      -- what caused you to want a WpFun? Something like "When converting ..."
         -> HsWrapper
-mkWpFun WpHole       WpHole       _  _  _ = WpHole
-mkWpFun WpHole       (WpCast co2) t1 _  _ = WpCast (mkTcFunCo Representational (mkTcRepReflCo t1) co2)
-mkWpFun (WpCast co1) WpHole       _  t2 _ = WpCast (mkTcFunCo Representational (mkTcSymCo co1) (mkTcRepReflCo t2))
-mkWpFun (WpCast co1) (WpCast co2) _  _  _ = WpCast (mkTcFunCo Representational (mkTcSymCo co1) co2)
-mkWpFun co1          co2          t1 _  d = WpFun co1 co2 t1 d
+mkWpFun (HsWrapper x) (HsWrapper y) = go x y
+  where
+    go []           []           _  _  _ = HsWrapper $ []
+    go []           [WpCast co2] t1 _  _ = HsWrapper $ [WpCast $ mkTcFunCo Representational (mkTcRepReflCo t1) co2]
+    go [WpCast co1] []           _  t2 _ = HsWrapper $ [WpCast $ mkTcFunCo Representational (mkTcSymCo co1) (mkTcRepReflCo t2)]
+    go [WpCast co1] [WpCast co2] _  _  _ = HsWrapper $ [WpCast $ mkTcFunCo Representational (mkTcSymCo co1) co2]
+    go co1          co2          t1 _  d = HsWrapper $ [WpFun (HsWrapper co1) (HsWrapper co2) t1 d]
 
 mkWpCastR :: TcCoercionR -> HsWrapper
 mkWpCastR co
-  | isTcReflCo co = WpHole
+  | isTcReflCo co = idHsWrapper
   | otherwise     = ASSERT2(tcCoercionRole co == Representational, ppr co)
-                    WpCast co
+                    wrapperFromStep (WpCast co)
 
 mkWpCastN :: TcCoercionN -> HsWrapper
 mkWpCastN co
-  | isTcReflCo co = WpHole
+  | isTcReflCo co = idHsWrapper
   | otherwise     = ASSERT2(tcCoercionRole co == Nominal, ppr co)
-                    WpCast (mkTcSubCo co)
+                    wrapperFromStep (WpCast $ mkTcSubCo co)
     -- The mkTcSubCo converts Nominal to Representational
 
 mkWpTyApps :: [Type] -> HsWrapper
@@ -352,23 +346,22 @@ mkWpLams ids = mk_co_lam_fn WpEvLam ids
 
 mkWpLet :: TcEvBinds -> HsWrapper
 -- This no-op is a quite a common case
-mkWpLet (EvBinds b) | isEmptyBag b = WpHole
-mkWpLet ev_binds                   = WpLet ev_binds
+mkWpLet (EvBinds b) | isEmptyBag b = idHsWrapper
+mkWpLet ev_binds                   = wrapperFromStep $ WpLet ev_binds
 
-mk_co_lam_fn :: (a -> HsWrapper) -> [a] -> HsWrapper
-mk_co_lam_fn f as = foldr (\x wrap -> f x <.> wrap) WpHole as
+mk_co_lam_fn :: (a -> HsWrapperStep) -> [a] -> HsWrapper
+mk_co_lam_fn f as = HsWrapper $ foldr (\x wrap -> f x `consOL` wrap) [] as
 
-mk_co_app_fn :: (a -> HsWrapper) -> [a] -> HsWrapper
+mk_co_app_fn :: (a -> HsWrapperStep) -> [a] -> HsWrapper
 -- For applications, the *first* argument must
 -- come *last* in the composition sequence
-mk_co_app_fn f as = foldr (\x wrap -> wrap <.> f x) WpHole as
+mk_co_app_fn f as = HsWrapper $ foldr (\x wrap -> wrap `snocOL` f x) [] as
 
 idHsWrapper :: HsWrapper
-idHsWrapper = WpHole
+idHsWrapper = HsWrapper nilOL
 
 isIdHsWrapper :: HsWrapper -> Bool
-isIdHsWrapper WpHole = True
-isIdHsWrapper _      = False
+isIdHsWrapper = isNilOL . unHsWrapper
 
 hsWrapDictBinders :: HsWrapper -> Bag DictId
 -- ^ Identifies the /lambda-bound/ dictionaries of an 'HsWrapper'. This is used
@@ -379,12 +372,10 @@ hsWrapDictBinders :: HsWrapper -> Bag DictId
 -- either superclasses of lambda-bound ones, or (extremely numerous) results of
 -- binding Wanted dictionaries.  We definitely don't want all those cluttering
 -- up the Given dictionaries for pattern-match overlap checking!
-hsWrapDictBinders wrap = go wrap
+hsWrapDictBinders = foldMap go . unHsWrapper
  where
    go (WpEvLam dict_id)   = unitBag dict_id
-   go (w1 `WpCompose` w2) = go w1 `unionBags` go w2
-   go (WpFun _ w _ _)     = go w
-   go WpHole              = emptyBag
+   go (WpFun _ w _ _)     = hsWrapDictBinders w
    go (WpCast  {})        = emptyBag
    go (WpEvApp {})        = emptyBag
    go (WpTyLam {})        = emptyBag
@@ -394,17 +385,17 @@ hsWrapDictBinders wrap = go wrap
 collectHsWrapBinders :: HsWrapper -> ([Var], HsWrapper)
 -- Collect the outer lambda binders of a HsWrapper,
 -- stopping as soon as you get to a non-lambda binder
-collectHsWrapBinders wrap = go wrap []
+collectHsWrapBinders = gos
   where
     -- go w ws = collectHsWrapBinders (w <.> w1 <.> ... <.> wn)
-    go :: HsWrapper -> [HsWrapper] -> ([Var], HsWrapper)
+    go :: HsWrapperStep -> HsWrapper -> ([Var], HsWrapper)
     go (WpEvLam v)       wraps = add_lam v (gos wraps)
     go (WpTyLam v)       wraps = add_lam v (gos wraps)
-    go (WpCompose w1 w2) wraps = go w1 (w2:wraps)
-    go wrap              wraps = ([], foldl' (<.>) wrap wraps)
+    go wrap              wraps = ([], HsWrapper $ consOL wrap $ unHsWrapper wraps)
 
-    gos []     = ([], WpHole)
-    gos (w:ws) = go w ws
+    gos (HsWrapper w) = case unConsOL w of
+      Nothing -> gos $ idHsWrapper
+      Just (w, ws) -> go w $ HsWrapper ws
 
     add_lam v (vs,w) = (v:vs, w)
 
@@ -932,16 +923,17 @@ pprHsWrapper :: HsWrapper -> (Bool -> SDoc) -> SDoc
 --    it's in a position that needs parens for a non-atomic thing
 pprHsWrapper wrap pp_thing_inside
   = sdocOption sdocPrintTypecheckerElaboration $ \case
-      True  -> help pp_thing_inside wrap False
+      True  -> helps pp_thing_inside wrap False
       False -> pp_thing_inside False
   where
-    help :: (Bool -> SDoc) -> HsWrapper -> Bool -> SDoc
+    helps :: (Bool -> SDoc) -> HsWrapper -> Bool -> SDoc
+    helps it = foldr (flip help) it . unHsWrapper
+
+    help :: (Bool -> SDoc) -> HsWrapperStep -> Bool -> SDoc
     -- True  <=> appears in function application position
     -- False <=> appears as body of let or lambda
-    help it WpHole             = it
-    help it (WpCompose f1 f2)  = help (help it f2) f1
     help it (WpFun f1 f2 t1 _) = add_parens $ text "\\(x" <> dcolon <> ppr t1 <> text ")." <+>
-                                              help (\_ -> it True <+> help (\_ -> text "x") f1 True) f2 False
+                                              helps (\_ -> it True <+> helps (\_ -> text "x") f1 True) f2 False
     help it (WpCast co)   = add_parens $ sep [it False, nest 2 (text "|>"
                                               <+> pprParendCo co)]
     help it (WpEvApp id)  = no_parens  $ sep [it True, nest 2 (ppr id)]
diff --git a/compiler/GHC/Tc/Utils/Zonk.hs b/compiler/GHC/Tc/Utils/Zonk.hs
index 4372a39e9d..ef2578978f 100644
--- a/compiler/GHC/Tc/Utils/Zonk.hs
+++ b/compiler/GHC/Tc/Utils/Zonk.hs
@@ -4,9 +4,11 @@
 
 -}
 
-{-# LANGUAGE CPP, TupleSections #-}
-{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE ViewPatterns #-}
 
 {-# OPTIONS_GHC -Wno-incomplete-record-updates #-}
@@ -85,7 +87,9 @@ import GHC.Core
 import {-# SOURCE #-} GHC.Tc.Gen.Splice (runTopSplice)
 
 import Control.Monad
+import Control.Monad.Trans.State
 import Data.List  ( partition )
+import Data.Tuple ( swap )
 import Control.Arrow ( second )
 
 {-
@@ -1037,27 +1041,32 @@ zonk_cmd_top env (HsCmdTop (CmdTopTc stack_tys ty ids) cmd)
 
 -------------------------------------------------------------------------
 zonkCoFn :: ZonkEnv -> HsWrapper -> TcM (ZonkEnv, HsWrapper)
-zonkCoFn env WpHole   = return (env, WpHole)
-zonkCoFn env (WpCompose c1 c2) = do { (env1, c1') <- zonkCoFn env c1
-                                    ; (env2, c2') <- zonkCoFn env1 c2
-                                    ; return (env2, WpCompose c1' c2') }
-zonkCoFn env (WpFun c1 c2 t1 d) = do { (env1, c1') <- zonkCoFn env c1
-                                     ; (env2, c2') <- zonkCoFn env1 c2
-                                     ; t1'         <- zonkTcTypeToTypeX env2 t1
-                                     ; return (env2, WpFun c1' c2' t1' d) }
-zonkCoFn env (WpCast co) = do { co' <- zonkCoToCo env co
-                              ; return (env, WpCast co') }
-zonkCoFn env (WpEvLam ev)   = do { (env', ev') <- zonkEvBndrX env ev
-                                 ; return (env', WpEvLam ev') }
-zonkCoFn env (WpEvApp arg)  = do { arg' <- zonkEvTerm env arg
-                                 ; return (env, WpEvApp arg') }
-zonkCoFn env (WpTyLam tv)   = ASSERT( isImmutableTyVar tv )
-                              do { (env', tv') <- zonkTyBndrX env tv
-                                 ; return (env', WpTyLam tv') }
-zonkCoFn env (WpTyApp ty)   = do { ty' <- zonkTcTypeToTypeX env ty
-                                 ; return (env, WpTyApp ty') }
-zonkCoFn env (WpLet bs)     = do { (env1, bs') <- zonkTcEvBinds env bs
-                                 ; return (env1, WpLet bs') }
+zonkCoFn env0 items0 = fmap swap $ runStateT (gos items0) env0
+  where
+    gos :: HsWrapper -> StateT ZonkEnv TcM HsWrapper
+    gos = fmap HsWrapper . traverse go . unHsWrapper
+
+    go :: HsWrapperStep -> StateT ZonkEnv TcM HsWrapperStep
+    go = \case
+      WpFun c1 c2 t1 d -> do
+        c1' <- gos c1
+        c2' <- gos c2
+        t1' <- liftWith zonkTcTypeToTypeX t1
+        return $ WpFun c1' c2' t1' d
+      WpCast co   -> WpCast <$> liftWith zonkCoToCo co
+      WpEvLam ev  -> WpEvLam <$> liftWith' zonkEvBndrX ev
+      WpEvApp arg -> WpEvApp <$> liftWith zonkEvTerm arg
+      WpTyLam tv -> do
+        MASSERT( isImmutableTyVar tv )
+        WpTyLam <$> liftWith' zonkTyBndrX tv
+      WpTyApp ty  -> WpTyApp <$> liftWith zonkTcTypeToTypeX ty
+      WpLet bs    -> WpLet <$> liftWith' zonkTcEvBinds bs
+
+    liftWith :: (ZonkEnv -> t -> TcM a) -> t -> StateT ZonkEnv TcM a
+    liftWith f x = StateT $ \env -> (,) <$> f env x <*> pure env
+
+    liftWith' :: (ZonkEnv -> t -> TcM (ZonkEnv, a)) -> t -> StateT ZonkEnv TcM a
+    liftWith' f x = StateT $ \env -> swap <$> f env x
 
 -------------------------------------------------------------------------
 zonkOverLit :: ZonkEnv -> HsOverLit GhcTcId -> TcM (HsOverLit GhcTc)
diff --git a/testsuite/tests/parser/should_compile/DumpTypecheckedAst.stderr b/testsuite/tests/parser/should_compile/DumpTypecheckedAst.stderr
index 80fc356925..a99f151995 100644
--- a/testsuite/tests/parser/should_compile/DumpTypecheckedAst.stderr
+++ b/testsuite/tests/parser/should_compile/DumpTypecheckedAst.stderr
@@ -541,10 +541,12 @@
             ({ <no location info> }
              (XExpr
               (HsWrap
-               (WpTyApp
-                (TyConApp
-                 ({abstract:TyCon})
-                 []))
+               (HsWrapper
+                (One
+                 (WpTyApp
+                  (TyConApp
+                   ({abstract:TyCon})
+                   []))))
                (HsConLikeOut
                 (NoExtField)
                 ({abstract:ConLike})))))
@@ -565,10 +567,12 @@
                 ({ <no location info> }
                  (XExpr
                   (HsWrap
-                   (WpTyApp
-                    (TyConApp
-                     ({abstract:TyCon})
-                     []))
+                   (HsWrapper
+                    (One
+                     (WpTyApp
+                      (TyConApp
+                       ({abstract:TyCon})
+                       []))))
                    (HsConLikeOut
                     (NoExtField)
                     ({abstract:ConLike})))))
@@ -589,10 +593,12 @@
                     ({ <no location info> }
                      (XExpr
                       (HsWrap
-                       (WpTyApp
-                        (TyConApp
-                         ({abstract:TyCon})
-                         []))
+                       (HsWrapper
+                        (One
+                         (WpTyApp
+                          (TyConApp
+                           ({abstract:TyCon})
+                           []))))
                        (HsConLikeOut
                         (NoExtField)
                         ({abstract:ConLike})))))
@@ -604,10 +610,12 @@
                   ({ <no location info> }
                    (XExpr
                     (HsWrap
-                     (WpTyApp
-                      (TyConApp
-                       ({abstract:TyCon})
-                       []))
+                     (HsWrapper
+                      (One
+                       (WpTyApp
+                        (TyConApp
+                         ({abstract:TyCon})
+                         []))))
                      (HsConLikeOut
                       (NoExtField)
                       ({abstract:ConLike})))))))))))))))))))))
@@ -633,10 +641,12 @@
       ({ <no location info> }
        (XExpr
         (HsWrap
-         (WpTyApp
-          (TyConApp
-           ({abstract:TyCon})
-           []))
+         (HsWrapper
+          (One
+           (WpTyApp
+            (TyConApp
+             ({abstract:TyCon})
+             []))))
          (HsConLikeOut
           (NoExtField)
           ({abstract:ConLike})))))))))
@@ -695,14 +705,16 @@
       (NoExtField)
       {Var: main}
       {Var: main}
-      (WpHole)
+      (HsWrapper
+       (None))
       (SpecPrags
        []))]
     [({abstract:TcEvBinds})]
     {Bag(Located (HsBind Var)):
      [({ DumpTypecheckedAst.hs:18:1-23 }
        (FunBind
-        (WpHole)
+        (HsWrapper
+         (None))
         ({ DumpTypecheckedAst.hs:18:1-4 }
          {Var: main})
         (MG