WorkWrap: Handle partial FUN apps in `isRecDataCon` (#21265)

Partial FUN apps like `(->) Bool` aren't detected by `splitFunTy_maybe`.
A silly oversight that is easily fixed by replacing `splitFunTy_maybe` with a
guard in the `splitTyConApp_maybe` case.

But fortunately, Simon nudged me into rewriting the whole `isRecDataCon`
function in a way that makes it much shorter and hence clearer which DataCons
are actually considered as recursive.

Fixes #21265.

4 files changed, 90 insertions, 93 deletions

diff --git a/compiler/GHC/Core/Opt/CprAnal.hs b/compiler/GHC/Core/Opt/CprAnal.hs
index 97cd36d15a..51bc507a20 100644
--- a/compiler/GHC/Core/Opt/CprAnal.hs
+++ b/compiler/GHC/Core/Opt/CprAnal.hs
@@ -989,43 +989,46 @@ What qualifies as a "recursive data constructor" as per
 Note [CPR for recursive data constructors]? That is up to
 'GHC.Core.Opt.WorkWrapW.Utils.isRecDataCon' to decide. It does a DFS search over
 the field types of the DataCon and looks for term-level recursion into the data
-constructor's type constructor. A few perhaps surprising points:
+constructor's type constructor. Assuming infinite fuel (point (4) below), it
+looks inside the following class of types, represented by `ty` (and responds
+`NonRecursiveOrUnsure` in all other cases):
+
+ A. If `ty = forall v. ty'`, then look into `ty'`
+ B. If `ty = Tc tc_args` and `Tc` is an `AlgTyCon`, look into the arg
+    types of its data constructors and check `tc_args` for recursion.
+ C. If `ty = F tc_args`, `F` is a `FamTyCon` and we can reduce `F tc_args` to
+    `rhs`, look into the `rhs` type.
+
+A few perhaps surprising points:
 
   1. It deems any function type as non-recursive, because it's unlikely that
      a recursion through a function type builds up a recursive data structure.
   2. It doesn't look into kinds or coercion types because there's nothing to unbox.
      Same for promoted data constructors.
-  3. We don't care whether a NewTyCon or DataTyCon App is fully saturated or not;
-     we simply look at its definition/DataCons and its field tys. Any recursive arg
-     occs will have been detected before (see the invariant of 'go_tc_app').
-     This is so that we expand the `ST` in `StateT Int (ST s) a`.
+  3. We don't care whether an AlgTyCon app `T tc_args` is fully saturated or not;
+     we simply look at its definition/DataCons and its field tys and look for
+     recursive occs in the `tc_args` we are given. This is so that we expand
+     the `ST` in `StateT Int (ST s) a`.
   4. We don't recurse deeper than 3 (at the moment of this writing) TyCons and
-     assume the DataCon is non-recursive after that. One reason is guaranteed
-     constant-time efficiency; the other is that it's fair to say that a recursion
-     over 3 or more TyCons doesn't really count as a list-like data structure
-     anymore and a bit of unboxing doesn't hurt much.
-  5. It checks TyConApps like `T <huge> <type>` by eagerly checking the
-     potentially huge argument types *before* it tries to expand the
-     DataCons/NewTyCon/TyFams/etc. so that it doesn't need to re-check those
-     argument types after having been substituted into every occurrence of
-     the the respective TyCon parameter binders. It's like call-by-value vs.
-     call-by-name: Eager checking of argument types means we only need to check
-     them exactly once.
-     There's one exception to that rule, namely when we are able to reduce a
-     TyFam by considering argument types. Then we pay the price of potentially
-     checking the same type arg twice (or more, if the TyFam is recursive).
-     It should hardly matter.
+     assume the DataCon is non-recursive after that. One reason for this "fuel"
+     approach is guaranteed constant-time efficiency; the other is that it's
+     fair to say that a recursion over 3 or more TyCons doesn't really count as
+     a list-like data structure anymore and a bit of unboxing doesn't hurt much.
+  5. It checks AlgTyCon apps like `T tc_args` by eagerly checking the `tc_args`
+     *before* it looks into the expanded DataCons/NewTyCon, so that it
+     terminates before doing a deep nest of expansions only to discover that the
+     first level already contained a recursion.
   6. As a result of keeping the implementation simple, it says "recursive"
      for `data T = MkT [T]`, even though we could argue that the inner recursion
      (through the `[]` TyCon) by way of which `T` is recursive will already be
      "broken" and thus never unboxed. Consequently, it might be OK to CPR a
      function returning `T`. Lacking arguments for or against the current simple
      behavior, we stick to it.
-  7. When the search hits an abstract TyCon (one without visible DataCons, e.g.,
-     from an .hs-boot file), it returns 'Nothing' for "inconclusive", the same
-     as when we run out of fuel. If there is ever a recursion through an
-     abstract TyCon, then it's not part of the same function we are looking at,
-     so we can treat it as if it wasn't recursive.
+  7. When the search hits an abstract TyCon (algebraic, but without visible
+     DataCons, e.g., from an .hs-boot file), it returns 'NonRecursiveOrUnsure',
+     the same as when we run out of fuel. If there is ever a recursion through
+     an abstract TyCon, then it's not part of the same function we are looking
+     at in CPR, so we can treat it as if it wasn't recursive.
      We handle stuck type and data families much the same.
 
 Here are a few examples of data constructors or data types with a single data
@@ -1049,7 +1052,7 @@ con and the answers of our function:
       E Int = Char
       E (a,b) = (E a, E b)
       E Char = Blub
-    data Blah = Blah (E (Int, (Int, Int)))     NonRec  (see point (5))
+    data Blah = Blah (E (Int, (Int, Int)))     NonRec
     data Blub = Blub (E (Char, Int))           Rec
     data Blub2 = Blub2 (E (Bool, Int))     }   Unsure, because stuck (see point (7))
 
