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.

6 files changed, 355 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
-
diff --git a/testsuite/tests/simplCore/should_compile/T18013.hs b/testsuite/tests/simplCore/should_compile/T18013.hs
new file mode 100644
index 0000000000..08b0a057d2
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T18013.hs
@@ -0,0 +1,18 @@
+{-# OPTIONS_GHC -ddump-rule-firings -ddump-simpl
+                -dsuppress-coercions -dsuppress-uniques #-}
+{-# LANGUAGE Arrows #-}
+
+module T18013 where
+
+import Control.Arrow
+import T18013a
+
+-- We want to ensure this generates good code. Uses of (.) should be
+-- specialized and inlined, and the rules defined on mkRule should fire.
+
+mapMaybeRule :: Rule IO a b -> Rule IO (Maybe a) (Maybe b)
+mapMaybeRule f = proc v -> case v of
+  Just x -> do
+    y <- f -< x
+    returnA -< Just y
+  Nothing -> returnA -< Nothing
diff --git a/testsuite/tests/simplCore/should_compile/T18013.stderr b/testsuite/tests/simplCore/should_compile/T18013.stderr
new file mode 100644
index 0000000000..677c08e7d9
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T18013.stderr
@@ -0,0 +1,210 @@
+Rule fired: Class op arr (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op arr (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op first (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Arrow (BUILTIN)
+Rule fired: Class op arr (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op ||| (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op . (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: mkRule @(_, ()) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @((), _) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @((), _) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @(_, ()) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @((), _) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @((), _) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @(_, ()) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @((), _) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @(_, ()) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @(_, ()) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: mkRule @((), _) (T18013a)
+Rule fired: Class op fmap (BUILTIN)
+
+==================== Tidy Core ====================
+Result size of Tidy Core
+  = {terms: 52, types: 106, coercions: 15, joins: 0/1}
+
+-- RHS size: {terms: 37, types: 87, coercions: 15, joins: 0/1}
+mapMaybeRule
+  :: forall a b. Rule IO a b -> Rule IO (Maybe a) (Maybe b)
+[GblId,
+ Arity=1,
+ Str=<S,1*U>,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [20] 150 30}]
+mapMaybeRule
+  = \ (@a) (@b) (f :: Rule IO a b) ->
+      case f of { Rule @s t0 g ->
+      let {
+        lvl :: Result s (Maybe b)
+        [LclId, Unf=OtherCon []]
+        lvl = T18013a.Result @s @(Maybe b) t0 (GHC.Maybe.Nothing @b) } in
+      T18013a.Rule
+        @IO
+        @(Maybe a)
+        @(Maybe b)
+        @s
+        t0
+        ((\ (s2 :: s)
+            (a1 :: Maybe a)
+            (s1 :: GHC.Prim.State# GHC.Prim.RealWorld) ->
+            case a1 of {
+              Nothing -> (# s1, lvl #);
+              Just x ->
+                case ((g s2 x) `cast` <Co:4>) s1 of { (# ipv, ipv1 #) ->
+                case ipv1 of { Result t2 c1 ->
+                (# ipv, T18013a.Result @s @(Maybe b) t2 (GHC.Maybe.Just @b c1) #)
+                }
+                }
+            })
+         `cast` <Co:11>)
+      }
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T18013.$trModule4 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 20 0}]
+T18013.$trModule4 = "main"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T18013.$trModule3 :: GHC.Types.TrName
+[GblId,
+ Cpr=m1,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 20}]
+T18013.$trModule3 = GHC.Types.TrNameS T18013.$trModule4
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T18013.$trModule2 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 30 0}]
+T18013.$trModule2 = "T18013"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T18013.$trModule1 :: GHC.Types.TrName
+[GblId,
+ Cpr=m1,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 20}]
+T18013.$trModule1 = GHC.Types.TrNameS T18013.$trModule2
+
+-- RHS size: {terms: 3, types: 0, coercions: 0, joins: 0/0}
+T18013.$trModule :: GHC.Types.Module
+[GblId,
+ Cpr=m1,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 30}]
+T18013.$trModule
+  = GHC.Types.Module T18013.$trModule3 T18013.$trModule1
+
+
+
diff --git a/testsuite/tests/simplCore/should_compile/T18013a.hs b/testsuite/tests/simplCore/should_compile/T18013a.hs
new file mode 100644
index 0000000000..f4c5eb46da
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T18013a.hs
@@ -0,0 +1,58 @@
+{-# LANGUAGE BlockArguments, GADTs, LambdaCase #-}
+
+module T18013a where
+
+import Prelude hiding ((.), id)
+
+import Control.Category
+import Control.Arrow
+import Data.Functor
+
+data Result s a = Result !s a
+
+data Rule m a b where
+  Rule :: !s -> !(s -> a -> m (Result s b)) -> Rule m a b
+
+mkRule :: Functor m => s -> (s -> a -> m (Result s b)) -> Rule m a b
+mkRule = Rule
+{-# INLINE CONLIKE [1] mkRule #-}
+{-# RULES
+"mkRule @((), _)" forall s f. mkRule ((), s) f =
+  Rule s (\s1 a -> f ((), s1) a <&> \(Result ((), s2) b) -> Result s2 b)
+"mkRule @(_, ())" forall s f. mkRule (s, ()) f =
+  Rule s (\s1 a -> f (s1, ()) a <&> \(Result (s2, ()) b) -> Result s2 b)
+#-}
+
+instance Monad m => Category (Rule m) where
+  id = arr id
+  {-# INLINE id #-}
+  Rule t0 g . Rule s0 f = mkRule (s0, t0) \(s1, t1) a -> do
+    Result s2 b <- f s1 a
+    Result t2 c <- g t1 b
+    pure $! Result (s2, t2) c
+  {-# INLINE (.) #-}
+
+instance Monad m => Arrow (Rule m) where
+  arr f = Rule () \_ a -> pure $! Result () (f a)
+  {-# INLINE arr #-}
+  first (Rule s0 f) = Rule s0 \s1 (a, c) -> do
+    Result s2 b <- f s1 a
+    pure $! Result s2 (b, c)
+  {-# INLINE first #-}
+
+instance Monad m => ArrowChoice (Rule m) where
+  left (Rule s0 f) = Rule s0 \s1 -> \case
+    Left a -> do
+      Result s2 b <- f s1 a
+      pure $! Result s2 (Left b)
+    Right a ->
+      pure $! Result s0 (Right a)
+  {-# INLINE left #-}
+  Rule s0 f ||| Rule t0 g = mkRule (s0, t0) \(s1, t1) -> \case
+    Left a -> do
+      Result s2 b <- f s1 a
+      pure $! Result (s2, t0) b
+    Right a -> do
+      Result t2 b <- g t1 a
+      pure $! Result (s0, t2) b
+  {-# INLINE (|||) #-}
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index b3262b8d19..396111fb4b 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -326,3 +326,4 @@ test('T17966',
      makefile_test, ['T17966'])
 # NB: T17810: -fspecialise-aggressively
 test('T17810', normal, multimod_compile, ['T17810', '-fspecialise-aggressively -dcore-lint -O -v0'])
+test('T18013', normal, multimod_compile, ['T18013', '-v0 -O'])