[commit: ghc] wip/T14152: Exitify: The rhs of join points are also in tail-call position (4c3157d)

[git at git.haskell.org]

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

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

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

commit 4c3157d2c31ce379141575d83a1eb97fedfac128
Author: Joachim Breitner <mail at joachim-breitner.de>
Date:   Sun Sep 3 09:26:44 2017 +0100

    Exitify: The rhs of join points are also in tail-call position


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

4c3157d2c31ce379141575d83a1eb97fedfac128
 compiler/simplCore/Exitify.hs | 36 ++++++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/compiler/simplCore/Exitify.hs b/compiler/simplCore/Exitify.hs
index de29170..eb882e52 100644
--- a/compiler/simplCore/Exitify.hs
+++ b/compiler/simplCore/Exitify.hs
@@ -144,12 +144,36 @@ exitify in_scope pairs =
             return (dc, pats, rhs')
         return $ Case (deAnnotate scrut) bndr ty alts'
 
-    go captured (_, AnnLet ann_bind body) = do
-        let bind = deAnnBind ann_bind
-        body' <- go (bindersOf bind ++ captured) body
-        return $ Let bind body'
-
-    go _ ann_e = return (deAnnotate ann_e) -- TODO: What else is a tail-call position?
+    go captured (_, AnnLet ann_bind body)
+        -- join point, RHS and body are tail-called
+        | AnnNonRec j rhs <- ann_bind
+        , Just join_arity <- isJoinId_maybe j
+        = do let (params, join_body) = collectNAnnBndrs join_arity rhs
+             join_body' <- go (params ++ captured) join_body
+             let rhs' = mkLams params join_body'
+             body' <- go (j : captured) body
+             return $ Let (NonRec j rhs') body'
+
+        -- rec join point, RHSs and body are tail-called
+        | AnnRec pairs <- ann_bind
+        , isJoinId (fst (head pairs))
+        = do let js = map fst pairs
+             pairs' <- forM pairs $ \(j,rhs) -> do
+                 let join_arity = idJoinArity j
+                     (params, join_body) = collectNAnnBndrs join_arity rhs
+                 join_body' <- go (params ++ js ++ captured) join_body
+                 let rhs' = mkLams params join_body'
+                 return (j, rhs')
+             body' <- go (js ++ captured) body
+             return $ Let (Rec pairs') body'
+
+        -- normal Let, only the body is tail-called
+        | otherwise
+        = do body' <- go (bindersOf bind ++ captured) body
+             return $ Let bind body'
+      where bind = deAnnBind ann_bind
+
+    go _ ann_e = return (deAnnotate ann_e)
 
 
     -- Picks a new unique, which is disjoint from