Add 'docWithStyle' to improve codegenwip/specialize-hdoc2

This new combinator

docWithStyle :: IsOutput doc => doc -> (PprStyle -> SDoc) -> doc

let us remove the need for code to be polymorphic in HDoc
when not used in code style.

Metric Decrease:
    ManyConstructors
    T13035
    T1969

5 files changed, 78 insertions, 68 deletions

diff --git a/compiler/GHC/Types/CostCentre.hs b/compiler/GHC/Types/CostCentre.hs
index e20a4977ec..a41496f83c 100644
--- a/compiler/GHC/Types/CostCentre.hs
+++ b/compiler/GHC/Types/CostCentre.hs
@@ -265,10 +265,8 @@ instance Outputable CostCentre where
   ppr = pprCostCentre
 
 pprCostCentre :: IsLine doc => CostCentre -> doc
-pprCostCentre cc = docWithContext $ \ sty ->
-  if codeStyle (sdocStyle sty)
-  then ppCostCentreLbl cc
-  else ftext (costCentreUserNameFS cc)
+pprCostCentre cc = docWithStyle (ppCostCentreLbl cc)
+                                (\_ -> ftext (costCentreUserNameFS cc))
 {-# SPECIALISE pprCostCentre :: CostCentre -> SDoc #-}
 {-# SPECIALISE pprCostCentre :: CostCentre -> HLine #-} -- see Note [SPECIALIZE to HDoc] in GHC.Utils.Outputable
 
diff --git a/compiler/GHC/Types/Name.hs b/compiler/GHC/Types/Name.hs
index 1b70c4d910..7a069a573d 100644
--- a/compiler/GHC/Types/Name.hs
+++ b/compiler/GHC/Types/Name.hs
@@ -627,21 +627,30 @@ instance OutputableBndr Name where
 
 pprName :: forall doc. IsLine doc => Name -> doc
 pprName name@(Name {n_sort = sort, n_uniq = uniq, n_occ = occ})
-  = docWithContext $ \ctx ->
-    let sty = sdocStyle ctx
-        debug = sdocPprDebug ctx
-        listTuplePuns = sdocListTuplePuns ctx
-    in handlePuns listTuplePuns (namePun_maybe name) $
-    case sort of
-      WiredIn mod _ builtin   -> pprExternal debug sty uniq mod occ True  builtin
-      External mod            -> pprExternal debug sty uniq mod occ False UserSyntax
-      System                  -> pprSystem   debug sty uniq occ
-      Internal                -> pprInternal debug sty uniq occ
+  = docWithStyle codeDoc normalDoc
   where
-    -- Print GHC.Types.List as [], etc.
-    handlePuns :: Bool -> Maybe FastString -> doc -> doc
-    handlePuns True (Just pun) _ = ftext pun
-    handlePuns _    _          r = r
+   codeDoc = case sort of
+               WiredIn mod _ _ -> pprModule mod <> char '_' <> ppr_z_occ_name occ
+               External mod    -> pprModule mod <> char '_' <> ppr_z_occ_name occ
+                                  -- In code style, always qualify
+                                  -- ToDo: maybe we could print all wired-in things unqualified
+                                  --       in code style, to reduce symbol table bloat?
+               System          -> pprUniqueAlways uniq
+               Internal        -> pprUniqueAlways uniq
+
+   normalDoc sty =
+     getPprDebug $ \debug ->
+     sdocOption sdocListTuplePuns $ \listTuplePuns ->
+       handlePuns listTuplePuns (namePun_maybe name) $
+       case sort of
+         WiredIn mod _ builtin   -> pprExternal debug sty uniq mod occ True  builtin
+         External mod            -> pprExternal debug sty uniq mod occ False UserSyntax
+         System                  -> pprSystem   debug sty uniq occ
+         Internal                -> pprInternal debug sty uniq occ
+
+   handlePuns :: Bool -> Maybe FastString -> SDoc -> SDoc
+   handlePuns True (Just pun) _ = ftext pun
+   handlePuns _    _          r = r
 {-# SPECIALISE pprName :: Name -> SDoc #-}
 {-# SPECIALISE pprName :: Name -> HLine #-} -- see Note [SPECIALIZE to HDoc] in GHC.Utils.Outputable
 
@@ -674,12 +683,8 @@ pprTickyName this_mod name
 pprNameUnqualified :: Name -> SDoc
 pprNameUnqualified Name { n_occ = occ } = ppr_occ_name occ
 
-pprExternal :: IsLine doc => Bool -> PprStyle -> Unique -> Module -> OccName -> Bool -> BuiltInSyntax -> doc
+pprExternal :: Bool -> PprStyle -> Unique -> Module -> OccName -> Bool -> BuiltInSyntax -> SDoc
 pprExternal debug sty uniq mod occ is_wired is_builtin
-  | codeStyle sty = pprModule mod <> char '_' <> ppr_z_occ_name occ
-        -- In code style, always qualify
-        -- ToDo: maybe we could print all wired-in things unqualified
-        --       in code style, to reduce symbol table bloat?
   | debug         = pp_mod <> ppr_occ_name occ
                      <> braces (hsep [if is_wired then text "(w)" else empty,
                                       pprNameSpaceBrief (occNameSpace occ),
@@ -695,9 +700,8 @@ pprExternal debug sty uniq mod occ is_wired is_builtin
     pp_mod = ppUnlessOption sdocSuppressModulePrefixes
                (pprModule mod <> dot)
 
-pprInternal :: IsLine doc => Bool -> PprStyle -> Unique -> OccName -> doc
+pprInternal :: Bool -> PprStyle -> Unique -> OccName -> SDoc
 pprInternal debug sty uniq occ
-  | codeStyle sty  = pprUniqueAlways uniq
   | debug          = ppr_occ_name occ <> braces (hsep [pprNameSpaceBrief (occNameSpace occ),
                                                        pprUnique uniq])
   | dumpStyle sty  = ppr_occ_name occ <> ppr_underscore_unique uniq
@@ -706,9 +710,8 @@ pprInternal debug sty uniq occ
   | otherwise      = ppr_occ_name occ   -- User style
 
 -- Like Internal, except that we only omit the unique in Iface style
-pprSystem :: IsLine doc => Bool -> PprStyle -> Unique -> OccName -> doc
-pprSystem debug sty uniq occ
-  | codeStyle sty  = pprUniqueAlways uniq
+pprSystem :: Bool -> PprStyle -> Unique -> OccName -> SDoc
+pprSystem debug _sty uniq occ
   | debug          = ppr_occ_name occ <> ppr_underscore_unique uniq
                      <> braces (pprNameSpaceBrief (occNameSpace occ))
   | otherwise      = ppr_occ_name occ <> ppr_underscore_unique uniq
@@ -717,7 +720,7 @@ pprSystem debug sty uniq occ
                                 -- so print the unique
 
 
-pprModulePrefix :: IsLine doc => PprStyle -> Module -> OccName -> doc
+pprModulePrefix :: PprStyle -> Module -> OccName -> SDoc
 -- Print the "M." part of a name, based on whether it's in scope or not
 -- See Note [Printing original names] in GHC.Types.Name.Ppr
 pprModulePrefix sty mod occ = ppUnlessOption sdocSuppressModulePrefixes $
@@ -728,20 +731,20 @@ pprModulePrefix sty mod occ = ppUnlessOption sdocSuppressModulePrefixes $
                           <> pprModuleName (moduleName mod) <> dot    -- scope either
       NameUnqual       -> empty                   -- In scope unqualified
 
-pprUnique :: IsLine doc => Unique -> doc
+pprUnique :: Unique -> SDoc
 -- Print a unique unless we are suppressing them
 pprUnique uniq
   = ppUnlessOption sdocSuppressUniques $
       pprUniqueAlways uniq
 
-ppr_underscore_unique :: IsLine doc => Unique -> doc
+ppr_underscore_unique :: Unique -> SDoc
 -- Print an underscore separating the name from its unique
 -- But suppress it if we aren't printing the uniques anyway
 ppr_underscore_unique uniq
   = ppUnlessOption sdocSuppressUniques $
       char '_' <> pprUniqueAlways uniq
 
-ppr_occ_name :: IsLine doc => OccName -> doc
+ppr_occ_name :: OccName -> SDoc
 ppr_occ_name occ = ftext (occNameFS occ)
         -- Don't use pprOccName; instead, just print the string of the OccName;
         -- we print the namespace in the debug stuff above
diff --git a/compiler/GHC/Types/Name/Occurrence.hs b/compiler/GHC/Types/Name/Occurrence.hs
index 36add1cfea..e8d42eb0cf 100644
--- a/compiler/GHC/Types/Name/Occurrence.hs
+++ b/compiler/GHC/Types/Name/Occurrence.hs
@@ -201,7 +201,7 @@ pprNonVarNameSpace :: NameSpace -> SDoc
 pprNonVarNameSpace VarName = empty
 pprNonVarNameSpace ns = pprNameSpace ns
 
-pprNameSpaceBrief :: IsLine doc => NameSpace -> doc
+pprNameSpaceBrief :: NameSpace -> SDoc
 pprNameSpaceBrief DataName  = char 'd'
 pprNameSpaceBrief VarName   = char 'v'
 pprNameSpaceBrief TvName    = text "tv"
@@ -287,10 +287,9 @@ instance OutputableBndr OccName where
 
 pprOccName :: IsLine doc => OccName -> doc
 pprOccName (OccName sp occ)
-  = docWithContext $ \ sty ->
-    if codeStyle (sdocStyle sty)
-    then ztext (zEncodeFS occ)
-    else ftext occ <> whenPprDebug (braces (pprNameSpaceBrief sp))
+  = docWithStyle (ztext (zEncodeFS occ)) (\_ -> ftext occ <> whenPprDebug (braces (pprNameSpaceBrief sp)))
+{-# SPECIALIZE pprOccName :: OccName -> SDoc #-}
+{-# SPECIALIZE pprOccName :: OccName -> HLine #-} -- see Note [SPECIALIZE to HDoc] in GHC.Utils.Outputable
 
 {-
 ************************************************************************
diff --git a/compiler/GHC/Unit/Types.hs b/compiler/GHC/Unit/Types.hs
index 4fe2b932f6..7439ab7dde 100644
--- a/compiler/GHC/Unit/Types.hs
+++ b/compiler/GHC/Unit/Types.hs
@@ -166,7 +166,7 @@ instance Outputable InstantiatedModule where
 instance Outputable InstantiatedUnit where
   ppr = pprInstantiatedUnit
 
-pprInstantiatedUnit :: IsLine doc => InstantiatedUnit -> doc
+pprInstantiatedUnit :: InstantiatedUnit -> SDoc
 pprInstantiatedUnit uid =
       -- getPprStyle $ \sty ->
       pprUnitId cid <>
@@ -180,8 +180,6 @@ pprInstantiatedUnit uid =
      where
       cid   = instUnitInstanceOf uid
       insts = instUnitInsts uid
-{-# SPECIALIZE pprInstantiatedUnit :: InstantiatedUnit -> SDoc #-}
-{-# SPECIALIZE pprInstantiatedUnit :: InstantiatedUnit -> HLine #-} -- see Note [SPECIALIZE to HDoc] in GHC.Utils.Outputable
 
 -- | Class for types that are used as unit identifiers (UnitKey, UnitId, Unit)
 --
@@ -203,14 +201,13 @@ instance IsUnitId u => IsUnitId (GenUnit u) where
    unitFS HoleUnit                = holeFS
 
 pprModule :: IsLine doc => Module -> doc
-pprModule mod@(Module p n) = docWithContext (doc . sdocStyle)
+pprModule mod@(Module p n) = docWithStyle code doc
  where
-  doc sty
-    | codeStyle sty =
-        (if p == mainUnit
+  code = (if p == mainUnit
                 then empty -- never qualify the main package in code
                 else ztext (zEncodeFS (unitFS p)) <> char '_')
             <> pprModuleName n
+  doc sty
     | qualModule sty mod =
         case p of
           HoleUnit -> angleBrackets (pprModuleName n)
@@ -352,12 +349,10 @@ stableUnitCmp p1 p2 = unitFS p1 `lexicalCompareFS` unitFS p2
 instance Outputable Unit where
    ppr pk = pprUnit pk
 
-pprUnit :: IsLine doc => Unit -> doc
+pprUnit :: Unit -> SDoc
 pprUnit (RealUnit (Definite d)) = pprUnitId d
 pprUnit (VirtUnit uid) = pprInstantiatedUnit uid
 pprUnit HoleUnit       = ftext holeFS
-{-# SPECIALIZE pprUnit :: Unit -> SDoc #-}
-{-# SPECIALIZE pprUnit :: Unit -> HLine #-} -- see Note [SPECIALIZE to HDoc] in GHC.Utils.Outputable
 
 instance Show Unit where
     show = unitString
@@ -535,12 +530,8 @@ instance Uniquable UnitId where
 instance Outputable UnitId where
     ppr = pprUnitId
 
-pprUnitId :: IsLine doc => UnitId -> doc
-pprUnitId (UnitId fs) = dualLine (sdocOption sdocUnitIdForUser ($ fs)) (ftext fs)
-                        -- see Note [Pretty-printing UnitId] in GHC.Unit
-                        -- also see Note [dualLine and dualDoc] in GHC.Utils.Outputable
-{-# SPECIALIZE pprUnitId :: UnitId -> SDoc #-}
-{-# SPECIALIZE pprUnitId :: UnitId -> HLine #-} -- see Note [SPECIALIZE to HDoc] in GHC.Utils.Outputable
+pprUnitId :: UnitId -> SDoc
+pprUnitId (UnitId fs) = sdocOption sdocUnitIdForUser ($ fs)
 
 -- | A 'DefUnitId' is an 'UnitId' with the invariant that
 -- it only refers to a definite library; i.e., one we have generated
diff --git a/compiler/GHC/Utils/Outputable.hs b/compiler/GHC/Utils/Outputable.hs
index b4a21a314e..6bd48605d9 100644
--- a/compiler/GHC/Utils/Outputable.hs
+++ b/compiler/GHC/Utils/Outputable.hs
@@ -126,6 +126,7 @@ import GHC.Data.FastString
 import qualified GHC.Utils.Ppr as Pretty
 import qualified GHC.Utils.Ppr.Colour as Col
 import GHC.Utils.Ppr       ( Doc, Mode(..) )
+import GHC.Utils.Panic.Plain (assert)
 import GHC.Serialized
 import GHC.LanguageExtensions (Extension)
 import GHC.Utils.GlobalVars( unsafeHasPprDebug )
@@ -855,9 +856,10 @@ ppWhenOption f doc = sdocOption f $ \case
    False -> empty
 
 {-# INLINE CONLIKE ppUnlessOption #-}
-ppUnlessOption :: IsLine doc => (SDocContext -> Bool) -> doc -> doc
-ppUnlessOption f doc = docWithContext $
-                          \ctx -> if f ctx then empty else doc
+ppUnlessOption :: (SDocContext -> Bool) -> SDoc -> SDoc
+ppUnlessOption f doc = sdocOption f $ \case
+   True  -> empty
+   False -> doc
 
 -- | Apply the given colour\/style for the argument.
 --
@@ -1040,10 +1042,7 @@ instance Outputable ModuleName where
 
 pprModuleName :: IsLine doc => ModuleName -> doc
 pprModuleName (ModuleName nm) =
-    docWithContext $ \ctx ->
-    if codeStyle (sdocStyle ctx)
-        then ztext (zEncodeFS nm)
-        else ftext nm
+    docWithStyle (ztext (zEncodeFS nm)) (\_ -> ftext nm)
 {-# SPECIALIZE pprModuleName :: ModuleName -> SDoc #-}
 {-# SPECIALIZE pprModuleName :: ModuleName -> HLine #-} -- see Note [SPECIALIZE to HDoc]
 
@@ -1633,6 +1632,7 @@ IsOutput, that allows these combinators to be generic over both variants:
     class IsOutput doc where
       empty :: doc
       docWithContext :: (SDocContext -> doc) -> doc
+      docWithStyle :: doc -> (PprStyle -> SDoc) -> doc
 
     class IsOutput doc => IsLine doc
     class (IsOutput doc, IsLine (Line doc)) => IsDoc doc
@@ -1669,13 +1669,22 @@ arguments depending on the type they are instantiated at. They serve as a
 difficult to make completely equivalent under both printer implementations.
 
 These operations should generally be avoided, as they can result in surprising
-changes in behavior when the printer implementation is changed. However, in
-certain cases, the alternative is even worse. For example, we use dualLine in
-the implementation of pprUnitId, as the hack we use for printing unit ids
-(see Note [Pretty-printing UnitId] in GHC.Unit) is difficult to adapt to HLine
-and is not necessary for code paths that use it, anyway.
-
-Use these operations wisely. -}
+changes in behavior when the printer implementation is changed.
+Right now, they are used only when outputting debugging comments in
+codegen, as it is difficult to adapt that code to use HLine and not necessary.
+
+Use these operations wisely.
+
+Note [docWithStyle]
+~~~~~~~~~~~~~~~~~~~
+Sometimes when printing, we consult the printing style. This can be done
+with 'docWithStyle c f'. This is similar to 'docWithContext (f . sdocStyle)',
+but:
+* For code style, 'docWithStyle c f' will return 'c'.
+* For other styles, 'docWithStyle c f', will call 'f style', but expect
+  an SDoc rather than doc. This removes the need to write code polymorphic
+  in SDoc and HDoc, since the latter is used only for code style.
+-}
 
 -- | Represents a single line of output that can be efficiently printed directly
 -- to a 'System.IO.Handle' (actually a 'BufHandle').
@@ -1700,7 +1709,7 @@ pattern HDoc f <- HDoc' f
 {-# COMPLETE HDoc #-}
 
 bPutHDoc :: BufHandle -> SDocContext -> HDoc -> IO ()
-bPutHDoc h ctx (HDoc f) = f ctx h
+bPutHDoc h ctx (HDoc f) = assert (codeStyle (sdocStyle ctx)) (f ctx h)
 
 -- | A superclass for 'IsLine' and 'IsDoc' that provides an identity, 'empty',
 -- as well as access to the shared 'SDocContext'.
@@ -1709,6 +1718,7 @@ bPutHDoc h ctx (HDoc f) = f ctx h
 class IsOutput doc where
   empty :: doc
   docWithContext :: (SDocContext -> doc) -> doc
+  docWithStyle :: doc -> (PprStyle -> SDoc) -> doc  -- see Note [docWithStyle]
 
 -- | A class of types that represent a single logical line of text, with support
 -- for horizontal composition.
@@ -1779,6 +1789,11 @@ instance IsOutput SDoc where
   {-# INLINE CONLIKE empty #-}
   docWithContext = sdocWithContext
   {-# INLINE docWithContext #-}
+  docWithStyle c f = sdocWithContext (\ctx -> let sty = sdocStyle ctx
+                                              in if codeStyle sty then c
+                                                                  else f sty)
+                     -- see Note [docWithStyle]
+  {-# INLINE CONLIKE docWithStyle #-}
 
 instance IsLine SDoc where
   char c = docToSDoc $ Pretty.char c
@@ -1823,12 +1838,16 @@ instance IsOutput HLine where
   {-# INLINE empty #-}
   docWithContext f = HLine $ \ctx h -> runHLine (f ctx) ctx h
   {-# INLINE CONLIKE docWithContext #-}
+  docWithStyle c _ = HLine $ \ctx h -> assert (codeStyle (sdocStyle ctx)) (runHLine c ctx h)  -- see Note [docWithStyle]
+  {-# INLINE CONLIKE docWithStyle #-}
 
 instance IsOutput HDoc where
   empty = HDoc (\_ _ -> pure ())
   {-# INLINE empty #-}
   docWithContext f = HDoc $ \ctx h -> runHDoc (f ctx) ctx h
   {-# INLINE CONLIKE docWithContext #-}
+  docWithStyle c _ = HDoc $ \ctx h -> assert (codeStyle (sdocStyle ctx)) (runHDoc c ctx h)  -- see Note [docWithStyle]
+  {-# INLINE CONLIKE docWithStyle #-}
 
 instance IsLine HLine where
   char c = HLine (\_ h -> bPutChar h c)