Introduce strict wrapper around TypeEnvwip/type-env-force

This reduces residency when compiling Cabal to a baseline of around
600mb live data from 850mb.

12 files changed, 60 insertions, 47 deletions

diff --git a/compiler/GHC/Driver/Env.hs b/compiler/GHC/Driver/Env.hs
index 8d9aa961fb..7c51ebf3e2 100644
--- a/compiler/GHC/Driver/Env.hs
+++ b/compiler/GHC/Driver/Env.hs
@@ -49,7 +49,7 @@ import GHC.Core.InstEnv ( ClsInst )
 import GHC.Types.Annotations ( Annotation, AnnEnv, mkAnnEnv, plusAnnEnv )
 import GHC.Types.CompleteMatch
 import GHC.Types.Name
-import GHC.Types.Name.Env
+import GHC.Types.TypeEnv
 import GHC.Types.TyThing
 
 import GHC.Builtin.Names ( gHC_PRIM )
@@ -259,10 +259,10 @@ lookupType hsc_env name = do
 
        !ty = if isOneShot (ghcMode (hsc_dflags hsc_env))
                -- in one-shot, we don't use the HPT
-               then lookupNameEnv pte name
+               then lookupTypeEnv pte name
                else case lookupHptByModule hpt mod of
-                Just hm -> lookupNameEnv (md_types (hm_details hm)) name
-                Nothing -> lookupNameEnv pte name
+                Just hm -> lookupTypeEnv (md_types (hm_details hm)) name
+                Nothing -> lookupTypeEnv pte name
    pure ty
 
 -- | Find the 'ModIface' for a 'Module', searching in both the loaded home
diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index d3695177d3..19639f04e0 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -193,12 +193,12 @@ import GHC.Types.Unique.Supply
 import GHC.Types.SourceFile
 import GHC.Types.SrcLoc
 import GHC.Types.Name
-import GHC.Types.Name.Env
 import GHC.Types.Name.Cache ( initNameCache )
 import GHC.Types.Name.Reader
 import GHC.Types.Name.Ppr
 import GHC.Types.TyThing
 import GHC.Types.HpcInfo
+import GHC.Types.TypeEnv
 
 import GHC.Utils.Fingerprint ( Fingerprint )
 import GHC.Utils.Panic
@@ -798,7 +798,7 @@ hscIncrementalCompile always_do_basic_recompilation_check m_tc_result
     -- One-shot mode needs a knot-tying mutable variable for interface
     -- files. See GHC.Tc.Utils.TcGblEnv.tcg_type_env_var.
     -- See also Note [hsc_type_env_var hack]
-    type_env_var <- newIORef emptyNameEnv
+    type_env_var <- newIORef emptyTypeEnv
     let mod = ms_mod mod_summary
         hsc_env | isOneShot (ghcMode (hsc_dflags hsc_env''))
                 = hsc_env'' { hsc_type_env_var = Just (mod, type_env_var) }
diff --git a/compiler/GHC/Driver/Make.hs b/compiler/GHC/Driver/Make.hs
index f13d13b198..4b99fb95b5 100644
--- a/compiler/GHC/Driver/Make.hs
+++ b/compiler/GHC/Driver/Make.hs
@@ -94,7 +94,7 @@ import GHC.Types.Unique.FM
 import GHC.Types.Unique.DSet
 import GHC.Types.Unique.Set
 import GHC.Types.Name
-import GHC.Types.Name.Env
+import GHC.Types.TypeEnv
 
 import GHC.Unit
 import GHC.Unit.External
@@ -1438,7 +1438,7 @@ parUpsweep_one mod home_mod_map comp_graph_loops lcl_logger lcl_dflags home_unit
                 -- Re-typecheck the loop
                 -- This is necessary to make sure the knot is tied when
                 -- we close a recursive module loop, see bug #12035.
