Be more careful about retaining KnotVars

It is quite easy to end up accidently retaining a KnotVars, which
contains pointers to a stale TypeEnv because they are placed in the
HscEnv.

One place in particular we have to be careful is when loading a module
into the EPS in `--make` mode, we have to remove the reference to
KnotVars as otherwise the interface loading thunks will forever retain
reference to the KnotVars which are live at the time the interface was
loaded.

These changes do not go as far as to enforce the invariant described in
Note [KnotVar invariants]

* At the end of upsweep, there should be no live KnotVars

but at least improve the situation.

This is left for future work (#20491)

4 files changed, 51 insertions, 5 deletions

diff --git a/compiler/GHC/Driver/Env/KnotVars.hs b/compiler/GHC/Driver/Env/KnotVars.hs
index 73f348835f..5c39381fba 100644
--- a/compiler/GHC/Driver/Env/KnotVars.hs
+++ b/compiler/GHC/Driver/Env/KnotVars.hs
@@ -13,30 +13,41 @@ import GHC.Prelude
 import GHC.Unit.Types ( Module )
 import GHC.Unit.Module.Env
 import Data.Maybe
+import GHC.Utils.Outputable
 
 -- See Note [Why is KnotVars not a ModuleEnv]
+-- See Note [KnotVars invariants]
 data KnotVars a = KnotVars { kv_domain :: [Module] -- Domain of the function , Note [KnotVars: Why store the domain?]
                            -- Invariant: kv_lookup is surjective relative to kv_domain
                            , kv_lookup :: Module -> Maybe a -- Lookup function
                            }
+                | NoKnotVars
                            deriving Functor
 
+instance Outputable (KnotVars a) where
+  ppr NoKnotVars = text "NoKnot"
+  ppr (KnotVars dom _lookup) = text "Knotty:" <+> ppr dom
+
 emptyKnotVars :: KnotVars a
-emptyKnotVars = KnotVars [] (const Nothing)
+emptyKnotVars = NoKnotVars
 
 knotVarsFromModuleEnv :: ModuleEnv a -> KnotVars a
+knotVarsFromModuleEnv me | isEmptyModuleEnv me = NoKnotVars
 knotVarsFromModuleEnv me = KnotVars (moduleEnvKeys me) (lookupModuleEnv me)
 
 knotVarElems :: KnotVars a -> [a]
 knotVarElems (KnotVars keys lookup) = mapMaybe lookup keys
+knotVarElems NoKnotVars = []
 
 lookupKnotVars :: KnotVars a -> Module -> Maybe a
-lookupKnotVars (KnotVars _ lookup) = lookup
+lookupKnotVars (KnotVars _ lookup) x = lookup x
+lookupKnotVars NoKnotVars _ = Nothing
 
 knotVarsWithout :: Module -> KnotVars a -> KnotVars a
 knotVarsWithout this_mod (KnotVars loop_mods lkup) = KnotVars
   (filter (/= this_mod) loop_mods)
   (\that_mod -> if that_mod == this_mod then Nothing else lkup that_mod)
+knotVarsWithout _ NoKnotVars = NoKnotVars
 
 {-
 Note [Why is KnotVars not a ModuleEnv]
@@ -67,5 +78,27 @@ the top-level identifiers from modules in the cycle might not be globalised prop
 This could be refactored so that the lint functions knew about 'KnotVars' and delayed
 this check until deciding whether a variable was local or not.
 
+
+Note [KnotVars invariants]
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+There is a simple invariant which should hold for the KnotVars constructor:
+
+* At the end of upsweep, there should be no live KnotVars
+
+This invariant is difficult to test but easy to check using ghc-debug. The usage of
+NoKnotVars is intended to make this invariant easier to check.
+
+The most common situation where a KnotVars is retained accidently is if a HscEnv
+which contains reference to a KnotVars is used during interface file loading. The
+thunks created during this process will retain a reference to the KnotVars. In theory,
+all these references should be removed by 'typecheckLoop' as that retypechecks all
+interface files in the loop without using KnotVars.
+
+At the time of writing (MP: Oct 21) the invariant doesn't actually hold but also
+doesn't seem to have too much of a negative consequence on compiler residency.
+In theory it could be quite bad as each KnotVars may retain a stale reference to an entire TypeEnv.
+
+See #20491
 -}
 
diff --git a/compiler/GHC/Iface/Load.hs b/compiler/GHC/Iface/Load.hs
index eb92b881e3..2b31074896 100644
--- a/compiler/GHC/Iface/Load.hs
+++ b/compiler/GHC/Iface/Load.hs
@@ -631,9 +631,18 @@ home-package modules however, so it's safe for the HPT to be empty.
 
 dontLeakTheHPT :: IfL a -> IfL a
 dontLeakTheHPT thing_inside = do
+  env <- getTopEnv
   let
+    inOneShot =
+      isOneShot (ghcMode (hsc_dflags env))
+    cleanGblEnv gbl_env
+      | inOneShot = gbl_env
+      | otherwise = gbl_env { if_rec_types = emptyKnotVars }
     cleanTopEnv hsc_env =
+
        let
+         !maybe_type_vars | inOneShot = Just (hsc_type_env_vars env)
+                          | otherwise = Nothing
          -- wrinkle: when we're typechecking in --backpack mode, the
          -- instantiation of a signature might reside in the HPT, so
          -- this case breaks the assumption that EPS interfaces only
@@ -655,11 +664,15 @@ dontLeakTheHPT thing_inside = do
        hsc_env {  hsc_targets      = panic "cleanTopEnv: hsc_targets"
                ,  hsc_mod_graph    = panic "cleanTopEnv: hsc_mod_graph"
                ,  hsc_IC           = panic "cleanTopEnv: hsc_IC"
+               ,  hsc_type_env_vars = case maybe_type_vars of
+                                          Just vars -> vars
+                                          Nothing -> panic "cleanTopEnv: hsc_type_env_vars"
                ,  hsc_unit_env     = unit_env
                }
 
-  updTopEnv cleanTopEnv $ do
+  updTopEnv cleanTopEnv $ updGblEnv cleanGblEnv $ do
   !_ <- getTopEnv        -- force the updTopEnv
+  !_ <- getGblEnv
   thing_inside
 
 
diff --git a/compiler/GHC/Tc/Types.hs b/compiler/GHC/Tc/Types.hs
index d3c4b44211..0ae52c6db3 100644
--- a/compiler/GHC/Tc/Types.hs
+++ b/compiler/GHC/Tc/Types.hs
@@ -317,7 +317,7 @@ data IfGblEnv
         -- We need the module name so we can test when it's appropriate
         -- to look in this env.
         -- See Note [Tying the knot] in GHC.IfaceToCore
-        if_rec_types :: !(KnotVars (IfG TypeEnv))
+        if_rec_types :: (KnotVars (IfG TypeEnv))
                 -- Allows a read effect, so it can be in a mutable
                 -- variable; c.f. handling the external package type env
                 -- Nothing => interactive stuff, no loops possible
diff --git a/compiler/GHC/Tc/Utils/Monad.hs b/compiler/GHC/Tc/Utils/Monad.hs
index 6b876d3121..1c80acdf9b 100644
--- a/compiler/GHC/Tc/Utils/Monad.hs
+++ b/compiler/GHC/Tc/Utils/Monad.hs
@@ -2087,7 +2087,7 @@ initIfaceLoad hsc_env do_this
                         if_doc = text "initIfaceLoad",
                         if_rec_types = emptyKnotVars
                     }
-      initTcRnIf 'i' hsc_env gbl_env () do_this
+      initTcRnIf 'i' (hsc_env { hsc_type_env_vars = emptyKnotVars }) gbl_env () do_this
 
 -- | This is used when we are doing to call 'typecheckModule' on an 'ModIface',
 -- if it's part of a loop with some other modules then we need to use their