Better loop detection in findTypeShape

Andreas pointed out, in !3466, that my fix for #18304 was not
quite right.  This patch fixes it properly, by having just one
RecTcChecker rather than (implicitly) two nested ones, in
findTypeShape.

3 files changed, 41 insertions, 34 deletions

diff --git a/compiler/GHC/Core/FamInstEnv.hs b/compiler/GHC/Core/FamInstEnv.hs
index a35f49b78f..8408bc5406 100644
--- a/compiler/GHC/Core/FamInstEnv.hs
+++ b/compiler/GHC/Core/FamInstEnv.hs
@@ -32,8 +32,8 @@ module GHC.Core.FamInstEnv (
 
         -- Normalisation
         topNormaliseType, topNormaliseType_maybe,
-        normaliseType, normaliseTcApp, normaliseTcArgs,
-        reduceTyFamApp_maybe,
+        normaliseType, normaliseTcApp,
+        topReduceTyFamApp_maybe, reduceTyFamApp_maybe,
 
         -- Flattening
         flattenTys
@@ -1100,7 +1100,7 @@ reduceTyFamApp_maybe :: FamInstEnvs
 --     the role we seek is representational
 -- It does *not* normalise the type arguments first, so this may not
 --     go as far as you want. If you want normalised type arguments,
---     use normaliseTcArgs first.
+--     use topReduceTyFamApp_maybe
 --
 -- The TyCon can be oversaturated.
 -- Works on both open and closed families
@@ -1308,10 +1308,9 @@ topNormaliseType_maybe env ty
       -- to the normalised type's kind
     tyFamStepper :: NormaliseStepper (Coercion, MCoercionN)
     tyFamStepper rec_nts tc tys  -- Try to step a type/data family
-      = let (args_co, ntys, res_co) = normaliseTcArgs env Representational tc tys in
-        case reduceTyFamApp_maybe env Representational tc ntys of
-          Just (co, rhs) -> NS_Step rec_nts rhs (args_co `mkTransCo` co, MCo res_co)
-          _              -> NS_Done
+      = case topReduceTyFamApp_maybe env tc tys of
+          Just (co, rhs, res_co) -> NS_Step rec_nts rhs (co, MCo res_co)
+          _                      -> NS_Done
 
 ---------------
 normaliseTcApp :: FamInstEnvs -> Role -> TyCon -> [Type] -> (Coercion, Type)
@@ -1366,18 +1365,23 @@ normalise_tc_app tc tys
         final_co     = mkCoherenceRightCo r nty (mkSymCo kind_co) orig_to_nty
 
 ---------------
--- | Normalise arguments to a tycon
-normaliseTcArgs :: FamInstEnvs          -- ^ env't with family instances
-                -> Role                 -- ^ desired role of output coercion
-                -> TyCon                -- ^ tc
-                -> [Type]               -- ^ tys
-                -> (Coercion, [Type], CoercionN)
-                                        -- ^ co :: tc tys ~ tc new_tys
-                                        -- NB: co might not be homogeneous
-                                        -- last coercion :: kind(tc tys) ~ kind(tc new_tys)
-normaliseTcArgs env role tc tys
-  = initNormM env role (tyCoVarsOfTypes tys) $
-    normalise_tc_args tc tys
+-- | Try to simplify a type-family application, by *one* step
+-- If topReduceTyFamApp_maybe env r F tys = Just (co, rhs, res_co)
+-- then    co     :: F tys ~R# rhs
+--         res_co :: typeKind(F tys) ~ typeKind(rhs)
+-- Type families and data families; always Representational role
+topReduceTyFamApp_maybe :: FamInstEnvs -> TyCon -> [Type]
+                        -> Maybe (Coercion, Type, Coercion)
+topReduceTyFamApp_maybe envs fam_tc arg_tys
+  | isFamilyTyCon fam_tc   -- type families and data families
+  , Just (co, rhs) <- reduceTyFamApp_maybe envs role fam_tc ntys
+  = Just (args_co `mkTransCo` co, rhs, res_co)
+  | otherwise
+  = Nothing
+  where
+    role = Representational
+    (args_co, ntys, res_co) = initNormM envs role (tyCoVarsOfTypes arg_tys) $
+                              normalise_tc_args fam_tc arg_tys
 
 normalise_tc_args :: TyCon -> [Type]             -- tc tys
                   -> NormM (Coercion, [Type], CoercionN)
diff --git a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
index 2357c4e3e3..1885de98f2 100644
--- a/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
+++ b/compiler/GHC/Core/Opt/WorkWrap/Utils.hs
@@ -1025,7 +1025,19 @@ findTypeShape fam_envs ty
        = TsFun (go rec_tc res)
 
        | Just (tc, tc_args)  <- splitTyConApp_maybe ty
-       , Just con <- isDataProductTyCon_maybe tc
+       = go_tc rec_tc tc tc_args
+
+       | Just (_, ty') <- splitForAllTy_maybe ty
+       = go rec_tc ty'
+
+       | otherwise
+       = TsUnk
+
+    go_tc rec_tc tc tc_args
+       | Just (_, rhs, _) <- topReduceTyFamApp_maybe fam_envs tc tc_args
+       = go rec_tc rhs
+
+       | Just con <- isDataProductTyCon_maybe tc
        , Just rec_tc <- if isTupleTyCon tc
                         then Just rec_tc
                         else checkRecTc rec_tc tc
@@ -1033,10 +1045,8 @@ findTypeShape fam_envs ty
          -- Maybe we should do so in checkRecTc.
        = TsProd (map (go rec_tc . scaledThing) (dataConInstArgTys con tc_args))
 
-       | Just (_, ty') <- splitForAllTy_maybe ty
-       = go rec_tc ty'
-
-       | Just (_, ty') <- topNormaliseType_maybe fam_envs ty
+       | Just (ty', _) <- instNewTyCon_maybe tc tc_args
+       , Just rec_tc <- checkRecTc rec_tc tc
        = go rec_tc ty'
 
        | otherwise
diff --git a/compiler/GHC/HsToCore/PmCheck/Oracle.hs b/compiler/GHC/HsToCore/PmCheck/Oracle.hs
index db56ecc98c..80a844ad19 100644
--- a/compiler/GHC/HsToCore/PmCheck/Oracle.hs
+++ b/compiler/GHC/HsToCore/PmCheck/Oracle.hs
@@ -372,16 +372,9 @@ pmTopNormaliseType (TySt inert) typ
 
     tyFamStepper :: FamInstEnvs -> NormaliseStepper ([Type] -> [Type], a -> a)
     tyFamStepper env rec_nts tc tys  -- Try to step a type/data family
-      = let (_args_co, ntys, _res_co) = normaliseTcArgs env Representational tc tys in
-          -- NB: It's OK to use normaliseTcArgs here instead of
-          -- normalise_tc_args (which takes the LiftingContext described
-          -- in Note [Normalising types]) because the reduceTyFamApp below
-          -- works only at top level. We'll never recur in this function
-          -- after reducing the kind of a bound tyvar.
-
-        case reduceTyFamApp_maybe env Representational tc ntys of
-          Just (_co, rhs) -> NS_Step rec_nts rhs ((rhs:), id)
-          _               -> NS_Done
+      = case topReduceTyFamApp_maybe env tc tys of
+          Just (_, rhs, _) -> NS_Step rec_nts rhs ((rhs:), id)
+          _                -> NS_Done
 
 -- | Returns 'True' if the argument 'Type' is a fully saturated application of
 -- a closed type constructor.