-                type_env_var <- liftIO $ newIORef emptyNameEnv
+                type_env_var <- liftIO $ newIORef emptyTypeEnv
                 let lcl_hsc_env' = lcl_hsc_env { hsc_type_env_var =
                                     Just (ms_mod lcl_mod, type_env_var) }
                 lcl_hsc_env'' <- case finish_loop of
@@ -1591,7 +1591,7 @@ upsweep mHscMessage old_hpt stable_mods cleanup sccs = do
         liftIO (cleanup hsc_env)
 
         -- Get ready to tie the knot
-        type_env_var <- liftIO $ newIORef emptyNameEnv
+        type_env_var <- liftIO $ newIORef emptyTypeEnv
         let hsc_env1 = hsc_env { hsc_type_env_var =
                                     Just (ms_mod mod, type_env_var) }
         setSession hsc_env1
diff --git a/compiler/GHC/Iface/Load.hs b/compiler/GHC/Iface/Load.hs
index b1a4f4d27c..01875cf522 100644
--- a/compiler/GHC/Iface/Load.hs
+++ b/compiler/GHC/Iface/Load.hs
@@ -191,7 +191,7 @@ importDecl name
                                 text "Use -ddump-if-trace to get an idea of which file caused the error"])
     found_things_msg eps =
         hang (text "Found the following declarations in" <+> ppr (nameModule name) <> colon)
-           2 (vcat (map ppr $ filter is_interesting $ nameEnvElts $ eps_PTE eps))
+           2 (vcat (map ppr $ filter is_interesting $ typeEnvElts $ eps_PTE eps))
       where
         is_interesting thing = nameModule name == nameModule (getName thing)
 
@@ -781,7 +781,7 @@ badSourceImport mod
 -----------------------------------------------------
 
 addDeclsToPTE :: PackageTypeEnv -> [(Name,TyThing)] -> PackageTypeEnv
-addDeclsToPTE pte things = extendNameEnvList pte things
+addDeclsToPTE (TypeEnv pte) things = TypeEnv (extendNameEnvList pte things)
 
 {-
 *********************************************************
diff --git a/compiler/GHC/Iface/UpdateIdInfos.hs b/compiler/GHC/Iface/UpdateIdInfos.hs
index 0c70b5caeb..a730d24e2c 100644
--- a/compiler/GHC/Iface/UpdateIdInfos.hs
+++ b/compiler/GHC/Iface/UpdateIdInfos.hs
@@ -44,8 +44,7 @@ updateModDetailsIdInfos cg_infos mod_details =
               , md_rules = rules
               } = mod_details
 
-    -- type TypeEnv = NameEnv TyThing
-    type_env' = mapNameEnv (updateTyThingIdInfos type_env' cg_infos) type_env
+    type_env' = mapTypeEnv (updateTyThingIdInfos type_env' cg_infos) type_env
     -- NB: Knot-tied! The result, type_env', is passed right back into into
     -- updateTyThingIdInfos, so that that occurrences of any Ids (e.g. in
     -- IdInfos, etc) can be looked up in the tidied env
@@ -115,17 +114,18 @@ updateIdInfo CgInfos{ cgNonCafs = NonCaffySet non_cafs, cgLFInfos = lf_infos } i
     in
       id2
 
+
 --------------------------------------------------------------------------------
 
-updateGlobalIds :: NameEnv TyThing -> CoreExpr -> CoreExpr
+updateGlobalIds :: TypeEnv -> CoreExpr -> CoreExpr
 -- Update occurrences of GlobalIds as directed by 'env'
 -- The 'env' maps a GlobalId to a version with accurate CAF info
 -- (and in due course perhaps other back-end-related info)
 updateGlobalIds env e = go env e
   where
-    go_id :: NameEnv TyThing -> Id -> Id
+    go_id :: TypeEnv -> Id -> Id
     go_id env var =
-      case lookupNameEnv env (varName var) of
+      case lookupTypeEnv env (varName var) of
         Nothing -> var
         Just (AnId id) -> id
         Just other -> pprPanic "UpdateIdInfos.updateGlobalIds" $
