[commit: ghc] wip/T14152: Do some ad-hoc CSE in Exitification (154fc88)

[git at git.haskell.org]

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

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

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

commit 154fc8861eb12656ba1bf56ea257ac4483d63d21
Author: Joachim Breitner <mail at joachim-breitner.de>
Date:   Wed Sep 6 15:17:12 2017 +0100

    Do some ad-hoc CSE in Exitification
    
    So that code like this
    
        safe :: Int -> [Int] -> [[Int]]
        safe x2 ds6
          = case ds6 of
              [] -> lvl6;
              q : l | x2     == q -> h x2 ds5
                    | x2 + 1 == q -> h x2 ds5
                    | x2 + 2 == q -> h x2 ds5
                    | x2 + 3 == q -> h x2 ds5
                    | otherwise   -> safe (x+10) l
    
    only gets one exit join point, instead of 4 identically.
    
    (There is currently no CSE for join points, and even if there were, it is cheap
    to do it here.)


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

154fc8861eb12656ba1bf56ea257ac4483d63d21
 compiler/simplCore/Exitify.hs | 37 +++++++++++++++++++++++++------------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/compiler/simplCore/Exitify.hs b/compiler/simplCore/Exitify.hs
index 916ae73..c1b64b1 100644
--- a/compiler/simplCore/Exitify.hs
+++ b/compiler/simplCore/Exitify.hs
@@ -11,6 +11,7 @@ import VarSet
 import VarEnv
 import CoreFVs
 import FastString
+import TrieMap
 import Type
 
 import Data.Bifunctor
@@ -71,13 +72,14 @@ exitify in_scope pairs =
     -- Which are the recursive calls?
     recursive_calls = mkVarSet $ map fst pairs
 
-    (pairs',exits) = (`runState` []) $ do
+    (pairs',st) = (`runState` ExitifyState [] emptyTM) $ do
         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
@@ -87,7 +89,7 @@ exitify in_scope pairs =
 
     go :: [Var]           -- ^ variables to abstract over
        -> CoreExprWithFVs -- ^ current expression in tail position
-       -> State [(Id, CoreExpr)] CoreExpr
+       -> ExitifyM CoreExpr
 
     go captured ann_e
         -- Do not touch an expression that is already a join call with no free variables
@@ -175,8 +177,8 @@ exitify in_scope pairs =
 --  * any exit join points created so far.
 mkExitJoinId :: InScopeSet -> Type -> JoinArity -> ExitifyM JoinId
 mkExitJoinId in_scope ty join_arity = do
-    fs <- get
-    let avoid = in_scope `extendInScopeSetList` (map fst fs)
+    st <- get
+    let avoid = in_scope `extendInScopeSetList` (map fst (es_joins st))
                          `extendInScopeSet` exit_id_tmpl -- just cosmetics
     return (uniqAway avoid exit_id_tmpl)
   where
@@ -190,13 +192,24 @@ mkExitJoinId in_scope ty join_arity = do
                , occ_int_cxt = False
                , occ_tail = AlwaysTailCalled join_arity }
 
+-- 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
-    -- Pick a suitable name
-    v <- mkExitJoinId in_scope ty join_arity
-    fs <- get
-    put ((v,rhs):fs)
-    return v
-
-
-type ExitifyM =  State [(JoinId, CoreExpr)]
+    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
+            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