refactor: uniqAway return type

12 files changed, 46 insertions, 43 deletions

diff --git a/compiler/GHC/Core/Coercion.hs b/compiler/GHC/Core/Coercion.hs
index 4c487401b0..012ac47193 100644
--- a/compiler/GHC/Core/Coercion.hs
+++ b/compiler/GHC/Core/Coercion.hs
@@ -2171,13 +2171,13 @@ liftCoSubstTyVarBndrUsing :: (LiftingContext -> Type -> (CoercionN, a))
                            -> (LiftingContext, TyVar, CoercionN, a)
 liftCoSubstTyVarBndrUsing fun lc@(LC subst cenv) old_var
   = ASSERT( isTyVar old_var )
-    ( LC (subst `extendTCvInScope` new_var) new_cenv
+    ( LC (subst `setTCvInScope` new_scope) new_cenv
     , new_var, eta, stuff )
   where
     old_kind     = tyVarKind old_var
     (eta, stuff) = fun lc old_kind
     k1           = coercionLKind eta
-    new_var      = uniqAway (getTCvInScope subst) (setVarType old_var k1)
+    (new_var, new_scope) = uniqAway (getTCvInScope subst) (setVarType old_var k1)
 
     lifted   = mkGReflRightCo Nominal (TyVarTy new_var) eta
                -- :: new_var ~ new_var |> eta
@@ -2189,13 +2189,13 @@ liftCoSubstCoVarBndrUsing :: (LiftingContext -> Type -> (CoercionN, a))
                            -> (LiftingContext, CoVar, CoercionN, a)
 liftCoSubstCoVarBndrUsing fun lc@(LC subst cenv) old_var
   = ASSERT( isCoVar old_var )
-    ( LC (subst `extendTCvInScope` new_var) new_cenv
+    ( LC (setTCvInScope subst new_scope) new_cenv
     , new_var, kind_co, stuff )
   where
     old_kind     = coVarKind old_var
     (eta, stuff) = fun lc old_kind
     k1           = coercionLKind eta
-    new_var      = uniqAway (getTCvInScope subst) (setVarType old_var k1)
+    (new_var, new_scope) = uniqAway (getTCvInScope subst) (setVarType old_var k1)
 
     -- old_var :: s1  ~r s2
     -- eta     :: (s1' ~r s2') ~N (t1 ~r t2)
diff --git a/compiler/GHC/Core/Lint.hs b/compiler/GHC/Core/Lint.hs
index 1ee9477a23..6ce1c84622 100644
--- a/compiler/GHC/Core/Lint.hs
+++ b/compiler/GHC/Core/Lint.hs
@@ -1510,7 +1510,7 @@ lintTyCoBndr :: TyCoVar -> (LintedTyCoVar -> LintM a) -> LintM a
 lintTyCoBndr tcv thing_inside
   = do { subst <- getTCvSubst
        ; kind' <- lintType (varType tcv)
-       ; let tcv' = uniqAway (getTCvInScope subst) $
+       ; let (tcv', _in_scope) = uniqAway (getTCvInScope subst) $
                     setVarType tcv kind'
              subst' = extendTCvSubstWithClone subst tcv tcv'
        ; when (isCoVar tcv) $
diff --git a/compiler/GHC/Core/Opt/Arity.hs b/compiler/GHC/Core/Opt/Arity.hs
index ef5a047184..384f99affc 100644
--- a/compiler/GHC/Core/Opt/Arity.hs
+++ b/compiler/GHC/Core/Opt/Arity.hs
@@ -66,6 +66,7 @@ import GHC.Utils.Panic
 import GHC.Data.FastString
 import GHC.Data.Pair
 import GHC.Utils.Misc
+import GHC.Core.TyCo.Subst
 
 {-
 ************************************************************************
@@ -1912,8 +1913,8 @@ freshEtaId n subst ty
       = (subst', eta_id')
       where
         Scaled mult' ty' = Type.substScaledTyUnchecked subst ty
-        eta_id' = uniqAway (getTCvInScope subst) $
+        (eta_id', new_scope) = uniqAway (getTCvInScope subst) $
                   mkSysLocalOrCoVar (fsLit "eta") (mkBuiltinUnique n) mult' ty'
                   -- "OrCoVar" since this can be used to eta-expand
                   -- coercion abstractions
-        subst'  = extendTCvInScope subst eta_id'
+        subst'  = setTCvInScope subst new_scope
diff --git a/compiler/GHC/Core/Opt/Exitify.hs b/compiler/GHC/Core/Opt/Exitify.hs
index 2b34992d72..2852981435 100644
--- a/compiler/GHC/Core/Opt/Exitify.hs
+++ b/compiler/GHC/Core/Opt/Exitify.hs
@@ -263,7 +263,8 @@ mkExitJoinId in_scope ty join_arity = do
     fs <- get
     let avoid = in_scope `extendInScopeSetList` (map fst fs)
                          `extendInScopeSet` exit_id_tmpl -- just cosmetics
-    return (uniqAway avoid exit_id_tmpl)
+    let (new_id, in_scope) = uniqAway avoid exit_id_tmpl
+    return new_id
   where
     exit_id_tmpl = mkSysLocal (fsLit "exit") initExitJoinUnique Many ty
                     `asJoinId` join_arity
diff --git a/compiler/GHC/Core/Opt/Simplify/Env.hs b/compiler/GHC/Core/Opt/Simplify/Env.hs
index 2430b3e234..06332be1cd 100644
--- a/compiler/GHC/Core/Opt/Simplify/Env.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Env.hs
@@ -855,7 +855,7 @@ subst_id_bndr env@(SimplEnv { seInScope = in_scope, seIdSubst = id_subst })
     -- afresh with both seInScope and seIdSubst
   where
     -- See Note [Bangs in the Simplifier]
-    !id1  = uniqAway in_scope old_id
+    !(id1, new_in_scope)  = uniqAway in_scope old_id
     !id2  = substIdType env id1
     !id3  = zapFragileIdInfo id2       -- Zaps rules, worker-info, unfolding
                                       -- and fragile OccInfo
@@ -869,8 +869,6 @@ subst_id_bndr env@(SimplEnv { seInScope = in_scope, seIdSubst = id_subst })
               | otherwise
               = delVarEnv id_subst old_id
 
-    !new_in_scope = in_scope `extendInScopeSet` new_id
-
 ------------------------------------
 seqTyVar :: TyVar -> ()
 seqTyVar b = b `seq` ()
diff --git a/compiler/GHC/Core/SimpleOpt.hs b/compiler/GHC/Core/SimpleOpt.hs
index 6353dcda6f..c5c3626d34 100644
--- a/compiler/GHC/Core/SimpleOpt.hs
+++ b/compiler/GHC/Core/SimpleOpt.hs
@@ -629,13 +629,12 @@ subst_opt_id_bndr env@(SOE { soe_subst = subst, soe_inl = inl }) old_id
   where
     Subst in_scope id_subst tv_subst cv_subst = subst
 
-    id1    = uniqAway in_scope old_id
+    (id1, new_in_scope) = uniqAway in_scope old_id
     id2    = updateIdTypeAndMult (substTy subst) id1
     new_id = zapFragileIdInfo id2
              -- Zaps rules, unfolding, and fragile OccInfo
              -- The unfolding and rules will get added back later, by add_info
 
-    new_in_scope = in_scope `extendInScopeSet` new_id
 
     no_change = new_id == old_id
 
diff --git a/compiler/GHC/Core/Subst.hs b/compiler/GHC/Core/Subst.hs
index 1c7d138574..56d09cf30f 100644
--- a/compiler/GHC/Core/Subst.hs
+++ b/compiler/GHC/Core/Subst.hs
@@ -470,9 +470,9 @@ substIdBndr :: SDoc
 
 substIdBndr _doc rec_subst subst@(Subst in_scope env tvs cvs) old_id
   = -- pprTrace "substIdBndr" (doc $$ ppr old_id $$ ppr in_scope) $
-    (Subst (in_scope `extendInScopeSet` new_id) new_env tvs cvs, new_id)
+    (Subst new_scope new_env tvs cvs, new_id)
   where
-    id1 = uniqAway in_scope old_id      -- id1 is cloned if necessary
+    (id1, new_scope) = uniqAway in_scope old_id      -- id1 is cloned if necessary
     id2 | no_type_change = id1
         | otherwise      = updateIdTypeAndMult (substTy subst) id1
 
diff --git a/compiler/GHC/Core/TyCo/Subst.hs b/compiler/GHC/Core/TyCo/Subst.hs
index e9c9b85a23..fb46b7ab33 100644
--- a/compiler/GHC/Core/TyCo/Subst.hs
+++ b/compiler/GHC/Core/TyCo/Subst.hs
@@ -20,7 +20,7 @@ module GHC.Core.TyCo.Subst
         getCvSubstEnv, getTCvInScope, getTCvSubstRangeFVs,
         isInScope, elemTCvSubst, notElemTCvSubst,
         setTvSubstEnv, setCvSubstEnv, zapTCvSubst,
-        extendTCvInScope, extendTCvInScopeList, extendTCvInScopeSet,
+        extendTCvInScope, extendTCvInScopeList, extendTCvInScopeSet, setTCvInScope,
         extendTCvSubst, extendTCvSubstWithClone,
         extendCvSubst, extendCvSubstWithClone,
         extendTvSubst, extendTvSubstBinderAndInScope, extendTvSubstWithClone,
@@ -316,6 +316,10 @@ extendTCvInScope :: TCvSubst -> Var -> TCvSubst
 extendTCvInScope (TCvSubst in_scope tenv cenv) var
   = TCvSubst (extendInScopeSet in_scope var) tenv cenv
 
+setTCvInScope :: TCvSubst -> InScopeSet -> TCvSubst
+setTCvInScope (TCvSubst _ tenv cenv) in_scope
+  = TCvSubst in_scope tenv cenv
+
 extendTCvInScopeList :: TCvSubst -> [Var] -> TCvSubst
 extendTCvInScopeList (TCvSubst in_scope tenv cenv) vars
   = TCvSubst (extendInScopeSetList in_scope vars) tenv cenv
@@ -888,7 +892,7 @@ substForAllCoTyVarBndrUsing :: Bool  -- apply sym to binder?
                             -> (TCvSubst, TyVar, KindCoercion)
 substForAllCoTyVarBndrUsing sym sco (TCvSubst in_scope tenv cenv) old_var old_kind_co
   = ASSERT( isTyVar old_var )
-    ( TCvSubst (in_scope `extendInScopeSet` new_var) new_env cenv
+    ( TCvSubst new_scope new_env cenv
     , new_var, new_kind_co )
   where
     new_env | no_change && not sym = delVarEnv tenv old_var
@@ -908,7 +912,7 @@ substForAllCoTyVarBndrUsing sym sco (TCvSubst in_scope tenv cenv) old_var old_ki
     -- we want. We don't want to do substitution once more. Also, in most cases,
     -- new_kind_co is a Refl, in which case coercionKind is really fast.
 
-    new_var  = uniqAway in_scope (setTyVarKind old_var new_ki1)
+    (new_var, new_scope)  = uniqAway in_scope (setTyVarKind old_var new_ki1)
 
 substForAllCoCoVarBndrUsing :: Bool  -- apply sym to binder?
                             -> (Coercion -> Coercion)  -- transformation to kind co
@@ -917,7 +921,7 @@ substForAllCoCoVarBndrUsing :: Bool  -- apply sym to binder?
 substForAllCoCoVarBndrUsing sym sco (TCvSubst in_scope tenv cenv)
                             old_var old_kind_co
   = ASSERT( isCoVar old_var )
-    ( TCvSubst (in_scope `extendInScopeSet` new_var) tenv new_cenv
+    ( TCvSubst new_scope tenv new_cenv
     , new_var, new_kind_co )
   where
     new_cenv | no_change && not sym = delVarEnv cenv old_var
@@ -931,7 +935,7 @@ substForAllCoCoVarBndrUsing sym sco (TCvSubst in_scope tenv cenv)
 
     Pair h1 h2 = coercionKind new_kind_co
 
-    new_var       = uniqAway in_scope $ mkCoVar (varName old_var) new_var_type
+    (new_var, new_scope)       = uniqAway in_scope $ mkCoVar (varName old_var) new_var_type
     new_var_type  | sym       = h2
                   | otherwise = h1
 
@@ -985,7 +989,7 @@ substTyVarBndrUsing
 substTyVarBndrUsing subst_fn subst@(TCvSubst in_scope tenv cenv) old_var
   = ASSERT2( _no_capture, pprTyVar old_var $$ pprTyVar new_var $$ ppr subst )
     ASSERT( isTyVar old_var )
-    (TCvSubst (in_scope `extendInScopeSet` new_var) new_env cenv, new_var)
+    (TCvSubst new_scope new_env cenv, new_var)
   where
     new_env | no_change = delVarEnv tenv old_var
             | otherwise = extendVarEnv tenv old_var (TyVarTy new_var)
@@ -1006,7 +1010,8 @@ substTyVarBndrUsing subst_fn subst@(TCvSubst in_scope tenv cenv) old_var
         --      (\x.e) with id_subst = [x |-> e']
         -- Here we must simply zap the substitution for x
 
-    new_var | no_kind_change = uniqAway in_scope old_var
+    (new_var, new_scope)
+            | no_kind_change = uniqAway in_scope old_var
             | otherwise = uniqAway in_scope $
                           setTyVarKind old_var (subst_fn subst old_ki)
         -- The uniqAway part makes sure the new variable is not already in scope
@@ -1019,7 +1024,7 @@ substCoVarBndrUsing
   -> TCvSubst -> CoVar -> (TCvSubst, CoVar)
 substCoVarBndrUsing subst_fn subst@(TCvSubst in_scope tenv cenv) old_var
   = ASSERT( isCoVar old_var )
-    (TCvSubst (in_scope `extendInScopeSet` new_var) tenv new_cenv, new_var)
+    (TCvSubst new_scope tenv new_cenv, new_var)
   where
     new_co         = mkCoVarCo new_var
     no_kind_change = noFreeVarsOfTypes [t1, t2]
@@ -1028,7 +1033,7 @@ substCoVarBndrUsing subst_fn subst@(TCvSubst in_scope tenv cenv) old_var
     new_cenv | no_change = delVarEnv cenv old_var
              | otherwise = extendVarEnv cenv old_var new_co
 
-    new_var = uniqAway in_scope subst_old_var
+    (new_var, new_scope) = uniqAway in_scope subst_old_var
     subst_old_var = mkCoVar (varName old_var) new_var_type
 
     (_, _, t1, t2, role) = coVarKindsTypesRole old_var
diff --git a/compiler/GHC/Core/Unify.hs b/compiler/GHC/Core/Unify.hs
index b3b7e269bd..1ffbe6b960 100644
--- a/compiler/GHC/Core/Unify.hs
+++ b/compiler/GHC/Core/Unify.hs
@@ -1997,9 +1997,9 @@ coreFlattenVarBndr subst env tv
     -- See Note [Flattening type-family applications when matching instances], wrinkle 2B.
     kind          = varType tv
     (env1, kind') = coreFlattenTy subst env kind
-    tv'           = uniqAway (fe_in_scope env1) (setVarType tv kind')
+    (tv', new_scope) = uniqAway (fe_in_scope env1) (setVarType tv kind')
     subst'        = extendVarEnv subst tv (mkTyVarTy tv')
-    env2          = updateInScopeSet env1 (flip extendInScopeSet tv')
+    env2          = updateInScopeSet env1 (const new_scope)
 
 coreFlattenTyFamApp :: TvSubstEnv -> FlattenEnv
                     -> TyCon         -- type family tycon
@@ -2010,13 +2010,13 @@ coreFlattenTyFamApp tv_subst env fam_tc fam_args
       Just (tv, _, _) -> (env', mkAppTys (mkTyVarTy tv) leftover_args')
       Nothing ->
         let tyvar_name = mkFlattenFreshTyName fam_tc
-            tv         = uniqAway in_scope $
+            (tv, new_in_scope) = uniqAway in_scope $
                          mkTyVar tyvar_name (typeKind fam_ty)
 
             ty'   = mkAppTys (mkTyVarTy tv) leftover_args'
             env'' = env' { fe_type_map = extendTypeMap type_map fam_ty
                                                        (tv, fam_tc, sat_fam_args)
-                         , fe_in_scope = extendInScopeSet in_scope tv }
+                         , fe_in_scope = new_in_scope }
         in (env'', ty')
   where
     arity = tyConArity fam_tc
diff --git a/compiler/GHC/Stg/CSE.hs b/compiler/GHC/Stg/CSE.hs
index bc266d20ba..fcc040597f 100644
--- a/compiler/GHC/Stg/CSE.hs
+++ b/compiler/GHC/Stg/CSE.hs
@@ -250,9 +250,9 @@ substBndr :: CseEnv -> InId -> (CseEnv, OutId)
 substBndr env old_id
   = (new_env, new_id)
   where
-    new_id = uniqAway (ce_in_scope env) old_id
+    (new_id, new_scope) = uniqAway (ce_in_scope env) old_id
     no_change = new_id == old_id
-    env' = env { ce_in_scope = ce_in_scope env `extendInScopeSet` new_id }
+    env' = env { ce_in_scope = new_scope }
     new_env | no_change = env'
             | otherwise = env' { ce_subst = extendVarEnv (ce_subst env) old_id new_id }
 
diff --git a/compiler/GHC/Stg/Subst.hs b/compiler/GHC/Stg/Subst.hs
index dce2859262..830522c16c 100644
--- a/compiler/GHC/Stg/Subst.hs
+++ b/compiler/GHC/Stg/Subst.hs
@@ -39,9 +39,8 @@ substBndr :: Id -> Subst -> (Id, Subst)
 substBndr id (Subst in_scope env)
   = (new_id, Subst new_in_scope new_env)
   where
-    new_id = uniqAway in_scope id
+    (new_id, new_in_scope) = uniqAway in_scope id
     no_change = new_id == id -- in case nothing shadowed
-    new_in_scope = in_scope `extendInScopeSet` new_id
     new_env
       | no_change = delVarEnv env id
       | otherwise = extendVarEnv env id new_id
diff --git a/compiler/GHC/Types/Var/Env.hs b/compiler/GHC/Types/Var/Env.hs
index ee9e2d399b..cca960ae01 100644
--- a/compiler/GHC/Types/Var/Env.hs
+++ b/compiler/GHC/Types/Var/Env.hs
@@ -183,13 +183,13 @@ since they are only locally unique. In particular, two successive calls to
 
 -- | @uniqAway in_scope v@ finds a unique that is not used in the
 -- in-scope set, and gives that to v. See Note [Local uniques].
-uniqAway :: InScopeSet -> Var -> Var
+uniqAway :: InScopeSet -> Var -> (Var, InScopeSet)
 -- It starts with v's current unique, of course, in the hope that it won't
 -- have to change, and thereafter uses the successor to the last derived unique
 -- found in the in-scope set.
 uniqAway in_scope var
-  | var `elemInScopeSet` in_scope = uniqAway' in_scope var      -- Make a new one
-  | otherwise                     = var                         -- Nothing to do
+  | var `elemInScopeSet` in_scope = let new_var = uniqAway' in_scope var in (new_var, extendInScopeSet  in_scope new_var)      -- Make a new one
+  | otherwise                     = (var, in_scope)                       -- Nothing to do
 
 uniqAway' :: InScopeSet -> Var -> Var
 -- This one *always* makes up a new variable
@@ -317,9 +317,9 @@ rnBndrL :: RnEnv2 -> Var -> (RnEnv2, Var)
 rnBndrL (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL
   = (RV2 { envL     = extendVarEnv envL bL new_b
          , envR     = envR
-         , in_scope = extendInScopeSet in_scope new_b }, new_b)
+         , in_scope = new_scope }, new_b)
   where
-    new_b = uniqAway in_scope bL
+    (new_b, new_scope) = uniqAway in_scope bL
 
 rnBndrR :: RnEnv2 -> Var -> (RnEnv2, Var)
 -- ^ Similar to 'rnBndr2' but used when there's a binder on the right
@@ -327,9 +327,9 @@ rnBndrR :: RnEnv2 -> Var -> (RnEnv2, Var)
 rnBndrR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
   = (RV2 { envR     = extendVarEnv envR bR new_b
          , envL     = envL
-         , in_scope = extendInScopeSet in_scope new_b }, new_b)
+         , in_scope = new_scope }, new_b)
   where
-    new_b = uniqAway in_scope bR
+    (new_b, new_scope) = uniqAway in_scope bR
 
 rnEtaL :: RnEnv2 -> Var -> (RnEnv2, Var)
 -- ^ Similar to 'rnBndrL' but used for eta expansion
@@ -337,9 +337,9 @@ rnEtaL :: RnEnv2 -> Var -> (RnEnv2, Var)
 rnEtaL (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL
   = (RV2 { envL     = extendVarEnv envL bL new_b
          , envR     = extendVarEnv envR new_b new_b     -- Note [Eta expansion]
-         , in_scope = extendInScopeSet in_scope new_b }, new_b)
+         , in_scope = new_scope }, new_b)
   where
-    new_b = uniqAway in_scope bL
+    (new_b, new_scope) = uniqAway in_scope bL
 
 rnEtaR :: RnEnv2 -> Var -> (RnEnv2, Var)
 -- ^ Similar to 'rnBndr2' but used for eta expansion
@@ -347,9 +347,9 @@ rnEtaR :: RnEnv2 -> Var -> (RnEnv2, Var)
 rnEtaR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
   = (RV2 { envL     = extendVarEnv envL new_b new_b     -- Note [Eta expansion]
          , envR     = extendVarEnv envR bR new_b
-         , in_scope = extendInScopeSet in_scope new_b }, new_b)
+         , in_scope = new_scope }, new_b)
   where
-    new_b = uniqAway in_scope bR
+    (new_b, new_scope) = uniqAway in_scope bR
 
 delBndrL, delBndrR :: RnEnv2 -> Var -> RnEnv2
 delBndrL rn@(RV2 { envL = env, in_scope = in_scope }) v