[commit: ghc] wip/T14068: Implement loopification for local bindings (#14068) (29a03ad)

[git at git.haskell.org]

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

On branch  : wip/T14068
Link       : http://ghc.haskell.org/trac/ghc/changeset/29a03ad7cc5734033bb1ccca1ca8c56a48ffeef3/ghc

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

commit 29a03ad7cc5734033bb1ccca1ca8c56a48ffeef3
Author: Joachim Breitner <mail at joachim-breitner.de>
Date:   Tue Aug 1 10:36:32 2017 -0400

    Implement loopification for local bindings (#14068)
    
    This is a relatively prelimary version. I am sure there is a huge number
    of invariants that this breaks, and conditions that I am not checking
    etc. I do not even know if the simplifier is the right place to
    implement this.
    
    But it works in this simple case:
    
        module T14068 where
        foo p f k =
          let bar a = if p a then bar (f a) else a
          in k bar
    
    so we can iterate from here.


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

29a03ad7cc5734033bb1ccca1ca8c56a48ffeef3
 compiler/coreSyn/CoreOpt.hs    | 34 +++++++++++++++++++++++++++-------
 compiler/simplCore/Simplify.hs | 17 ++++++++++++++++-
 2 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/compiler/coreSyn/CoreOpt.hs b/compiler/coreSyn/CoreOpt.hs
index 4a19605..0dae086 100644
--- a/compiler/coreSyn/CoreOpt.hs
+++ b/compiler/coreSyn/CoreOpt.hs
@@ -10,6 +10,7 @@ module CoreOpt (
 
         -- ** Join points
         joinPointBinding_maybe, joinPointBindings_maybe,
+        loopificationJoinPointBinding_maybe ,
 
         -- ** Predicates on expressions
         exprIsConApp_maybe, exprIsLiteral_maybe, exprIsLambda_maybe,
@@ -642,22 +643,41 @@ joinPointBinding_maybe bndr rhs
   = Just (bndr, rhs)
 
   | AlwaysTailCalled join_arity <- tailCallInfo (idOccInfo bndr)
-  , not (bad_unfolding join_arity (idUnfolding bndr))
+  , not (badUnfoldingForJoin join_arity bndr)
   , (bndrs, body) <- etaExpandToJoinPoint join_arity rhs
   = Just (bndr `asJoinId` join_arity, mkLams bndrs body)
 
   | otherwise
   = Nothing
 
+-- | like joinPointBinding_maybe, but looks for RecursiveTailCalled
+loopificationJoinPointBinding_maybe :: InBndr -> InExpr -> Maybe (InBndr, InExpr)
+loopificationJoinPointBinding_maybe bndr rhs
+  | not (isId bndr)
+  = Nothing
+
+  | isJoinId bndr
+  = Nothing -- do not loopificate again
+
+  | RecursiveTailCalled join_arity <- tailCallInfo (idOccInfo bndr)
+  , not (badUnfoldingForJoin join_arity bndr)
+  , (bndrs, body) <- etaExpandToJoinPoint join_arity rhs
+  = Just (bndr `asJoinId` join_arity, mkLams bndrs body)
+
+  | otherwise
+  = Nothing
+
+-- | badUnfoldingForJoin returns True if we should /not/ convert a non-join-id
+--   into a join-id, even though it is AlwaysTailCalled
+--   See Note [Join points and INLINE pragmas]
+badUnfoldingForJoin :: JoinArity -> Id -> Bool
+badUnfoldingForJoin join_arity bndr = bad_unfolding (idUnfolding bndr)
   where
-    -- bad_unfolding returns True if we should /not/ convert a non-join-id
-    -- into a join-id, even though it is AlwaysTailCalled
-    -- See Note [Join points and INLINE pragmas]
-    bad_unfolding join_arity (CoreUnfolding { uf_src = src, uf_tmpl = rhs })
+    bad_unfolding (CoreUnfolding { uf_src = src, uf_tmpl = rhs })
       = isStableSource src && join_arity > joinRhsArity rhs
-    bad_unfolding _ (DFunUnfolding {})
+    bad_unfolding (DFunUnfolding {})
       = True
-    bad_unfolding _ _
+    bad_unfolding _
       = False
 
 joinPointBindings_maybe :: [(InBndr, InExpr)] -> Maybe [(InBndr, InExpr)]
diff --git a/compiler/simplCore/Simplify.hs b/compiler/simplCore/Simplify.hs
index 1fc9112..51b93ff 100644
--- a/compiler/simplCore/Simplify.hs
+++ b/compiler/simplCore/Simplify.hs
@@ -36,7 +36,8 @@ import CoreUnfold
 import CoreUtils
 import CoreArity
 import CoreOpt          ( pushCoTyArg, pushCoValArg
-                        , joinPointBinding_maybe, joinPointBindings_maybe )
+                        , joinPointBinding_maybe, joinPointBindings_maybe
+                        , loopificationJoinPointBinding_maybe )
 --import PrimOp           ( tagToEnumKey ) -- temporalily commented out. See #8326
 import Rules            ( mkRuleInfo, lookupRule, getRules )
 --import TysPrim          ( intPrimTy ) -- temporalily commented out. See #8326
@@ -1637,6 +1638,7 @@ simplRecE :: SimplEnv
 -- simplRecE is used for
 --  * non-top-level recursive lets in expressions
 simplRecE env pairs body cont
+
   | Just pairs' <- joinPointBindings_maybe pairs
   = do  { (env1, cont') <- prepareJoinCont env cont
         ; let bndrs' = map fst pairs'
@@ -1647,6 +1649,19 @@ simplRecE env pairs body cont
         ; env3 <- simplRecBind env2 NotTopLevel (Just cont') pairs'
         ; simplExprF env3 body cont' }
 
+  -- Is this a tail-recursive function that we want to loopify? Then
+  -- lets loopify it and re-analyse.
+  | [(bndr,rhs)] <- pairs
+  , Just (join_bndr, join_rhs) <- loopificationJoinPointBinding_maybe bndr rhs
+  , let Just arity = isJoinId_maybe join_bndr
+  = do  { let (join_params, _join_body) = collectNBinders arity join_rhs
+        ; let bndr' = zapFragileIdInfo bndr -- TODO: What do we have to zap here?
+        ; let rhs' = mkLams join_params $
+                     mkLetRec [(join_bndr,join_rhs)] $
+                     mkVarApps (Var join_bndr) join_params
+        ; simplNonRecE env bndr' (rhs', env) ([], body) cont
+        }
+
   | otherwise
   = do  { let bndrs = map fst pairs
         ; MASSERT(all (not . isJoinId) bndrs)