Add an INLINE pragma to Control.Category.>>>

This fixes #18013 by adding INLINE pragmas to both Control.Category.>>>
and GHC.Desugar.>>>. The functional change in this patch is tiny (just
two lines of pragmas!), but an accompanying Note explains in gory
detail what’s going on.

2 files changed, 68 insertions, 4 deletions

diff --git a/libraries/base/Control/Category.hs b/libraries/base/Control/Category.hs
index 96f0c33aed..c033c7618e 100644
--- a/libraries/base/Control/Category.hs
+++ b/libraries/base/Control/Category.hs
@@ -77,3 +77,67 @@ instance Category Coercion where
 -- | Left-to-right composition
 (>>>) :: Category cat => cat a b -> cat b c -> cat a c
 f >>> g = g . f
+{-# INLINE (>>>) #-} -- see Note [INLINE on >>>]
+
+{- Note [INLINE on >>>]
+~~~~~~~~~~~~~~~~~~~~~~~
+It’s crucial that we include an INLINE pragma on >>>, which may be
+surprising. After all, its unfolding is tiny, so GHC will be extremely
+keen to inline it even without the pragma. Indeed, it is actually
+/too/ keen: unintuitively, the pragma is needed to rein in inlining,
+not to encourage it.
+
+How is that possible? The difference lies entirely in whether GHC will
+inline unsaturated calls. With no pragma at all, we get the following
+unfolding guidance:
+    ALWAYS_IF(arity=3,unsat_ok=True,boring_ok=True)
+But with the pragma, we restrict inlining to saturated calls:
+    ALWAYS_IF(arity=3,unsat_ok=False,boring_ok=True)
+Why does this matter? Because the programmer may have put an INLINE
+pragma on (.):
+
+    instance Functor f => Category (Blah f) where
+      id = ...
+      Blah f . Blah g = buildBlah (\x -> ...)
+      {-# INLINE (.) #-}
+
+The intent here is to inline (.) at all saturated call sites. Perhaps
+there is a RULE on buildBlah the programmer wants to fire, or maybe
+they just expect the inlining to expose further simplifications.
+Either way, code that uses >>> should not defeat this inlining, but if
+we inline unsaturated calls, it might! Consider:
+
+    let comp = (>>>) ($fCategoryBlah $dFunctor) in f `comp` (g `comp` h)
+
+While simplifying this expression, we’ll start with the RHS of comp.
+Without the INLINE pragma on >>>, we’ll inline it immediately, even
+though it isn’t saturated:
+
+    let comp = \f g -> $fCategoryBlah_$c. $dFunctor g f
+     in f `comp` (g `comp` h)
+
+Now `$fCategoryBlah_$c. $dFunctor g f` /is/ a fully-saturated call, so
+it will get inlined immediately, too:
+
+    let comp = \(Blah g) (Blah f) -> buildBlah (\x -> ...)
+     in f `comp` (g `comp` h)
+
+All okay so far. But if the RHS of (.) is large, comp won’t be inlined
+at its use sites, and any RULEs on `buildBlah` will never fire. Bad!
+
+What happens differently with the INLINE pragma on >>>? Well, we won’t
+inline >>> immediately, since it isn’t saturated, which means comp’s
+unfolding will be tiny. GHC will inline it at both use sites:
+
+    (>>>) ($fCategoryBlah $dFunctor) f
+          ((>>>) ($fCategoryBlah $dFunctor) g h)
+
+And now the calls to >>> are saturated, so they’ll be inlined,
+followed by (.), and any RULEs can fire as desired. Problem solved.
+
+This situation might seem academic --- who would ever write a
+definition like comp? Probably nobody, but GHC generates such
+definitions when desugaring proc notation, which causes real problems
+(see #18013). That could be fixed by changing the proc desugaring, but
+fixing it this way is the Right Thing, it might benefit other programs
+in more subtle ways too, and it’s easier to boot. -}
diff --git a/libraries/base/GHC/Desugar.hs b/libraries/base/GHC/Desugar.hs
index 89278e99b9..bf25c99bb1 100644
--- a/libraries/base/GHC/Desugar.hs
+++ b/libraries/base/GHC/Desugar.hs
@@ -10,13 +10,13 @@
 -- Module      :  GHC.Desugar
 -- Copyright   :  (c) The University of Glasgow, 2007
 -- License     :  see libraries/base/LICENSE
--- 
+--
 -- Maintainer  :  cvs-ghc@haskell.org
 -- Stability   :  internal
 -- Portability :  non-portable (GHC extensions)
 --
 -- Support code for desugaring in GHC
--- 
+--
 -----------------------------------------------------------------------------
 
 module GHC.Desugar ((>>>), AnnotationWrapper(..), toAnnotationWrapper) where
@@ -28,14 +28,14 @@ import Data.Data        (Data)
 -- A version of Control.Category.>>> overloaded on Arrow
 (>>>) :: forall arr. Arrow arr => forall a b c. arr a b -> arr b c -> arr a c
 -- NB: the type of this function is the "shape" that GHC expects
---     in tcInstClassOp.  So don't put all the foralls at the front!  
+--     in tcInstClassOp.  So don't put all the foralls at the front!
 --     Yes, this is a bit grotesque, but heck it works and the whole
 --     arrows stuff needs reworking anyway!
 f >>> g = g . f
+{-# INLINE (>>>) #-} -- see Note [INLINE on >>>] in Control.Category
 
 -- A wrapper data type that lets the typechecker get at the appropriate dictionaries for an annotation
 data AnnotationWrapper = forall a. (Data a) => AnnotationWrapper a
 
 toAnnotationWrapper :: (Data a) => a -> AnnotationWrapper
 toAnnotationWrapper what = AnnotationWrapper what
-