@@ -133,7 +133,7 @@ updateGlobalIds env e = go env e
           nest 4 (text "Id:" <+> ppr var $$
                   text "TyThing:" <+> ppr other)
 
-    go :: NameEnv TyThing -> CoreExpr -> CoreExpr
+    go :: TypeEnv -> CoreExpr -> CoreExpr
     go env (Var v) = Var (go_id env v)
     go _ e@Lit{} = e
     go env (App e1 e2) = App (go env e1) (go env e2)
@@ -148,7 +148,7 @@ updateGlobalIds env e = go env e
     go _ e@Type{} = e
     go _ e@Coercion{} = e
 
-    go_binds :: NameEnv TyThing -> CoreBind -> CoreBind
+    go_binds :: TypeEnv -> CoreBind -> CoreBind
     go_binds env (NonRec b e) =
       assertNotInNameEnv env [b] (NonRec b (go env e))
     go_binds env (Rec prs) =
@@ -156,5 +156,5 @@ updateGlobalIds env e = go env e
 
 -- In `updateGlobaLIds` Names of local binders should not shadow Name of
 -- globals. This assertion is to check that.
-assertNotInNameEnv :: NameEnv a -> [Id] -> b -> b
-assertNotInNameEnv env ids x = ASSERT(not (any (\id -> elemNameEnv (idName id) env) ids)) x
+assertNotInNameEnv :: TypeEnv -> [Id] -> b -> b
+assertNotInNameEnv env ids x = ASSERT(not (any (\id -> elemTypeEnv (idName id) env) ids)) x
diff --git a/compiler/GHC/IfaceToCore.hs b/compiler/GHC/IfaceToCore.hs
index 76079ae8ff..d619dbc01d 100644
--- a/compiler/GHC/IfaceToCore.hs
+++ b/compiler/GHC/IfaceToCore.hs
@@ -212,7 +212,7 @@ typecheckIface iface
                          -- we'll infinite loop with hs-boot.  See #10083 for
                          -- an example where this would cause non-termination.
                          text "Type envt:" <+> ppr (map fst names_w_things)])
