Implement late lambda lift

Summary:
This implements a selective lambda-lifting pass late in the STG
pipeline.

Lambda lifting has the effect of avoiding closure allocation at the cost
of having to make former free vars available at call sites, possibly
enlarging closures surrounding call sites in turn.

We identify beneficial cases by means of an analysis that estimates
closure growth.

There's a Wiki page at
https://ghc.haskell.org/trac/ghc/wiki/LateLamLift.

Reviewers: simonpj, bgamari, simonmar

Reviewed By: simonpj

Subscribers: rwbarton, carter

GHC Trac Issues: #9476

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

1 files changed, 28 insertions, 23 deletions

diff --git a/compiler/stgSyn/StgFVs.hs b/compiler/stgSyn/StgFVs.hs
index 80ce33ff7a..edfc94ed2d 100644
--- a/compiler/stgSyn/StgFVs.hs
+++ b/compiler/stgSyn/StgFVs.hs
@@ -1,6 +1,7 @@
 -- | Free variable analysis on STG terms.
 module StgFVs (
-    annTopBindingsFreeVars
+    annTopBindingsFreeVars,
+    annBindingFreeVars
   ) where
 
 import GhcPrelude
@@ -26,13 +27,17 @@ addLocals :: [Id] -> Env -> Env
 addLocals bndrs env
   = env { locals = extendVarSetList (locals env) bndrs }
 
--- | Annotates a top-level STG binding with its free variables.
+-- | Annotates a top-level STG binding group with its free variables.
 annTopBindingsFreeVars :: [StgTopBinding] -> [CgStgTopBinding]
 annTopBindingsFreeVars = map go
   where
     go (StgTopStringLit id bs) = StgTopStringLit id bs
     go (StgTopLifted bind)
-      = StgTopLifted (fst (binding emptyEnv emptyVarSet bind))
+      = StgTopLifted (annBindingFreeVars bind)
+
+-- | Annotates an STG binding with its free variables.
+annBindingFreeVars :: StgBinding -> CgStgBinding
+annBindingFreeVars = fst . binding emptyEnv emptyDVarSet
 
 boundIds :: StgBinding -> [Id]
 boundIds (StgNonRec b _) = [b]
@@ -53,35 +58,35 @@ boundIds (StgRec pairs)  = map fst pairs
 --      knot-tying.
 
 -- | This makes sure that only local, non-global free vars make it into the set.
-mkFreeVarSet :: Env -> [Id] -> IdSet
-mkFreeVarSet env = mkVarSet . filter (`elemVarSet` locals env)
+mkFreeVarSet :: Env -> [Id] -> DIdSet
+mkFreeVarSet env = mkDVarSet . filter (`elemVarSet` locals env)
 
-args :: Env -> [StgArg] -> IdSet
+args :: Env -> [StgArg] -> DIdSet
 args env = mkFreeVarSet env . mapMaybe f
   where
     f (StgVarArg occ) = Just occ
     f _               = Nothing
 
-binding :: Env -> IdSet -> StgBinding -> (CgStgBinding, IdSet)
+binding :: Env -> DIdSet -> StgBinding -> (CgStgBinding, DIdSet)
 binding env body_fv (StgNonRec bndr r) = (StgNonRec bndr r', fvs)
   where
     -- See Note [Tacking local binders]
     (r', rhs_fvs) = rhs env r
-    fvs = delVarSet body_fv bndr `unionVarSet` rhs_fvs
+    fvs = delDVarSet body_fv bndr `unionDVarSet` rhs_fvs
 binding env body_fv (StgRec pairs) = (StgRec pairs', fvs)
   where
     -- See Note [Tacking local binders]
     bndrs = map fst pairs
     (rhss, rhs_fvss) = mapAndUnzip (rhs env . snd) pairs
     pairs' = zip bndrs rhss
-    fvs = delVarSetList (unionVarSets (body_fv:rhs_fvss)) bndrs
+    fvs = delDVarSetList (unionDVarSets (body_fv:rhs_fvss)) bndrs
 
-expr :: Env -> StgExpr -> (CgStgExpr, IdSet)
+expr :: Env -> StgExpr -> (CgStgExpr, DIdSet)
 expr env = go
   where
     go (StgApp occ as)
-      = (StgApp occ as, unionVarSet (args env as) (mkFreeVarSet env [occ]))
-    go (StgLit lit) = (StgLit lit, emptyVarSet)
+      = (StgApp occ as, unionDVarSet (args env as) (mkFreeVarSet env [occ]))
+    go (StgLit lit) = (StgLit lit, emptyDVarSet)
     go (StgConApp dc as tys) = (StgConApp dc as tys, args env as)
     go (StgOpApp op as ty) = (StgOpApp op as ty, args env as)
     go StgLam{} = pprPanic "StgFVs: StgLam" empty
@@ -90,16 +95,16 @@ expr env = go
         (scrut', scrut_fvs) = go scrut
         -- See Note [Tacking local binders]
         (alts', alt_fvss) = mapAndUnzip (alt (addLocals [bndr] env)) alts
-        alt_fvs = unionVarSets alt_fvss
-        fvs = delVarSet (unionVarSet scrut_fvs alt_fvs) bndr
-    go (StgLet bind body) = go_bind StgLet bind body
-    go (StgLetNoEscape bind body) = go_bind StgLetNoEscape bind body
+        alt_fvs = unionDVarSets alt_fvss
+        fvs = delDVarSet (unionDVarSet scrut_fvs alt_fvs) bndr
+    go (StgLet ext bind body) = go_bind (StgLet ext) bind body
+    go (StgLetNoEscape ext bind body) = go_bind (StgLetNoEscape ext) bind body
     go (StgTick tick e) = (StgTick tick e', fvs')
       where
         (e', fvs) = go e
-        fvs' = unionVarSet (tickish tick) fvs
-        tickish (Breakpoint _ ids) = mkVarSet ids
-        tickish _                  = emptyVarSet
+        fvs' = unionDVarSet (tickish tick) fvs
+        tickish (Breakpoint _ ids) = mkDVarSet ids
+        tickish _                  = emptyDVarSet
 
     go_bind dc bind body = (dc bind' body', fvs)
       where
@@ -108,18 +113,18 @@ expr env = go
         (body', body_fvs) = expr env' body
         (bind', fvs) = binding env' body_fvs bind
 
-rhs :: Env -> StgRhs -> (CgStgRhs, IdSet)
+rhs :: Env -> StgRhs -> (CgStgRhs, DIdSet)
 rhs env (StgRhsClosure _ ccs uf bndrs body)
   = (StgRhsClosure fvs ccs uf bndrs body', fvs)
   where
     -- See Note [Tacking local binders]
     (body', body_fvs) = expr (addLocals bndrs env) body
-    fvs = delVarSetList body_fvs bndrs
+    fvs = delDVarSetList body_fvs bndrs
 rhs env (StgRhsCon ccs dc as) = (StgRhsCon ccs dc as, args env as)
 
-alt :: Env -> StgAlt -> (CgStgAlt, IdSet)
+alt :: Env -> StgAlt -> (CgStgAlt, DIdSet)
 alt env (con, bndrs, e) = ((con, bndrs, e'), fvs)
   where
     -- See Note [Tacking local binders]
     (e', rhs_fvs) = expr (addLocals bndrs env) e
-    fvs = delVarSetList rhs_fvs bndrs
+    fvs = delDVarSetList rhs_fvs bndrs