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

[git at git.haskell.org]

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

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

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

commit b727e282d2332e43dcb73237dcb379aad4ffc2e7
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


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

b727e282d2332e43dcb73237dcb379aad4ffc2e7
 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 8913c3e..6fe17ac 100644
--- a/compiler/simplCore/Exitify.hs
+++ b/compiler/simplCore/Exitify.hs
@@ -135,12 +135,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