Produce all DerivInfo in tcTyAndClassDecls

Before this refactoring:

* DerivInfo for data family instances was returned from tcTyAndClassDecls
* DerivInfo for data declarations was generated with mkDerivInfos and added at a
  later stage of the pipeline in tcInstDeclsDeriv

After this refactoring:

* DerivInfo for both data family instances and data declarations is returned from
  tcTyAndClassDecls in a single list.

This uniform treatment results in a more convenient arrangement to fix #16731.

4 files changed, 62 insertions, 48 deletions

diff --git a/compiler/typecheck/TcDeriv.hs b/compiler/typecheck/TcDeriv.hs
index b7c1478da3..224a6a713a 100644
--- a/compiler/typecheck/TcDeriv.hs
+++ b/compiler/typecheck/TcDeriv.hs
@@ -9,7 +9,7 @@ Handles @deriving@ clauses on @data@ declarations.
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE TypeFamilies #-}
 
-module TcDeriv ( tcDeriving, DerivInfo(..), mkDerivInfos ) where
+module TcDeriv ( tcDeriving, DerivInfo(..) ) where
 
 #include "HsVersions.h"
 
@@ -23,7 +23,7 @@ import FamInst
 import TcDerivInfer
 import TcDerivUtils
 import TcValidity( allDistinctTyVars )
-import TcClassDcl( instDeclCtxt3, tcATDefault, tcMkDeclCtxt )
+import TcClassDcl( instDeclCtxt3, tcATDefault )
 import TcEnv
 import TcGenDeriv                       -- Deriv stuff
 import TcValidity( checkValidInstHead )
@@ -199,19 +199,6 @@ data DerivInfo = DerivInfo { di_rep_tc  :: TyCon
                            , di_ctxt    :: SDoc -- ^ error context
                            }
 