-        ; return $ ModDetails { md_types     = type_env
+        ; return $ ModDetails { md_types     = TypeEnv type_env
                               , md_insts     = insts
                               , md_fam_insts = fam_insts
                               , md_rules     = rules
@@ -399,7 +399,7 @@ typecheckIfacesForMerging mod ifaces tc_env_var =
     -- TODO: change tcIfaceDecls to accept w/o Fingerprint
     names_w_things <- tcIfaceDecls ignore_prags (map (\x -> (fingerprint0, x))
                                                   (occEnvElts decl_env))
-    let global_type_env = mkNameEnv names_w_things
+    let global_type_env = TypeEnv (mkNameEnv names_w_things)
     writeMutVar tc_env_var global_type_env
 
     -- OK, now typecheck each ModIface using this environment
@@ -408,7 +408,7 @@ typecheckIfacesForMerging mod ifaces tc_env_var =
         type_env <- fixM $ \type_env ->
             setImplicitEnvM type_env $ do
                 decls <- tcIfaceDecls ignore_prags (mi_decls iface)
-                return (mkNameEnv decls)
+                return (TypeEnv $ mkNameEnv decls)
         -- But note that we use this type_env to typecheck references to DFun
         -- in 'IfaceInst'
         setImplicitEnvM type_env $ do
@@ -448,7 +448,7 @@ typecheckIfaceForInstantiate nsubst iface =
     type_env <- fixM $ \type_env ->
         setImplicitEnvM type_env $ do
             decls     <- tcIfaceDecls ignore_prags (mi_decls iface)
-            return (mkNameEnv decls)
+            return (TypeEnv $ mkNameEnv decls)
     -- See Note [rnIfaceNeverExported]
     setImplicitEnvM type_env $ do
     insts     <- mapM tcIfaceInst (mi_insts iface)
@@ -1790,7 +1790,7 @@ tcIfaceGlobal name
                 | nameIsLocalOrFrom mod name
                 -> do           -- It's defined in the module being compiled
                 { type_env <- setLclEnv () get_type_env         -- yuk
-                ; case lookupNameEnv type_env name of
+                ; case lookupTypeEnv type_env name of
                     Just thing -> return thing
                     -- See Note [Knot-tying fallback on boot]
                     Nothing   -> via_external
diff --git a/compiler/GHC/Tc/Gen/Splice.hs b/compiler/GHC/Tc/Gen/Splice.hs
index ab45f3f373..d0884d8a75 100644
--- a/compiler/GHC/Tc/Gen/Splice.hs
+++ b/compiler/GHC/Tc/Gen/Splice.hs
@@ -109,6 +109,7 @@ import GHC.Types.Error
 import GHC.Types.Fixity as Hs
 import GHC.Types.Annotations
 import GHC.Types.Name
+import GHC.Types.TypeEnv
 import GHC.Serialized
 
 import GHC.Unit.Finder
@@ -1567,7 +1568,7 @@ tcLookupTh name
                 Just thing -> return thing;
                 Nothing    ->
 
-          case lookupNameEnv (tcg_type_env gbl_env) name of {
+          case lookupTypeEnv (tcg_type_env gbl_env) name of {
                 Just thing -> return (AGlobal thing);
                 Nothing    ->
 
diff --git a/compiler/GHC/Tc/Utils/Backpack.hs b/compiler/GHC/Tc/Utils/Backpack.hs
index 066755e8f7..8981ed1fdc 100644
--- a/compiler/GHC/Tc/Utils/Backpack.hs
+++ b/compiler/GHC/Tc/Utils/Backpack.hs
@@ -154,7 +154,7 @@ checkHsigIface tcg_env gr sig_iface
     -- have to look up the right name.
     sig_type_occ_env = mkOccEnv
                      . map (\t -> (nameOccName (getName t), t))
-                     $ nameEnvElts sig_type_env
+                     $ typeEnvElts sig_type_env
     dfun_names = map getName sig_insts
     check_export name
       -- Skip instances, we'll check them later
diff --git a/compiler/GHC/Tc/Utils/Env.hs b/compiler/GHC/Tc/Utils/Env.hs
index 526bb489ac..7d2dd97304 100644
--- a/compiler/GHC/Tc/Utils/Env.hs
+++ b/compiler/GHC/Tc/Utils/Env.hs
@@ -241,7 +241,7 @@ tcLookupGlobal :: Name -> TcM TyThing
 tcLookupGlobal name
   = do  {    -- Try local envt
           env <- getGblEnv
-        ; case lookupNameEnv (tcg_type_env env) name of {
+        ; case lookupTypeEnv (tcg_type_env env) name of {
                 Just thing -> return thing ;
                 Nothing    ->
 
@@ -264,7 +264,7 @@ tcLookupGlobal name
 tcLookupGlobalOnly :: Name -> TcM TyThing
 tcLookupGlobalOnly name
   = do { env <- getGblEnv
-       ; return $ case lookupNameEnv (tcg_type_env env) name of
+       ; return $ case lookupTypeEnv (tcg_type_env env) name of
                     Just thing -> thing
                     Nothing    -> pprPanic "tcLookupGlobalOnly" (ppr name) }
 
@@ -407,8 +407,9 @@ tcExtendRecEnv :: [(Name,TyThing)] -> TcM r -> TcM r
 -- Just like tcExtendGlobalEnv, except the argument is a list of pairs
 tcExtendRecEnv gbl_stuff thing_inside
  = do  { tcg_env <- getGblEnv
-       ; let ge'      = extendNameEnvList (tcg_type_env tcg_env) gbl_stuff
-             tcg_env' = tcg_env { tcg_type_env = ge' }
+       ; let (TypeEnv te) = (tcg_type_env tcg_env)
+       ; let ge'      = extendNameEnvList te gbl_stuff
+             tcg_env' = tcg_env { tcg_type_env = TypeEnv ge' }
          -- No need for setGlobalTypeEnv (which side-effects the
          -- tcg_type_env_var); tcExtendRecEnv is used just
          -- when kind-check a group of type/class decls. It would
diff --git a/compiler/GHC/Tc/Utils/Monad.hs b/compiler/GHC/Tc/Utils/Monad.hs
index a3c087c4da..fa3c30723b 100644
--- a/compiler/GHC/Tc/Utils/Monad.hs
+++ b/compiler/GHC/Tc/Utils/Monad.hs
@@ -245,7 +245,7 @@ initTc hsc_env hsc_src keep_rn_syntax mod loc do_this
         dfun_n_var   <- newIORef emptyOccSet ;
         type_env_var <- case hsc_type_env_var hsc_env of {
                            Just (_mod, te_var) -> return te_var ;
-                           Nothing             -> newIORef emptyNameEnv } ;
+                           Nothing             -> newIORef emptyTypeEnv } ;
 
         dependent_files_var <- newIORef [] ;
         static_wc_var       <- newIORef emptyWC ;
@@ -295,7 +295,7 @@ initTc hsc_env hsc_src keep_rn_syntax mod loc do_this
                                      || moduleUnit mod == bignumUnit
                                      then Just []  -- See Note [Default types]
                                      else Nothing,
-                tcg_type_env       = emptyNameEnv,
+                tcg_type_env       = emptyTypeEnv,
                 tcg_type_env_var   = type_env_var,
                 tcg_inst_env       = emptyInstEnv,
                 tcg_fam_inst_env   = emptyFamInstEnv,
diff --git a/compiler/GHC/Tc/Utils/Zonk.hs b/compiler/GHC/Tc/Utils/Zonk.hs
index 4d4860c7e1..3cde568511 100644
--- a/compiler/GHC/Tc/Utils/Zonk.hs
+++ b/compiler/GHC/Tc/Utils/Zonk.hs
@@ -79,7 +79,6 @@ import GHC.Core
 import GHC.Core.Predicate
 
 import GHC.Types.Name
-import GHC.Types.Name.Env
 import GHC.Types.Var
 import GHC.Types.Var.Env
 import GHC.Types.Id
@@ -408,7 +407,7 @@ setZonkType ze flexi = ze { ze_flexi = flexi }
 
 zonkEnvIds :: ZonkEnv -> TypeEnv
 zonkEnvIds (ZonkEnv { ze_id_env = id_env})
-  = mkNameEnv [(getName id, AnId id) | id <- nonDetEltsUFM id_env]
+  = mkTypeEnv [AnId id | id <- nonDetEltsUFM id_env]
   -- It's OK to use nonDetEltsUFM here because we forget the ordering
   -- immediately by creating a TypeEnv
 
diff --git a/compiler/GHC/Types/TypeEnv.hs b/compiler/GHC/Types/TypeEnv.hs
index 1b8fcd0b35..3cf0dcd537 100644
--- a/compiler/GHC/Types/TypeEnv.hs
+++ b/compiler/GHC/Types/TypeEnv.hs
@@ -1,5 +1,5 @@
 module GHC.Types.TypeEnv
-   ( TypeEnv
+   ( TypeEnv(..)
    , emptyTypeEnv
    , lookupTypeEnv
    , mkTypeEnv
@@ -16,6 +16,8 @@ module GHC.Types.TypeEnv
    , typeEnvDataCons
    , typeEnvCoAxioms
    , typeEnvClasses
+   , elemTypeEnv
+   , mapTypeEnv
    )
 where
 
@@ -33,10 +35,11 @@ import GHC.Types.Name
 import GHC.Types.Name.Env
 import GHC.Types.Var
 import GHC.Types.TyThing
+import GHC.Utils.Outputable
 
 -- | A map from 'Name's to 'TyThing's, constructed by typechecking
 -- local declarations or interface files
-type TypeEnv = NameEnv TyThing
+data TypeEnv = TypeEnv !(NameEnv TyThing)
 
 emptyTypeEnv    :: TypeEnv
 typeEnvElts     :: TypeEnv -> [TyThing]
@@ -48,8 +51,8 @@ typeEnvDataCons :: TypeEnv -> [DataCon]
 typeEnvClasses  :: TypeEnv -> [Class]
 lookupTypeEnv   :: TypeEnv -> Name -> Maybe TyThing
 
-emptyTypeEnv        = emptyNameEnv
-typeEnvElts     env = nameEnvElts env
+emptyTypeEnv        = TypeEnv emptyNameEnv
+typeEnvElts     (TypeEnv env) = nameEnvElts env
 typeEnvTyCons   env = [tc | ATyCon tc   <- typeEnvElts env]
 typeEnvCoAxioms env = [ax | ACoAxiom ax <- typeEnvElts env]
 typeEnvIds      env = [id | AnId id     <- typeEnvElts env]
@@ -61,12 +64,21 @@ typeEnvClasses  env = [cl | tc <- typeEnvTyCons env,
 mkTypeEnv :: [TyThing] -> TypeEnv
 mkTypeEnv things = extendTypeEnvList emptyTypeEnv things
 
+elemTypeEnv :: Name -> TypeEnv -> Bool
+elemTypeEnv n (TypeEnv env) = elemNameEnv n env
+
+mapTypeEnv :: (TyThing -> TyThing) -> TypeEnv -> TypeEnv
+mapTypeEnv f (TypeEnv env) = TypeEnv (mapNameEnv f env)
+
 mkTypeEnvWithImplicits :: [TyThing] -> TypeEnv
 mkTypeEnvWithImplicits things =
   mkTypeEnv things
-    `plusNameEnv`
+    `plusTypeEnv`
   mkTypeEnv (concatMap implicitTyThings things)
 
+plusTypeEnv :: TypeEnv -> TypeEnv -> TypeEnv
+plusTypeEnv (TypeEnv t1) (TypeEnv t2) = TypeEnv (t1 `plusNameEnv` t2)
+
 typeEnvFromEntities :: [Id] -> [TyCon] -> [PatSyn] -> [FamInst] -> TypeEnv
 typeEnvFromEntities ids tcs patsyns famInsts =
   mkTypeEnv (   map AnId ids
@@ -78,19 +90,19 @@ typeEnvFromEntities ids tcs patsyns famInsts =
  where
   all_tcs = tcs ++ famInstsRepTyCons famInsts
 
-lookupTypeEnv = lookupNameEnv
+lookupTypeEnv (TypeEnv t) = lookupNameEnv t
 
 -- Extend the type environment
 extendTypeEnv :: TypeEnv -> TyThing -> TypeEnv
-extendTypeEnv env thing = extendNameEnv env (getName thing) thing
+extendTypeEnv (TypeEnv env) thing = TypeEnv (extendNameEnv env (getName thing) thing)
 
 extendTypeEnvList :: TypeEnv -> [TyThing] -> TypeEnv
-extendTypeEnvList env things = foldl' extendTypeEnv env things
+extendTypeEnvList env things = (foldl' extendTypeEnv env things)
 
 extendTypeEnvWithIds :: TypeEnv -> [Id] -> TypeEnv
-extendTypeEnvWithIds env ids
-  = extendNameEnvList env [(getName id, AnId id) | id <- ids]
+extendTypeEnvWithIds (TypeEnv env) ids
+  = TypeEnv (extendNameEnvList env [(getName id, AnId id) | id <- ids])
 
-plusTypeEnv :: TypeEnv -> TypeEnv -> TypeEnv
-plusTypeEnv env1 env2 = plusNameEnv env1 env2
+instance Outputable TypeEnv where
+  ppr (TypeEnv t) = ppr t