diff --git a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
index 471a3a3569..298bb3202a 100644
--- a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
+++ b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
@@ -35,6 +35,7 @@ import GHC.Core.Coercion
 import GHC.Core.Reduction
 import GHC.Core.FamInstEnv
 import GHC.Core.TyCon
+import GHC.Core.TyCon.Set
 import GHC.Core.TyCon.RecWalk
 import GHC.Core.SimpleOpt( SimpleOpts )
 
@@ -49,7 +50,6 @@ import GHC.Types.Unique.Supply
 import GHC.Types.Name ( getOccFS )
 
 import GHC.Data.FastString
-import GHC.Data.Maybe
 import GHC.Data.OrdList
 import GHC.Data.List.SetOps
 
@@ -1267,87 +1267,65 @@ combineIRDCRs = foldl' combineIRDCR NonRecursiveOrUnsure
 -- See Note [Detecting recursive data constructors] for which recursive DataCons
 -- we want to flag.
 isRecDataCon :: FamInstEnvs -> IntWithInf -> DataCon -> IsRecDataConResult
-isRecDataCon fam_envs fuel dc
-  | isTupleDataCon dc || isUnboxedSumDataCon dc
+isRecDataCon fam_envs fuel orig_dc
+  | isTupleDataCon orig_dc || isUnboxedSumDataCon orig_dc
   = NonRecursiveOrUnsure
   | otherwise
   = -- pprTraceWith "isRecDataCon" (\answer -> ppr dc <+> dcolon <+> ppr (dataConRepType dc) $$ ppr fuel $$ ppr answer) $
-    go_dc fuel (setRecTcMaxBound 1 initRecTc) dc
+    go_dc fuel emptyTyConSet orig_dc
   where
