Inline argument closures to inline methods

If an argumnet to an inline method refers to a closure
that is the result of eta-expanding another inline method
inline the argument method.

3 files changed, 48 insertions, 7 deletions

diff --git a/src/dotty/tools/dotc/ast/Desugar.scala b/src/dotty/tools/dotc/ast/Desugar.scala
index 4e27da2ca..ecb6a3212 100644
--- a/src/dotty/tools/dotc/ast/Desugar.scala
+++ b/src/dotty/tools/dotc/ast/Desugar.scala
@@ -607,11 +607,17 @@ object desugar {
    *  ==>
    *      def $anonfun(params) = body
    *      Closure($anonfun)
+   *
+   *  If `inlineable` is true, tag $anonfun with an @inline annotation.
    */
-  def makeClosure(params: List[ValDef], body: Tree, tpt: Tree = TypeTree()) =
+  def makeClosure(params: List[ValDef], body: Tree, tpt: Tree = TypeTree(), inlineable: Boolean)(implicit ctx: Context) = {
+    var mods = synthetic
+    if (inlineable)
+      mods = mods.withAddedAnnotation(New(ref(defn.InlineAnnotType), Nil).withPos(body.pos))
     Block(
-      DefDef(nme.ANON_FUN, Nil, params :: Nil, tpt, body).withMods(synthetic),
+      DefDef(nme.ANON_FUN, Nil, params :: Nil, tpt, body).withMods(mods),
       Closure(Nil, Ident(nme.ANON_FUN), EmptyTree))
+  }
 
   /** If `nparams` == 1, expand partial function
    *
diff --git a/src/dotty/tools/dotc/typer/Inliner.scala b/src/dotty/tools/dotc/typer/Inliner.scala
index ca2daa4a1..546988d21 100644
--- a/src/dotty/tools/dotc/typer/Inliner.scala
+++ b/src/dotty/tools/dotc/typer/Inliner.scala
@@ -470,29 +470,53 @@ class Inliner(call: tpd.Tree, rhs: tpd.Tree)(implicit ctx: Context) {
     // the owner from the inlined method to the current owner.
     val inliner = new TreeTypeMap(typeMap, treeMap, meth :: Nil, ctx.owner :: Nil)
 
-    val bindings = bindingsBuf.toList.map(_.withPos(call.pos))
     val expansion = inliner(rhs.withPos(call.pos))
 
     // The final expansion runs a typing pass over the inlined tree. See InlineTyper for details.
     val expansion1 = InlineTyper.typed(expansion, pt)(inlineContext(call))
-    val result = tpd.Inlined(call, bindings, expansion1)
 
+    /** Does given definition bind a closure that will be inlined? */
+    def bindsInlineableClosure(defn: ValOrDefDef) = Ident(defn.symbol.termRef) match {
+      case InlineableClosure(_) => true
+      case _ => false
+    }
+
+    /** All bindings in `bindingsBuf` except bindings of inlineable closures */
+    val bindings = bindingsBuf.toList.filterNot(bindsInlineableClosure).map(_.withPos(call.pos))
+
+    val result = tpd.Inlined(call, bindings, expansion1)
     inlining.println(i"inlined $call\n --> \n$result")
     result
   }
 
+  /** An extractor for references to closure arguments that refer to `@inline` methods */
+  private object InlineableClosure {
+    lazy val paramProxies = paramProxy.values.toSet
+    def unapply(tree: Ident)(implicit ctx: Context): Option[Tree] =
+      if (paramProxies.contains(tree.tpe)) {
+        bindingsBuf.find(_.name == tree.name).get.rhs match {
+          case Closure(_, meth, _) if meth.symbol.isInlineMethod => Some(meth)
+          case Block(_, Closure(_, meth, _)) if meth.symbol.isInlineMethod => Some(meth)
+          case _ => None
+        }
+      } else None
+  }
+
   /** A typer for inlined code. Its purpose is:
    *  1. Implement constant folding over inlined code
    *  2. Selectively expand ifs with constant conditions
-   *  3. Make sure inlined code is type-correct.
-   *  4. Make sure that the tree's typing is idempotent (so that future -Ycheck passes succeed)
+   *  3. Inline arguments that are inlineable closures
+   *  4. Make sure inlined code is type-correct.
+   *  5. Make sure that the tree's typing is idempotent (so that future -Ycheck passes succeed)
    */
   private object InlineTyper extends ReTyper {
+
     override def typedSelect(tree: untpd.Select, pt: Type)(implicit ctx: Context): Tree = {
       val res = super.typedSelect(tree, pt)
       ensureAccessible(res.tpe, tree.qualifier.isInstanceOf[untpd.Super], tree.pos)
       res
     }
+
     override def typedIf(tree: untpd.If, pt: Type)(implicit ctx: Context) = {
       val cond1 = typed(tree.cond, defn.BooleanType)
       cond1.tpe.widenTermRefExpr match {
@@ -505,5 +529,12 @@ class Inliner(call: tpd.Tree, rhs: tpd.Tree)(implicit ctx: Context) {
           super.typedIf(if1, pt)
       }
     }
+
+    override def typedApply(tree: untpd.Apply, pt: Type)(implicit ctx: Context) = tree.asInstanceOf[tpd.Tree] match {
+      case Apply(Select(InlineableClosure(fn), nme.apply), args) =>
+        typed(fn.appliedToArgs(args), pt)
+      case _ =>
+        super.typedApply(tree, pt)
+    }
   }
 }
diff --git a/src/dotty/tools/dotc/typer/Typer.scala b/src/dotty/tools/dotc/typer/Typer.scala
index 8830537e3..c8362522d 100644
--- a/src/dotty/tools/dotc/typer/Typer.scala
+++ b/src/dotty/tools/dotc/typer/Typer.scala
@@ -744,7 +744,11 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
             case WildcardType(_) => untpd.TypeTree()
             case _ => untpd.TypeTree(protoResult)
           }
-          desugar.makeClosure(inferredParams, fnBody, resultTpt)
+          val inlineable = fnBody match {
+            case Apply(untpd.TypedSplice(fn), _) => Inliner.hasBodyToInline(fn.symbol)
+            case _ => false
+          }
+          desugar.makeClosure(inferredParams, fnBody, resultTpt, inlineable)
         }
       typed(desugared, pt)
     }