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

4 files changed, 35 insertions, 29 deletions

diff --git a/compiler/codeGen/StgCmm.hs b/compiler/codeGen/StgCmm.hs
index 5b80ba61d9..59ceba8706 100644
--- a/compiler/codeGen/StgCmm.hs
+++ b/compiler/codeGen/StgCmm.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE DataKinds #-}
 
 -----------------------------------------------------------------------------
 --
@@ -44,6 +45,7 @@ import Module
 import Outputable
 import Stream
 import BasicTypes
+import VarSet ( isEmptyVarSet )
 
 import OrdList
 import MkGraph
@@ -57,10 +59,10 @@ codeGen :: DynFlags
         -> Module
         -> [TyCon]
         -> CollectedCCs                -- (Local/global) cost-centres needing declaring/registering.
-        -> [StgTopBinding]             -- Bindings to convert
+        -> [CgStgTopBinding]           -- Bindings to convert
         -> HpcInfo
         -> Stream IO CmmGroup ()       -- Output as a stream, so codegen can
-                                        -- be interleaved with output
+                                       -- be interleaved with output
 
 codeGen dflags this_mod data_tycons
         cost_centre_info stg_binds hpc_info
@@ -117,7 +119,7 @@ This is so that we can write the top level processing in a compositional
 style, with the increasing static environment being plumbed as a state
 variable. -}
 