-    _pp_dc_ty = ppr dc
-    (<||>) = combineIRDCR
-
-    go_dc :: IntWithInf -> RecTcChecker -> DataCon -> IsRecDataConResult
-    go_dc fuel rec_tc dc =
-      combineIRDCRs [ go_arg_ty fuel rec_tc arg_ty
+    go_dc :: IntWithInf -> TyConSet -> DataCon -> IsRecDataConResult
+    go_dc fuel visited_tcs dc =
+      combineIRDCRs [ go_arg_ty fuel visited_tcs arg_ty
                     | arg_ty <- map scaledThing (dataConRepArgTys dc) ]
 
-    go_arg_ty :: IntWithInf -> RecTcChecker -> Type -> IsRecDataConResult
-    go_arg_ty fuel rec_tc ty
+    go_arg_ty :: IntWithInf -> TyConSet -> Type -> IsRecDataConResult
+    go_arg_ty fuel visited_tcs ty
       --- | pprTrace "arg_ty" (ppr ty) False = undefined
 
-      | Just (_, _arg_ty, _res_ty) <- splitFunTy_maybe ty
-      -- = go_arg_ty fuel rec_tc _arg_ty <||> go_arg_ty fuel rec_tc _res_ty
-          -- Plausible, but unnecessary for CPR.
-          -- See Note [Detecting recursive data constructors], point (1)
-      = NonRecursiveOrUnsure
-
       | Just (_tcv, ty') <- splitForAllTyCoVar_maybe ty
-      = go_arg_ty fuel rec_tc ty'
-          -- See Note [Detecting recursive data constructors], point (2)
+      = go_arg_ty fuel visited_tcs ty'
+          -- See Note [Detecting recursive data constructors], point (A)
 
       | Just (tc, tc_args) <- splitTyConApp_maybe ty
-      = combineIRDCRs (map (go_arg_ty fuel rec_tc) tc_args)
-        <||> go_tc_app fuel rec_tc tc tc_args
+      = go_tc_app fuel visited_tcs tc tc_args
 
       | otherwise
       = NonRecursiveOrUnsure
 
-    -- | PRECONDITION: tc_args has no recursive occs
-    -- See Note [Detecting recursive data constructors], point (5)
-    go_tc_app :: IntWithInf -> RecTcChecker -> TyCon -> [Type] -> IsRecDataConResult
-    go_tc_app fuel rec_tc tc tc_args
-      --- | pprTrace "tc_app" (vcat [ppr tc, ppr tc_args]) False = undefined
-
-      | isPrimTyCon tc
-      = NonRecursiveOrUnsure
-
-      | not $ tcIsRuntimeTypeKind $ tyConResKind tc
-      = NonRecursiveOrUnsure
-
-      | isAbstractTyCon tc   -- When tc has no DataCons, from an hs-boot file
-      = NonRecursiveOrUnsure -- See Note [Detecting recursive data constructors], point (7)
-
-      | isFamilyTyCon tc
-      -- This is the only place where we look at tc_args
-      -- See Note [Detecting recursive data constructors], point (5)
-      = case topReduceTyFamApp_maybe fam_envs tc tc_args of
-          Just (HetReduction (Reduction _ rhs) _) ->
-            go_arg_ty fuel rec_tc rhs
-          Nothing ->
-            NonRecursiveOrUnsure -- NB: We simply give up here. Better return
-                                 -- Unsure, as for abstract TyCons, point (7)
-
-      | tc == dataConTyCon dc
-      = DefinitelyRecursive -- loop found!
-
-      | otherwise
-      = assertPpr (isAlgTyCon tc) (ppr tc <+> ppr dc) $
-        case checkRecTc rec_tc tc of
-          Nothing -> NonRecursiveOrUnsure
-            -- we expanded this TyCon once already, no need to test it multiple times
-
-          Just rec_tc'
-            | Just (_tvs, rhs, _co) <- unwrapNewTyCon_maybe tc
-                -- See Note [Detecting recursive data constructors], points (2) and (3)
-            -> go_arg_ty fuel rec_tc' rhs
-
-            | fuel < 0
-            -> NonRecursiveOrUnsure -- that's why we track fuel!
-
-            | let dcs = expectJust "isRecDataCon:go_tc_app" $ tyConDataCons_maybe tc
-            -> combineIRDCRs (map (\dc -> go_dc (subWithInf fuel 1) rec_tc' dc) dcs)
-                -- See Note [Detecting recursive data constructors], point (4)
+    go_tc_app :: IntWithInf -> TyConSet -> TyCon -> [Type] -> IsRecDataConResult
+    go_tc_app fuel visited_tcs tc tc_args =
+      case tyConDataCons_maybe tc of
+        --- | pprTrace "tc_app" (vcat [ppr tc, ppr tc_args]) False -> undefined
+
+        _ | Just (HetReduction (Reduction _ rhs) _) <- topReduceTyFamApp_maybe fam_envs tc tc_args
+          -- This is the only place where we look at tc_args, which might have
+          -- See Note [Detecting recursive data constructors], point (C) and (5)
+          -> go_arg_ty fuel visited_tcs rhs
+
+        _ | tc == dataConTyCon orig_dc
+          -> DefinitelyRecursive -- loop found!
+
+        Just dcs
+          | DefinitelyRecursive <- combineIRDCRs [ go_arg_ty fuel visited_tcs' ty | ty <- tc_args ]
+              -- Check tc_args, See Note [Detecting recursive data constructors], point (5)
+              -- The new visited_tcs', so that we don't recursively check tc,
+              -- promising that we will check it below.
+              -- Do the tc_args check *before* the dcs check below, otherwise
+              -- we might miss an obvious rec occ in tc_args when we run out of
+              -- fuel and respond NonRecursiveOrUnsure instead
+          -> DefinitelyRecursive
+
+          | fuel >= 0
+              -- See Note [Detecting recursive data constructors], point (4)
+          , not (tc `elemTyConSet` visited_tcs)
+              -- only need to check tc if we haven't visited it already. NB: original visited_tcs
+          -> combineIRDCRs [ go_dc (subWithInf fuel 1) visited_tcs' dc | dc <- dcs ]
+              -- Finally: check ds
+
+        _ -> NonRecursiveOrUnsure
+        where
+          visited_tcs' = extendTyConSet visited_tcs tc
 
 {-
 ************************************************************************
diff --git a/testsuite/tests/stranal/should_compile/T21265.hs b/testsuite/tests/stranal/should_compile/T21265.hs
new file mode 100644
index 0000000000..d1a7c5a0b8
--- /dev/null
+++ b/testsuite/tests/stranal/should_compile/T21265.hs
@@ -0,0 +1,15 @@
+{-# LANGUAGE RankNTypes #-}
+
+module T21265 (extractorProduct') where
+
+class GSerialiseProduct f where
+  dummy :: f x -> ()
+  productExtractor :: TransFusion [] ((->) Bool) (f Int)
+
+extractorProduct' :: GSerialiseProduct f => Maybe (f Int)
+extractorProduct' = unTransFusion productExtractor go
+
+go :: f x -> Maybe (g x)
+go _ = Nothing
+
+newtype TransFusion f g a = TransFusion { unTransFusion :: forall h. Applicative h => (forall x. f x -> h (g x)) -> h a }
diff --git a/testsuite/tests/stranal/should_compile/all.T b/testsuite/tests/stranal/should_compile/all.T
index 2698a3a851..47d2130346 100644
--- a/testsuite/tests/stranal/should_compile/all.T
+++ b/testsuite/tests/stranal/should_compile/all.T
@@ -83,3 +83,4 @@ test('T20817', [ grep_errmsg(r'Str') ], compile, ['-dsuppress-uniques -ddump-str
 test('T21150', [ grep_errmsg(r'( t\d? :: Int)') ], compile, ['-dsuppress-uniques -ddump-exitify'])
 # T21128: Check that y is not reboxed in $wtheresCrud
 test('T21128', [ grep_errmsg(r'let { y = I\#') ], multimod_compile, ['T21128', '-v0 -dsuppress-uniques -dsuppress-all -ddump-simpl'])
+test('T21265', normal, compile, [''])