Support polymorphic beta-reduction while inlining

Author: nicolasstucki

compiler/src/dotty/tools/dotc/inlines/InlineReducer.scala

@@ -12,8 +12,6 @@ import NameKinds.{InlineAccessorName, InlineBinderName, InlineScrutineeName}
 import config.Printers.inlining
 import util.SimpleIdentityMap
 
-import dotty.tools.dotc.transform.BetaReduce
-
 import collection.mutable
 
 /** A utility class offering methods for rewriting inlined code */
@@ -150,45 +148,6 @@ class InlineReducer(inliner: Inliner)(using Context):
     binding1.withSpan(call.span)
   }
 
-  /** Rewrite an application
-    *
-    *    ((x1, ..., xn) => b)(e1, ..., en)
-    *
-    *  to
-    *
-    *    val/def x1 = e1; ...; val/def xn = en; b
-    *
-    *  where `def` is used for call-by-name parameters. However, we shortcut any NoPrefix
-    *  refs among the ei's directly without creating an intermediate binding.
-    *
-    *  This variant of beta-reduction preserves the integrity of `Inlined` tree nodes.
-    */
-  def betaReduce(tree: Tree)(using Context): Tree = tree match {
-    case Apply(Select(cl, nme.apply), args) if defn.isFunctionType(cl.tpe) =>
-      val bindingsBuf = new mutable.ListBuffer[DefTree]
-      def recur(cl: Tree): Option[Tree] = cl match
-        case Block((ddef : DefDef) :: Nil, closure: Closure) if ddef.symbol == closure.meth.symbol =>
-          ddef.tpe.widen match
-            case mt: MethodType if ddef.paramss.head.length == args.length =>
-              // TODO beta reduce PolyType
-              Some(BetaReduce.reduceApplication(ddef, List(args), bindingsBuf))
-            case _ => None
-        case Block(stats, expr) if stats.forall(isPureBinding) =>
-          recur(expr).map(cpy.Block(cl)(stats, _))
-        case Inlined(call, bindings, expr) if bindings.forall(isPureBinding) =>
-          recur(expr).map(cpy.Inlined(cl)(call, bindings, _))
-        case Typed(expr, tpt) =>
-          recur(expr)
-        case _ => None
-      recur(cl) match
-        case Some(reduced) =>
-          seq(bindingsBuf.result(), reduced).withSpan(tree.span)
-        case None =>
-          tree
-    case _ =>
-      tree
-  }
-
   /** The result type of reducing a match. It consists optionally of a list of bindings
    *  for the pattern-bound variables and the RHS of the selected case.
    *  Returns `None` if no case was selected.

compiler/src/dotty/tools/dotc/inlines/Inliner.scala

@@ -21,6 +21,7 @@ import collection.mutable
 import reporting.trace
 import util.Spans.Span
 import dotty.tools.dotc.transform.Splicer
+import dotty.tools.dotc.transform.BetaReduce
 import quoted.QuoteUtils
 import scala.annotation.constructorOnly
 
@@ -811,7 +812,7 @@ class Inliner(val call: tpd.Tree)(using Context):
           case Quoted(Spliced(inner)) => inner
           case _ => tree
       val locked = ctx.typerState.ownedVars
-      val res = cancelQuotes(constToLiteral(betaReduce(super.typedApply(tree, pt)))) match {
+      val res = cancelQuotes(constToLiteral(BetaReduce(super.typedApply(tree, pt)))) match {
         case res: Apply if res.symbol == defn.QuotedRuntime_exprSplice
                         && StagingContext.level == 0
                         && !hasInliningErrors =>
@@ -824,7 +825,7 @@ class Inliner(val call: tpd.Tree)(using Context):
 
     override def typedTypeApply(tree: untpd.TypeApply, pt: Type)(using Context): Tree =
       val locked = ctx.typerState.ownedVars
-      val tree1 = inlineIfNeeded(constToLiteral(betaReduce(super.typedTypeApply(tree, pt))), pt, locked)
+      val tree1 = inlineIfNeeded(constToLiteral(BetaReduce(super.typedTypeApply(tree, pt))), pt, locked)
       if tree1.symbol.isQuote then
         ctx.compilationUnit.needsStaging = true
       tree1
@@ -1005,7 +1006,7 @@ class Inliner(val call: tpd.Tree)(using Context):
             super.transform(t1)
           case t: Apply =>
             val t1 = super.transform(t)
-            if (t1 `eq` t) t else reducer.betaReduce(t1)
+            if (t1 `eq` t) t else BetaReduce(t1)
           case Block(Nil, expr) =>
             super.transform(expr)
           case _ =>

compiler/src/dotty/tools/dotc/transform/BetaReduce.scala

@@ -56,6 +56,68 @@ object BetaReduce:
   val name: String = "betaReduce"
   val description: String = "reduce closure applications"
 
+  /** Rewrite an application
+   *
+   *    ((x1, ..., xn) => b)(e1, ..., en)
+   *
+   *  to
+   *
+   *    val/def x1 = e1; ...; val/def xn = en; b
+   *
+   *  where `def` is used for call-by-name parameters. However, we shortcut any NoPrefix
+   *  refs among the ei's directly without creating an intermediate binding.
+   *
+   *  Similarly, rewrites type applications
+   *
+   *    ([X1, ..., Xm] => (x1, ..., xn) => b).apply[T1, .., Tm](e1, ..., en)
+   *
+   *  to
+   *
+   *    type X1 = T1; ...; type Xm = Tm;val/def x1 = e1; ...; val/def xn = en; b
+   *
+   *  This beta-reduction preserves the integrity of `Inlined` tree nodes.
+   */
+  def apply(tree: Tree)(using Context): Tree =
+    val bindingsBuf = new ListBuffer[DefTree]
+    def recur(fn: Tree, argss: List[List[Tree]]): Option[Tree] = fn match
+      case Block((ddef : DefDef) :: Nil, closure: Closure) if ddef.symbol == closure.meth.symbol =>
+        ddef.tpe.widen match // TODO can these guards be removed?
+          case mt: MethodType if ddef.paramss.head.length == argss.head.length =>
+            Some(reduceApplication(ddef, argss, bindingsBuf))
+          case _ => None
+      case Block((TypeDef(_, template: Template)) :: Nil, Typed(Apply(Select(New(_), _), _), _)) if template.constr.rhs.isEmpty =>
+        template.body match
+          case (ddef: DefDef) :: Nil =>
+            ddef.tpe.widen match // TODO can these guards be removed?
+              case mt: MethodType if ddef.paramss.head.length == argss.head.length =>
+                Some(reduceApplication(ddef, argss, bindingsBuf))
+              case mt: PolyType if ddef.paramss.head.length == argss.head.length && ddef.paramss.last.length == argss.last.length =>
+                Some(reduceApplication(ddef, argss, bindingsBuf))
+              case _ => None
+          case _ => None
+      case Block(stats, expr) if stats.forall(isPureBinding) =>
+        recur(expr, argss).map(cpy.Block(fn)(stats, _))
+      case Inlined(call, bindings, expr) if bindings.forall(isPureBinding) =>
+        recur(expr, argss).map(cpy.Inlined(fn)(call, bindings, _))
+      case Typed(expr, tpt) =>
+        recur(expr, argss)
+      case _ => None
+    tree match
+      case Apply(Select(fn, nme.apply), args) if defn.isFunctionType(fn.tpe) =>
+        recur(fn, List(args)) match
+          case Some(reduced) =>
+            seq(bindingsBuf.result(), reduced).withSpan(tree.span)
+          case None =>
+            tree
+      case Apply(TypeApply(Select(fn, nme.apply), targs), args) if fn.tpe.typeSymbol eq dotc.core.Symbols.defn.PolyFunctionClass =>
+        recur(fn, List(targs, args)) match
+          case Some(reduced) =>
+            seq(bindingsBuf.result(), reduced).withSpan(tree.span)
+          case None =>
+            tree
+      case _ =>
+        tree
+
   /** Beta-reduces a call to `fn` with arguments `argSyms` or returns `tree` */
   def apply(original: Tree, fn: Tree, argss: List[List[Tree]])(using Context): Tree =
     fn match

tests/run-macros/inline-beta-reduce-polyfunction.check

@@ -0,0 +1,7 @@
+{
+  type X = Int
+  {
+    println(1)
+    1
+  }
+}

tests/run-macros/inline-beta-reduce-polyfunction.scala

@@ -0,0 +1,5 @@
+transparent inline def foo(inline f: [X] => X => X): Int = f[Int](1)
+
+@main def Test: Unit =
+  val code = compiletime.codeOf(foo([X] => (x: X) => { println(x); x }))
+  println(code)