-cgTopBinding :: DynFlags -> StgTopBinding -> FCode ()
+cgTopBinding :: DynFlags -> CgStgTopBinding -> FCode ()
 cgTopBinding dflags (StgTopLifted (StgNonRec id rhs))
   = do  { id' <- maybeExternaliseId dflags id
         ; let (info, fcode) = cgTopRhs dflags NonRecursive id' rhs
@@ -144,7 +146,7 @@ cgTopBinding dflags (StgTopStringLit id str)
         ; addBindC (litIdInfo dflags id' mkLFStringLit lit)
         }
 
-cgTopRhs :: DynFlags -> RecFlag -> Id -> StgRhs -> (CgIdInfo, FCode ())
+cgTopRhs :: DynFlags -> RecFlag -> Id -> CgStgRhs -> (CgIdInfo, FCode ())
         -- The Id is passed along for setting up a binding...
         -- It's already been externalised if necessary
 
@@ -153,8 +155,8 @@ cgTopRhs dflags _rec bndr (StgRhsCon _cc con args)
       -- con args are always non-void,
       -- see Note [Post-unarisation invariants] in UnariseStg
 
-cgTopRhs dflags rec bndr (StgRhsClosure cc fvs upd_flag args body)
-  = ASSERT(null fvs)    -- There should be no free variables
+cgTopRhs dflags rec bndr (StgRhsClosure fvs cc upd_flag args body)
+  = ASSERT(isEmptyVarSet fvs)    -- There should be no free variables
     cgTopRhsClosure dflags rec bndr cc upd_flag args body
 
 
diff --git a/compiler/codeGen/StgCmmBind.hs b/compiler/codeGen/StgCmmBind.hs
index 004bf90c67..dba122fd0c 100644
--- a/compiler/codeGen/StgCmmBind.hs
+++ b/compiler/codeGen/StgCmmBind.hs
@@ -44,6 +44,7 @@ import Name
 import Module
 import ListSetOps
 import Util
+import UniqSet ( nonDetEltsUniqSet )
 import BasicTypes
 import Outputable
 import FastString
@@ -64,7 +65,7 @@ cgTopRhsClosure :: DynFlags
                 -> CostCentreStack      -- Optional cost centre annotation
                 -> UpdateFlag
                 -> [Id]                 -- Args
-                -> StgExpr
+                -> CgStgExpr
                 -> (CgIdInfo, FCode ())
 
 cgTopRhsClosure dflags rec id ccs upd_flag args body =
@@ -121,7 +122,7 @@ cgTopRhsClosure dflags rec id ccs upd_flag args body =
 --              Non-top-level bindings
 ------------------------------------------------------------------------
 
-cgBind :: StgBinding -> FCode ()
+cgBind :: CgStgBinding -> FCode ()
 cgBind (StgNonRec name rhs)
   = do  { (info, fcode) <- cgRhs name rhs
         ; addBindC info
@@ -190,7 +191,7 @@ cgBind (StgRec pairs)
  -}
 
 cgRhs :: Id
-      -> StgRhs
+      -> CgStgRhs
       -> FCode (
                  CgIdInfo         -- The info for this binding
                , FCode CmmAGraph  -- A computation which will generate the
@@ -206,9 +207,12 @@ cgRhs id (StgRhsCon cc con args)
       -- see Note [Post-unarisation invariants] in UnariseStg
 
 {- See Note [GC recovery] in compiler/codeGen/StgCmmClosure.hs -}
-cgRhs id (StgRhsClosure cc fvs upd_flag args body)
+cgRhs id (StgRhsClosure fvs cc upd_flag args body)
   = do dflags <- getDynFlags
-       mkRhsClosure dflags id cc (nonVoidIds fvs) upd_flag args body
+       mkRhsClosure dflags id cc (nonVoidIds (nonDetEltsUniqSet fvs)) upd_flag args body
+       -- It's OK to use nonDetEltsUniqSet here because we're not aiming for
+       -- bit-for-bit determinism.
+       -- See Note [Unique Determinism and code generation]
 
 ------------------------------------------------------------------------
 --              Non-constructor right hand sides
@@ -218,7 +222,7 @@ mkRhsClosure :: DynFlags -> Id -> CostCentreStack
              -> [NonVoid Id]                    -- Free vars
              -> UpdateFlag
              -> [Id]                            -- Args
-             -> StgExpr
+             -> CgStgExpr
              -> FCode (CgIdInfo, FCode CmmAGraph)
 
 {- mkRhsClosure looks for two special forms of the right-hand side:
@@ -436,7 +440,7 @@ closureCodeBody :: Bool            -- whether this is a top-level binding
                 -> CostCentreStack -- Optional cost centre attached to closure
                 -> [NonVoid Id]    -- incoming args to the closure
                 -> Int             -- arity, including void args
-                -> StgExpr
+                -> CgStgExpr
                 -> [(NonVoid Id, ByteOff)] -- the closure's free vars
                 -> FCode ()
 
@@ -560,7 +564,7 @@ mkSlowEntryCode bndr cl_info arg_regs -- function closure is already in `Node'
 
 -----------------------------------------
 thunkCode :: ClosureInfo -> [(NonVoid Id, ByteOff)] -> CostCentreStack
-          -> LocalReg -> Int -> StgExpr -> FCode ()
+          -> LocalReg -> Int -> CgStgExpr -> FCode ()
 thunkCode cl_info fv_details _cc node arity body
   = do { dflags <- getDynFlags
        ; let node_points = nodeMustPointToIt dflags (closureLFInfo cl_info)
diff --git a/compiler/codeGen/StgCmmBind.hs-boot b/compiler/codeGen/StgCmmBind.hs-boot
index 5840e990c8..8e3dd38ad8 100644
--- a/compiler/codeGen/StgCmmBind.hs-boot
+++ b/compiler/codeGen/StgCmmBind.hs-boot
@@ -1,6 +1,6 @@
 module StgCmmBind where
 
 import StgCmmMonad( FCode )
-import StgSyn( StgBinding )
+import StgSyn( CgStgBinding )
 
-cgBind :: StgBinding -> FCode ()
+cgBind :: CgStgBinding -> FCode ()
diff --git a/compiler/codeGen/StgCmmExpr.hs b/compiler/codeGen/StgCmmExpr.hs
index 7fc9dfc829..e8d111f9d5 100644
--- a/compiler/codeGen/StgCmmExpr.hs
+++ b/compiler/codeGen/StgCmmExpr.hs
@@ -56,7 +56,7 @@ import Data.Function ( on )
 --              cgExpr: the main function
 ------------------------------------------------------------------------
 
-cgExpr  :: StgExpr -> FCode ReturnKind
+cgExpr  :: CgStgExpr -> FCode ReturnKind
 
 cgExpr (StgApp fun args)     = cgIdApp fun args
 
@@ -114,7 +114,7 @@ bound only to stable things like stack locations..  The 'e' part will
 execute *next*, just like the scrutinee of a case. -}
 
 -------------------------
-cgLneBinds :: BlockId -> StgBinding -> FCode ()
+cgLneBinds :: BlockId -> CgStgBinding -> FCode ()
 cgLneBinds join_id (StgNonRec bndr rhs)
   = do  { local_cc <- saveCurrentCostCentre
                 -- See Note [Saving the current cost centre]
@@ -135,7 +135,7 @@ cgLetNoEscapeRhs
     :: BlockId          -- join point for successor of let-no-escape
     -> Maybe LocalReg   -- Saved cost centre
     -> Id
-    -> StgRhs
+    -> CgStgRhs
     -> FCode (CgIdInfo, FCode ())
 
 cgLetNoEscapeRhs join_id local_cc bndr rhs =
@@ -149,9 +149,9 @@ cgLetNoEscapeRhs join_id local_cc bndr rhs =
 cgLetNoEscapeRhsBody
     :: Maybe LocalReg   -- Saved cost centre
     -> Id
-    -> StgRhs
+    -> CgStgRhs
     -> FCode (CgIdInfo, FCode ())
-cgLetNoEscapeRhsBody local_cc bndr (StgRhsClosure cc _ _upd args body)
+cgLetNoEscapeRhsBody local_cc bndr (StgRhsClosure _ cc _upd args body)
   = cgLetNoEscapeClosure bndr local_cc cc (nonVoidIds args) body
 cgLetNoEscapeRhsBody local_cc bndr (StgRhsCon cc con args)
   = cgLetNoEscapeClosure bndr local_cc cc []
@@ -168,7 +168,7 @@ cgLetNoEscapeClosure
         -> Maybe LocalReg       -- Slot for saved current cost centre
         -> CostCentreStack      -- XXX: *** NOT USED *** why not?
         -> [NonVoid Id]         -- Args (as in \ args -> body)
-        -> StgExpr              -- Body (as in above)
+        -> CgStgExpr            -- Body (as in above)
         -> FCode (CgIdInfo, FCode ())
 
 cgLetNoEscapeClosure bndr cc_slot _unused_cc args body
@@ -298,7 +298,7 @@ data GcPlan
                         -- of the case alternative(s) into the upstream check
 
 -------------------------------------
-cgCase :: StgExpr -> Id -> AltType -> [StgAlt] -> FCode ReturnKind
+cgCase :: CgStgExpr -> Id -> AltType -> [CgStgAlt] -> FCode ReturnKind
 
 cgCase (StgOpApp (StgPrimOp op) args _) bndr (AlgAlt tycon) alts
   | isEnumerationTyCon tycon -- Note [case on bool]
@@ -547,7 +547,7 @@ maybeSaveCostCentre simple_scrut
 
 
 -----------------
-isSimpleScrut :: StgExpr -> AltType -> FCode Bool
+isSimpleScrut :: CgStgExpr -> AltType -> FCode Bool
 -- Simple scrutinee, does not block or allocate; hence safe to amalgamate
 -- heap usage from alternatives into the stuff before the case
 -- NB: if you get this wrong, and claim that the expression doesn't allocate
@@ -570,7 +570,7 @@ isSimpleOp (StgPrimOp op) stg_args                  = do
 isSimpleOp (StgPrimCallOp _) _                           = return False
 
 -----------------
-chooseReturnBndrs :: Id -> AltType -> [StgAlt] -> [NonVoid Id]
+chooseReturnBndrs :: Id -> AltType -> [CgStgAlt] -> [NonVoid Id]
 -- These are the binders of a case that are assigned by the evaluation of the
 -- scrutinee.
 -- They're non-void, see Note [Post-unarisation invariants] in UnariseStg.
@@ -591,7 +591,7 @@ chooseReturnBndrs _ _ _ = panic "chooseReturnBndrs"
                              -- MultiValAlt has only one alternative
 
 -------------------------------------
-cgAlts :: (GcPlan,ReturnKind) -> NonVoid Id -> AltType -> [StgAlt]
+cgAlts :: (GcPlan,ReturnKind) -> NonVoid Id -> AltType -> [CgStgAlt]
        -> FCode ReturnKind
 -- At this point the result of the case are in the binders
 cgAlts gc_plan _bndr PolyAlt [(_, _, rhs)]
@@ -666,7 +666,7 @@ cgAlts _ _ _ _ = panic "cgAlts"
 --   goto L1
 
 -------------------
-cgAlgAltRhss :: (GcPlan,ReturnKind) -> NonVoid Id -> [StgAlt]
+cgAlgAltRhss :: (GcPlan,ReturnKind) -> NonVoid Id -> [CgStgAlt]
              -> FCode ( Maybe CmmAGraphScoped
                       , [(ConTagZ, CmmAGraphScoped)] )
 cgAlgAltRhss gc_plan bndr alts
@@ -686,13 +686,13 @@ cgAlgAltRhss gc_plan bndr alts
 
 
 -------------------
-cgAltRhss :: (GcPlan,ReturnKind) -> NonVoid Id -> [StgAlt]
+cgAltRhss :: (GcPlan,ReturnKind) -> NonVoid Id -> [CgStgAlt]
           -> FCode [(AltCon, CmmAGraphScoped)]
 cgAltRhss gc_plan bndr alts = do
   dflags <- getDynFlags
   let
     base_reg = idToReg dflags bndr
-    cg_alt :: StgAlt -> FCode (AltCon, CmmAGraphScoped)
+    cg_alt :: CgStgAlt -> FCode (AltCon, CmmAGraphScoped)
     cg_alt (con, bndrs, rhs)
       = getCodeScoped             $
         maybeAltHeapCheck gc_plan $