--- | Extract `deriving` clauses of proper data type (skips data families)
-mkDerivInfos :: [LTyClDecl GhcRn] -> TcM [DerivInfo]
-mkDerivInfos decls = concatMapM (mk_deriv . unLoc) decls
-  where
-
-    mk_deriv decl@(DataDecl { tcdLName = L _ data_name
-                            , tcdDataDefn =
-                                HsDataDefn { dd_derivs = L _ clauses } })
-      = do { tycon <- tcLookupTyCon data_name
-           ; return [DerivInfo { di_rep_tc = tycon, di_clauses = clauses
-                               , di_ctxt = tcMkDeclCtxt decl }] }
-    mk_deriv _ = return []
-
 {-
 
 ************************************************************************
diff --git a/compiler/typecheck/TcInstDcls.hs b/compiler/typecheck/TcInstDcls.hs
index c2f7a1100a..716acb6942 100644
--- a/compiler/typecheck/TcInstDcls.hs
+++ b/compiler/typecheck/TcInstDcls.hs
@@ -392,17 +392,14 @@ tcInstDecls1 inst_decls
 --   (DerivDecl) to check and process all derived class instances.
 tcInstDeclsDeriv
   :: [DerivInfo]
-  -> [LTyClDecl GhcRn]
   -> [LDerivDecl GhcRn]
   -> TcM (TcGblEnv, [InstInfo GhcRn], HsValBinds GhcRn)
-tcInstDeclsDeriv datafam_deriv_infos tyclds derivds
+tcInstDeclsDeriv deriv_infos derivds
   = do th_stage <- getStage -- See Note [Deriving inside TH brackets]
        if isBrackStage th_stage
        then do { gbl_env <- getGblEnv
                ; return (gbl_env, bagToList emptyBag, emptyValBindsOut) }
-       else do { data_deriv_infos <- mkDerivInfos tyclds
-               ; let deriv_infos = datafam_deriv_infos ++ data_deriv_infos
-               ; (tcg_env, info_bag, valbinds) <- tcDeriving deriv_infos derivds
+       else do { (tcg_env, info_bag, valbinds) <- tcDeriving deriv_infos derivds
                ; return (tcg_env, bagToList info_bag, valbinds) }
 
 addClsInsts :: [InstInfo GhcRn] -> TcM a -> TcM a
diff --git a/compiler/typecheck/TcRnDriver.hs b/compiler/typecheck/TcRnDriver.hs
index 55c229766f..96240e6092 100644
--- a/compiler/typecheck/TcRnDriver.hs
+++ b/compiler/typecheck/TcRnDriver.hs
@@ -1682,7 +1682,7 @@ tcTyClsInstDecls :: [TyClGroup GhcRn]
 tcTyClsInstDecls tycl_decls deriv_decls binds
  = tcAddDataFamConPlaceholders (tycl_decls >>= group_instds) $
    tcAddPatSynPlaceholders (getPatSynBinds binds) $
-   do { (tcg_env, inst_info, datafam_deriv_info)
+   do { (tcg_env, inst_info, deriv_info)
           <- tcTyAndClassDecls tycl_decls ;
       ; setGblEnv tcg_env $ do {
           -- With the @TyClDecl@s and @InstDecl@s checked we're ready to
@@ -1692,9 +1692,8 @@ tcTyClsInstDecls tycl_decls deriv_decls binds
           -- Careful to quit now in case there were instance errors, so that
           -- the deriving errors don't pile up as well.
           ; failIfErrsM
-          ; let tyclds = tycl_decls >>= group_tyclds
           ; (tcg_env', inst_info', val_binds)
-              <- tcInstDeclsDeriv datafam_deriv_info tyclds deriv_decls
+              <- tcInstDeclsDeriv deriv_info deriv_decls
           ; setGblEnv tcg_env' $ do {
                 failIfErrsM
               ; pure (tcg_env', inst_info' ++ inst_info, val_binds)
diff --git a/compiler/typecheck/TcTyClsDecls.hs b/compiler/typecheck/TcTyClsDecls.hs
index 06a730519b..67fc558af1 100644
--- a/compiler/typecheck/TcTyClsDecls.hs
+++ b/compiler/typecheck/TcTyClsDecls.hs
@@ -36,7 +36,7 @@ import TcHsSyn
 import TcTyDecls
 import TcClassDcl
 import {-# SOURCE #-} TcInstDcls( tcInstDecls1 )
-import TcDeriv (DerivInfo)
+import TcDeriv (DerivInfo(..))
 import TcUnify ( unifyKind )
 import TcHsType
 import ClsInst( AssocInstInfo(..) )
@@ -124,7 +124,7 @@ tcTyAndClassDecls :: [TyClGroup GhcRn]      -- Mutually-recursive groups in
                                             -- classes
                                             -- and their implicit Ids,DataCons
                          , [InstInfo GhcRn] -- Source-code instance decls info
-                         , [DerivInfo]      -- data family deriving info
+                         , [DerivInfo]      -- Deriving info
                          )
 -- Fails if there are any errors
 tcTyAndClassDecls tyclds_s
@@ -160,7 +160,7 @@ tcTyClGroup (TyClGroup { group_tyclds = tyclds
            -- Step 1: Typecheck the type/class declarations
        ; traceTc "---- tcTyClGroup ---- {" empty
        ; traceTc "Decls for" (ppr (map (tcdName . unLoc) tyclds))
-       ; tyclss <- tcTyClDecls tyclds role_annots
+       ; (tyclss, data_deriv_info) <- tcTyClDecls tyclds role_annots
 
            -- Step 1.5: Make sure we don't have any type synonym cycles
        ; traceTc "Starting synonym cycle check" (ppr tyclss)
@@ -186,12 +186,20 @@ tcTyClGroup (TyClGroup { group_tyclds = tyclds
        ; gbl_env <- addTyConsToGblEnv tyclss
 
            -- Step 4: check instance declarations
-       ; setGblEnv gbl_env $
-         tcInstDecls1 instds }
+       ; (gbl_env', inst_info, datafam_deriv_info) <-
+         setGblEnv gbl_env $
+         tcInstDecls1 instds
+
+       ; let deriv_info = datafam_deriv_info ++ data_deriv_info
+       ; return (gbl_env', inst_info, deriv_info) }
+
 
 tcTyClGroup (XTyClGroup _) = panic "tcTyClGroup"
 
-tcTyClDecls :: [LTyClDecl GhcRn] -> RoleAnnotEnv -> TcM [TyCon]
+tcTyClDecls
+  :: [LTyClDecl GhcRn]
+  -> RoleAnnotEnv
+  -> TcM ([TyCon], [DerivInfo])
 tcTyClDecls tyclds role_annots
   = tcExtendKindEnv promotion_err_env $   --- See Note [Type environment evolution]
     do {    -- Step 1: kind-check this group and returns the final
@@ -206,7 +214,7 @@ tcTyClDecls tyclds role_annots
             -- NB: We have to be careful here to NOT eagerly unfold
             -- type synonyms, as we have not tested for type synonym
             -- loops yet and could fall into a black hole.
-       ; fixM $ \ ~rec_tyclss -> do
+       ; fixM $ \ ~(rec_tyclss, _) -> do
            { tcg_env <- getGblEnv
            ; let roles = inferRoles (tcg_src tcg_env) role_annots rec_tyclss
 
@@ -214,7 +222,8 @@ tcTyClDecls tyclds role_annots
                  -- NB: if the decls mention any ill-staged data cons
                  -- (see Note [Recursion and promoting data constructors])
                  -- we will have failed already in kcTyClGroup, so no worries here
-           ; tcExtendRecEnv (zipRecTyClss tc_tycons rec_tyclss) $
+           ; (tycons, data_deriv_infos) <-
+             tcExtendRecEnv (zipRecTyClss tc_tycons rec_tyclss) $
 
                  -- Also extend the local type envt with bindings giving
                  -- a TcTyCon for each each knot-tied TyCon or Class
@@ -223,7 +232,8 @@ tcTyClDecls tyclds role_annots
              tcExtendKindEnvWithTyCons tc_tycons $
 
                  -- Kind and type check declarations for this group
-               mapM (tcTyClDecl roles) tyclds
+               mapAndUnzipM (tcTyClDecl roles) tyclds
+           ; return (tycons, concat data_deriv_infos)
            } }
   where
     promotion_err_env = mkPromotionErrorEnv tyclds
@@ -1521,40 +1531,55 @@ unlifted types, resolving #13595.
 
 -}
 
-tcTyClDecl :: RolesInfo -> LTyClDecl GhcRn -> TcM TyCon
+tcTyClDecl :: RolesInfo -> LTyClDecl GhcRn -> TcM (TyCon, [DerivInfo])
 tcTyClDecl roles_info (dL->L loc decl)
   | Just thing <- wiredInNameTyThing_maybe (tcdName decl)
   = case thing of -- See Note [Declarations for wired-in things]
-      ATyCon tc -> return tc
+      ATyCon tc -> return (tc, wiredInDerivInfo tc decl)
       _ -> pprPanic "tcTyClDecl" (ppr thing)
 
   | otherwise
   = setSrcSpan loc $ tcAddDeclCtxt decl $
     do { traceTc "---- tcTyClDecl ---- {" (ppr decl)
-       ; tc <- tcTyClDecl1 Nothing roles_info decl
+       ; (tc, deriv_infos) <- tcTyClDecl1 Nothing roles_info decl
        ; traceTc "---- tcTyClDecl end ---- }" (ppr tc)
-       ; return tc }
+       ; return (tc, deriv_infos) }
+
+noDerivInfos :: a -> (a, [DerivInfo])
+noDerivInfos a = (a, [])
+
+wiredInDerivInfo :: TyCon -> TyClDecl GhcRn -> [DerivInfo]
+wiredInDerivInfo tycon decl
+  | DataDecl { tcdDataDefn = dataDefn } <- decl
+  , HsDataDefn { dd_derivs = derivs } <- dataDefn
+  = [ DerivInfo { di_rep_tc = tycon
+                , di_clauses = unLoc derivs
+                , di_ctxt = tcMkDeclCtxt decl } ]
+wiredInDerivInfo _ _ = []
 
   -- "type family" declarations
-tcTyClDecl1 :: Maybe Class -> RolesInfo -> TyClDecl GhcRn -> TcM TyCon
+tcTyClDecl1 :: Maybe Class -> RolesInfo -> TyClDecl GhcRn -> TcM (TyCon, [DerivInfo])
 tcTyClDecl1 parent _roles_info (FamDecl { tcdFam = fd })
-  = tcFamDecl1 parent fd
+  = fmap noDerivInfos $
+    tcFamDecl1 parent fd
 
   -- "type" synonym declaration
 tcTyClDecl1 _parent roles_info
             (SynDecl { tcdLName = (dL->L _ tc_name)
                      , tcdRhs   = rhs })
   = ASSERT( isNothing _parent )
+    fmap noDerivInfos $
     bindTyClTyVars tc_name $ \ binders res_kind ->
     tcTySynRhs roles_info tc_name binders res_kind rhs
 
   -- "data/newtype" declaration
 tcTyClDecl1 _parent roles_info
-            (DataDecl { tcdLName = (dL->L _ tc_name)
-                      , tcdDataDefn = defn })
+            decl@(DataDecl { tcdLName = (dL->L _ tc_name)
+                           , tcdDataDefn = defn })
   = ASSERT( isNothing _parent )
     bindTyClTyVars tc_name $ \ tycon_binders res_kind ->
-    tcDataDefn roles_info tc_name tycon_binders res_kind defn
+    tcDataDefn (tcMkDeclCtxt decl) roles_info tc_name
+               tycon_binders res_kind defn
 
 tcTyClDecl1 _parent roles_info
             (ClassDecl { tcdLName = (dL->L _ class_name)
@@ -1567,7 +1592,7 @@ tcTyClDecl1 _parent roles_info
   = ASSERT( isNothing _parent )
     do { clas <- tcClassDecl1 roles_info class_name hs_ctxt
                               meths fundeps sigs ats at_defs
-       ; return (classTyCon clas) }
+       ; return (noDerivInfos (classTyCon clas)) }
 
 tcTyClDecl1 _ _ (XTyClDecl _) = panic "tcTyClDecl1"
 
@@ -2009,17 +2034,20 @@ tcTySynRhs roles_info tc_name binders res_kind hs_ty
              tycon = buildSynTyCon tc_name binders res_kind roles rhs_ty
        ; return tycon }
 
-tcDataDefn :: RolesInfo -> Name
+tcDataDefn :: SDoc
+           -> RolesInfo -> Name
            -> [TyConBinder] -> Kind
-           -> HsDataDefn GhcRn -> TcM TyCon
+           -> HsDataDefn GhcRn -> TcM (TyCon, [DerivInfo])
   -- NB: not used for newtype/data instances (whether associated or not)
-tcDataDefn roles_info
+tcDataDefn err_ctxt
+           roles_info
            tc_name tycon_binders res_kind
            (HsDataDefn { dd_ND = new_or_data, dd_cType = cType
                        , dd_ctxt = ctxt
                        , dd_kindSig = mb_ksig  -- Already in tc's kind
                                                -- via getInitialKinds
-                       , dd_cons = cons })
+                       , dd_cons = cons
+                       , dd_derivs = derivs })
  =  do { gadt_syntax <- dataDeclChecks tc_name new_or_data ctxt cons
 
        ; tcg_env <- getGblEnv
@@ -2057,8 +2085,11 @@ tcDataDefn roles_info
                                   stupid_theta tc_rhs
                                   (VanillaAlgTyCon tc_rep_nm)
                                   gadt_syntax) }
+       ; let deriv_info = DerivInfo { di_rep_tc = tycon
+                                    , di_clauses = unLoc derivs
+                                    , di_ctxt = err_ctxt }
        ; traceTc "tcDataDefn" (ppr tc_name $$ ppr tycon_binders $$ ppr extra_bndrs)
-       ; return tycon }
+       ; return (tycon, [deriv_info]) }
   where
     -- Abstract data types in hsig files can have arbitrary kinds,
     -- because they may be implemented by type synonyms
@@ -2084,7 +2115,7 @@ tcDataDefn roles_info
           DataType -> return (mkDataTyConRhs data_cons)
           NewType  -> ASSERT( not (null data_cons) )
                       mkNewTyConRhs tc_name tycon (head data_cons)
-tcDataDefn _ _ _ _ (XHsDataDefn _) = panic "tcDataDefn"
+tcDataDefn _ _ _ _ _ (XHsDataDefn _) = panic "tcDataDefn"
 
 
 -------------------------