Don't track free variables in STG syntax by default

Summary:
Currently, `CoreToStg` annotates `StgRhsClosure`s with their set of non-global
free variables.  This free variable information is only needed in the final
code generation step (i.e. `StgCmm.codeGen`), which leads to transformations
such as `StgCse` and `StgUnarise` having to maintain this information.

This is tiresome and unnecessary, so this patch introduces a trees-to-grow-like
approach that only introduces the free variable set into the syntax tree in the
code gen pass, along with a free variable analysis on STG terms to generate
that information.

Fixes #15754.

Reviewers: simonpj, osa1, bgamari, simonmar

Reviewed By: osa1

Subscribers: rwbarton, carter

GHC Trac Issues: #15754

Differential Revision: https://phabricator.haskell.org/D5324

3 files changed, 8 insertions, 35 deletions

diff --git a/compiler/simplStg/StgCse.hs b/compiler/simplStg/StgCse.hs
index fe7943c7d8..a22a7c1400 100644
--- a/compiler/simplStg/StgCse.hs
+++ b/compiler/simplStg/StgCse.hs
@@ -227,9 +227,6 @@ substArg :: CseEnv -> InStgArg -> OutStgArg
 substArg env (StgVarArg from) = StgVarArg (substVar env from)
 substArg _   (StgLitArg lit)  = StgLitArg lit
 
-substVars :: CseEnv -> [InId] -> [OutId]
-substVars env = map (substVar env)
-
 substVar :: CseEnv -> InId -> OutId
 substVar env id = fromMaybe id $ lookupVarEnv (ce_subst env) id
 
@@ -284,9 +281,9 @@ stgCseTopLvl in_scope (StgTopLifted (StgRec eqs))
   where in_scope' = in_scope `extendInScopeSetList` [ bndr | (bndr, _) <- eqs ]
 
 stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
-stgCseTopLvlRhs in_scope (StgRhsClosure ccs occs upd args body)
+stgCseTopLvlRhs in_scope (StgRhsClosure ext ccs upd args body)
     = let body' = stgCseExpr (initEnv in_scope) body
-      in  StgRhsClosure ccs occs upd args body'
+      in  StgRhsClosure ext ccs upd args body'
 stgCseTopLvlRhs _ (StgRhsCon ccs dataCon args)
     = StgRhsCon ccs dataCon args
 
@@ -402,12 +399,11 @@ stgCseRhs env bndr (StgRhsCon ccs dataCon args)
           pair = (bndr, StgRhsCon ccs dataCon args')
       in (Just pair, env')
   where args' = substArgs env args
-stgCseRhs env bndr (StgRhsClosure ccs occs upd args body)
+stgCseRhs env bndr (StgRhsClosure ext ccs upd args body)
     = let (env1, args') = substBndrs env args
           env2 = forgetCse env1 -- See note [Free variables of an StgClosure]
           body' = stgCseExpr env2 body
-      in (Just (substVar env bndr, StgRhsClosure ccs occs' upd args' body'), env)
-  where occs' = substVars env occs
+      in (Just (substVar env bndr, StgRhsClosure ext ccs upd args' body'), env)
 
 
 mkStgCase :: StgExpr -> OutId -> AltType -> [StgAlt] -> StgExpr
diff --git a/compiler/simplStg/StgStats.hs b/compiler/simplStg/StgStats.hs
index c548d80707..a2a9a8530f 100644
--- a/compiler/simplStg/StgStats.hs
+++ b/compiler/simplStg/StgStats.hs
@@ -66,9 +66,6 @@ combineSEs = foldr combineSE emptySE
 countOne :: CounterType -> StatEnv
 countOne c = Map.singleton c 1
 
-countN :: CounterType -> Int -> StatEnv
-countN = Map.singleton
-
 {-
 ************************************************************************
 *                                                                      *
@@ -131,9 +128,8 @@ statRhs :: Bool -> (Id, StgRhs) -> StatEnv
 statRhs top (_, StgRhsCon _ _ _)
   = countOne (ConstructorBinds top)
 
-statRhs top (_, StgRhsClosure _ fv u _ body)
-  = statExpr body                       `combineSE`
-    countN FreeVariables (length fv)    `combineSE`
+statRhs top (_, StgRhsClosure _ _ u _ body)
+  = statExpr body `combineSE`
     countOne (
       case u of
         ReEntrant   -> ReEntrantBinds   top
diff --git a/compiler/simplStg/UnariseStg.hs b/compiler/simplStg/UnariseStg.hs
index a46497452f..c3a8bc76e2 100644
--- a/compiler/simplStg/UnariseStg.hs
+++ b/compiler/simplStg/UnariseStg.hs
@@ -281,11 +281,10 @@ unariseBinding rho (StgRec xrhss)
   = StgRec <$> mapM (\(x, rhs) -> (x,) <$> unariseRhs rho rhs) xrhss
 
 unariseRhs :: UnariseEnv -> StgRhs -> UniqSM StgRhs
-unariseRhs rho (StgRhsClosure ccs fvs update_flag args expr)
+unariseRhs rho (StgRhsClosure ext ccs update_flag args expr)
   = do (rho', args1) <- unariseFunArgBinders rho args
        expr' <- unariseExpr rho' expr
-       let fvs' = unariseFreeVars rho fvs
-       return (StgRhsClosure ccs fvs' update_flag args1 expr')
+       return (StgRhsClosure ext ccs update_flag args1 expr')
 
 unariseRhs rho (StgRhsCon ccs con args)
   = ASSERT(not (isUnboxedTupleCon con || isUnboxedSumCon con))
@@ -723,24 +722,6 @@ unariseConArgBinders rho xs = second concat <$> mapAccumLM unariseConArgBinder r
 unariseConArgBinder :: UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
 unariseConArgBinder = unariseArgBinder True
 
-unariseFreeVars :: UnariseEnv -> [InId] -> [OutId]
-unariseFreeVars rho fvs
- = [ v | fv <- fvs, StgVarArg v <- unariseFreeVar rho fv ]
-   -- Notice that we filter out any StgLitArgs
-   -- e.g.   case e of (x :: (# Int | Bool #))
-   --           (# v | #) ->  ... let {g = \y. ..x...} in ...
-   --           (# | w #) -> ...
-   --     Here 'x' is free in g's closure, and the env will have
-   --       x :-> [1, v]
-   --     we want to capture 'v', but not 1, in the free vars
-
-unariseFreeVar :: UnariseEnv -> Id -> [StgArg]
-unariseFreeVar rho x =
-  case lookupVarEnv rho x of
-    Just (MultiVal args) -> args
-    Just (UnaryVal arg)  -> [arg]
-    Nothing              -> [StgVarArg x]
-
 --------------------------------------------------------------------------------
 
 mkIds :: FastString -> [UnaryType] -> UniqSM [Id]