[commit: ghc] wip/T14152: Get in_scope_set right for CSE’ed exitification (bfa9c26)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/ghc

On branch  : wip/T14152
Link       : http://ghc.haskell.org/trac/ghc/changeset/bfa9c262a87480e1852a16422d0b9f5a2601db5b/ghc

>---------------------------------------------------------------

commit bfa9c262a87480e1852a16422d0b9f5a2601db5b
Author: Joachim Breitner <mail at joachim-breitner.de>
Date:   Thu Sep 7 00:23:13 2017 +0100

    Get in_scope_set right for CSE’ed exitification
    
    (this patch will be squashed with the CSE patch)


>---------------------------------------------------------------

bfa9c262a87480e1852a16422d0b9f5a2601db5b
 compiler/simplCore/Exitify.hs | 58 ++++++++++++++++++++++++++++---------------
 1 file changed, 38 insertions(+), 20 deletions(-)

diff --git a/compiler/simplCore/Exitify.hs b/compiler/simplCore/Exitify.hs
index c1b64b1..b9973e2 100644
--- a/compiler/simplCore/Exitify.hs
+++ b/compiler/simplCore/Exitify.hs
@@ -60,7 +60,7 @@ exitifyProgram binds = map goTopLvl binds
 -- join-points outside the joinrec.
 exitify :: InScopeSet -> [(Var,CoreExpr)] -> (CoreExpr -> CoreExpr)
 exitify in_scope pairs =
-    \body ->mkExitLets exits (mkLetRec pairs' body)
+    \body -> mkExitLets exits (mkLetRec pairs' body)
   where
     mkExitLets ((exitId, exitRhs):exits') = mkLetNonRec exitId exitRhs . mkExitLets exits'
     mkExitLets [] = id
@@ -72,14 +72,13 @@ exitify in_scope pairs =
     -- Which are the recursive calls?
     recursive_calls = mkVarSet $ map fst pairs
 
-    (pairs',st) = (`runState` ExitifyState [] emptyTM) $ do
+    (pairs',exits) = runExitifyState in_scope $
         forM ann_pairs $ \(x,rhs) -> do
             -- go past the lambdas of the join point
             let (args, body) = collectNAnnBndrs (idJoinArity x) rhs
             body' <- go args body
             let rhs' = mkLams args body'
             return (x, rhs')
-    exits = es_joins st
 
     -- main working function. Goes through the RHS (tail-call positions only),
     -- checks if there are no more recursive calls, if so, abstracts over
@@ -109,10 +108,10 @@ exitify in_scope pairs =
         | is_exit = do
             -- Assemble the RHS of the exit join point
             let rhs = mkLams args e
-                ty = exprType rhs
-            let avoid = in_scope `extendInScopeSetList` captured
+            -- Remember what is in scope here
+            nowInScope captured
             -- Remember this binding under a suitable name
-            v <- addExit avoid ty (length args) rhs
+            v <- addExit (length args) rhs
             -- And call it from here
             return $ mkVarApps (Var v) args
       where
@@ -171,16 +170,41 @@ exitify in_scope pairs =
     go _ ann_e = return (deAnnotate ann_e)
 
 
+type ExitifyM = State ExitifyState
+data ExitifyState = ExitifyState
+        { es_in_scope_acc :: InScopeSet -- ^ combined in_scope_set of all call sites
+        , es_in_scope     :: InScopeSet -- ^ final in_scope_set
+        , es_joins        :: [(JoinId, CoreExpr)] -- ^ exit join points
+        , es_map          :: CoreMap JoinId
+            -- ^ reverse lookup map, see Note [Avoid duplicate exit points]
+        }
+
+-- Runs the ExitifyM monad, and feeds in the final es_in_scope_acc as the
+-- es_in_scope to use
+runExitifyState :: InScopeSet -> ExitifyM a -> (a, [(JoinId, CoreExpr)])
+runExitifyState in_scope_init f = (res, es_joins state)
+  where
+    (res, state) = runState f (ExitifyState in_scope_init in_scope [] emptyTM)
+    in_scope = es_in_scope_acc state
+
+-- Keeps track of what is in scope at all the various positions where
+-- we want to call an exit join point
+nowInScope :: [Var] -> ExitifyM ()
+nowInScope captured = do
+    st <- get
+    put (st { es_in_scope_acc = es_in_scope_acc st `extendInScopeSetList` captured})
+
 -- Picks a new unique, which is disjoint from
 --  * the free variables of the whole joinrec
 --  * any bound variables (captured)
 --  * any exit join points created so far.
-mkExitJoinId :: InScopeSet -> Type -> JoinArity -> ExitifyM JoinId
-mkExitJoinId in_scope ty join_arity = do
+mkExitJoinId :: Type -> JoinArity -> ExitifyM JoinId
+mkExitJoinId ty join_arity = do
     st <- get
-    let avoid = in_scope `extendInScopeSetList` (map fst (es_joins st))
-                         `extendInScopeSet` exit_id_tmpl -- just cosmetics
-    return (uniqAway avoid exit_id_tmpl)
+    let in_scope = es_in_scope st `extendInScopeSet` exit_id_tmpl -- cosmetic only
+    let v = uniqAway in_scope exit_id_tmpl
+    put (st { es_in_scope = es_in_scope st `extendInScopeSet` v})
+    return v
   where
     exit_id_tmpl = mkSysLocal (fsLit "exit") initExitJoinUnique ty
                     `asJoinId` join_arity
@@ -194,22 +218,16 @@ mkExitJoinId in_scope ty join_arity = do
 
 -- Adds a new exit join point
 -- (or re-uses an existing one)
-addExit :: InScopeSet -> Type -> JoinArity -> CoreExpr -> ExitifyM JoinId
-addExit in_scope ty join_arity rhs = do
+addExit :: JoinArity -> CoreExpr -> ExitifyM JoinId
+addExit join_arity rhs = do
     st <- get
     case lookupTM rhs (es_map st) of
         Just v -> return v
         Nothing -> do
             -- Pick a suitable name
-            v <- mkExitJoinId in_scope ty join_arity
+            v <- mkExitJoinId (exprType rhs) join_arity
             st <- get
             put (st { es_joins = (v,rhs) : es_joins st
                     , es_map = insertTM rhs v (es_map st)
                     })
             return v
-
-data ExitifyState = ExitifyState { es_joins :: [(JoinId, CoreExpr)]
-                                 , es_map :: CoreMap JoinId
-                                 }
-
-type ExitifyM = State ExitifyState