Beta-reduce directly applied PolymorphicFunction

nicolasstucki · nicolasstucki · commit e02c9e1086fb · 2023-01-05T17:56:29.000Z
diff --git a/compiler/src/dotty/tools/dotc/inlines/InlineReducer.scala b/compiler/src/dotty/tools/dotc/inlines/InlineReducer.scala
@@ -165,12 +165,13 @@ class InlineReducer(inliner: Inliner)(using Context):
     */
   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[ValDef]
+      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 =>
-              Some(BetaReduce.reduceApplication(ddef, args, bindingsBuf))
+              // 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, _))
diff --git a/compiler/src/dotty/tools/dotc/transform/BetaReduce.scala b/compiler/src/dotty/tools/dotc/transform/BetaReduce.scala
@@ -40,49 +40,74 @@ class BetaReduce extends MiniPhase:
 
   override def transformApply(app: Apply)(using Context): Tree = app.fun match
     case Select(fn, nme.apply) if defn.isFunctionType(fn.tpe) =>
-      val app1 = BetaReduce(app, fn, app.args)
+      val app1 = BetaReduce(app, fn, List(app.args))
+      if app1 ne app then report.log(i"beta reduce $app -> $app1")
+      app1
+    case TypeApply(Select(fn, nme.apply), targs) if fn.tpe.typeSymbol eq defn.PolyFunctionClass =>
+      val app1 = BetaReduce(app, fn, List(targs, app.args))
       if app1 ne app then report.log(i"beta reduce $app -> $app1")
       app1
     case _ =>
       app
 
-
 object BetaReduce:
   import ast.tpd._
 
   val name: String = "betaReduce"
   val description: String = "reduce closure applications"
 
   /** Beta-reduces a call to `fn` with arguments `argSyms` or returns `tree` */
-  def apply(original: Tree, fn: Tree, args: List[Tree])(using Context): Tree =
+  def apply(original: Tree, fn: Tree, argss: List[List[Tree]])(using Context): Tree =
     fn match
       case Typed(expr, _) =>
-        BetaReduce(original, expr, args)
+        BetaReduce(original, expr, argss)
       case Block((anonFun: DefDef) :: Nil, closure: Closure) =>
-        BetaReduce(anonFun, args)
+        BetaReduce(anonFun, argss)
+      case Block((TypeDef(_, template: Template)) :: Nil, Typed(Apply(Select(New(_), _), _), _)) if template.constr.rhs.isEmpty =>
+        template.body match
+          case (anonFun: DefDef) :: Nil =>
+            BetaReduce(anonFun, argss)
+          case _ =>
+            original
       case Block(stats, expr) =>
-        val tree = BetaReduce(original, expr, args)
+        val tree = BetaReduce(original, expr, argss)
         if tree eq original then original
         else cpy.Block(fn)(stats, tree)
       case Inlined(call, bindings, expr) =>
-        val tree = BetaReduce(original, expr, args)
+        val tree = BetaReduce(original, expr, argss)
         if tree eq original then original
         else cpy.Inlined(fn)(call, bindings, tree)
       case _ =>
         original
   end apply
 
   /** Beta-reduces a call to `ddef` with arguments `args` */
-  def apply(ddef: DefDef, args: List[Tree])(using Context) =
-    val bindings = new ListBuffer[ValDef]()
-    val expansion1 = reduceApplication(ddef, args, bindings)
+  def apply(ddef: DefDef, argss: List[List[Tree]])(using Context) =
+    val bindings = new ListBuffer[DefTree]()
+    val expansion1 = reduceApplication(ddef, argss, bindings)
     val bindings1 = bindings.result()
     seq(bindings1, expansion1)
 
   /** Beta-reduces a call to `ddef` with arguments `args` and registers new bindings */
-  def reduceApplication(ddef: DefDef, args: List[Tree], bindings: ListBuffer[ValDef])(using Context): Tree =
-    val vparams = ddef.termParamss.iterator.flatten.toList
+  def reduceApplication(ddef: DefDef, argss: List[List[Tree]], bindings: ListBuffer[DefTree])(using Context): Tree =
+    assert(argss.size == 1 || argss.size == 2)
+    val targs = if argss.size == 2 then argss.head else Nil
+    val args = argss.last
+    val tparams = ddef.leadingTypeParams
+    val vparams = ddef.termParamss.flatten
+    assert(targs.hasSameLengthAs(tparams))
     assert(args.hasSameLengthAs(vparams))
+
+    val targSyms =
+      for (targ, tparam) <- targs.zip(tparams) yield
+        targ.tpe.dealias match
+          case ref @ TypeRef(NoPrefix, _) =>
+            ref.symbol
+          case _ =>
+            val binding = TypeDef(newSymbol(ctx.owner, tparam.name, EmptyFlags, targ.tpe, coord = targ.span)).withSpan(targ.span)
+            bindings += binding
+            binding.symbol
+
     val argSyms =
       for (arg, param) <- args.zip(vparams) yield
         arg.tpe.dealias match
@@ -99,8 +124,8 @@ object BetaReduce:
     val expansion = TreeTypeMap(
       oldOwners = ddef.symbol :: Nil,
       newOwners = ctx.owner :: Nil,
-      substFrom = vparams.map(_.symbol),
-      substTo = argSyms
+      substFrom = (tparams ::: vparams).map(_.symbol),
+      substTo = targSyms ::: argSyms
     ).transform(ddef.rhs)
 
     val expansion1 = new TreeMap {
diff --git a/compiler/src/dotty/tools/dotc/transform/InlinePatterns.scala b/compiler/src/dotty/tools/dotc/transform/InlinePatterns.scala
@@ -51,11 +51,12 @@ class InlinePatterns extends MiniPhase:
         case Apply(App(fn, argss), args) => (fn, argss :+ args)
         case _ => (app, Nil)
 
+  // TODO merge with BetaReduce.scala
   private def betaReduce(tree: Apply, fn: Tree, name: Name, args: List[Tree])(using Context): Tree =
     fn match
       case Block(TypeDef(_, template: Template) :: Nil, Apply(Select(New(_),_), Nil)) if template.constr.rhs.isEmpty =>
         template.body match
-          case List(ddef @ DefDef(`name`, _, _, _)) => BetaReduce(ddef, args)
+          case List(ddef @ DefDef(`name`, _, _, _)) => BetaReduce(ddef, List(args))
           case _ => tree
       case _ => tree
 
diff --git a/compiler/src/dotty/tools/dotc/transform/PickleQuotes.scala b/compiler/src/dotty/tools/dotc/transform/PickleQuotes.scala
@@ -322,7 +322,7 @@ object PickleQuotes {
                   }
                   val Block(List(ddef: DefDef), _) = splice: @unchecked
                   // TODO: beta reduce inner closure? Or wait until BetaReduce phase?
-                  BetaReduce(ddef, spliceArgs).select(nme.apply).appliedTo(args(2).asInstance(quotesType))
+                  BetaReduce(ddef, List(spliceArgs)).select(nme.apply).appliedTo(args(2).asInstance(quotesType))
                 }
                 CaseDef(Literal(Constant(idx)), EmptyTree, rhs)
               }
diff --git a/compiler/src/scala/quoted/runtime/impl/QuotesImpl.scala b/compiler/src/scala/quoted/runtime/impl/QuotesImpl.scala
@@ -362,8 +362,9 @@ class QuotesImpl private (using val ctx: Context) extends Quotes, QuoteUnpickler
     object Term extends TermModule:
       def betaReduce(tree: Term): Option[Term] =
         tree match
+          // TODO support TypeApply. Would fix #15968.
           case app @ tpd.Apply(tpd.Select(fn, nme.apply), args) if dotc.core.Symbols.defn.isFunctionType(fn.tpe) =>
-            val app1 = dotc.transform.BetaReduce(app, fn, args)
+            val app1 = dotc.transform.BetaReduce(app, fn, List(args))
             if app1 eq app then None
             else Some(app1.withSpan(tree.span))
           case tpd.Block(Nil, expr) =>
diff --git a/compiler/test/dotty/tools/backend/jvm/InlineBytecodeTests.scala b/compiler/test/dotty/tools/backend/jvm/InlineBytecodeTests.scala
@@ -582,6 +582,32 @@ class InlineBytecodeTests extends DottyBytecodeTest {
     }
   }
 
+  @Test def beta_reduce_polymorphic_function = {
+    val source = """class Test:
+                   |  def test =
+                   |    ([Z] => (arg: Z) => { val a: Z = arg; a }).apply[Int](2)
+                 """.stripMargin
+
+    checkBCode(source) { dir =>
+      val clsIn      = dir.lookupName("Test.class", directory = false).input
+      val clsNode    = loadClassNode(clsIn)
+
+      val fun = getMethod(clsNode, "test")
+      val instructions = instructionsFromMethod(fun)
+      val expected =
+        List(
+          Op(ICONST_2),
+          VarOp(ISTORE, 1),
+          VarOp(ILOAD, 1),
+          Op(IRETURN)
+        )
+
+      assert(instructions == expected,
+        "`i was not properly beta-reduced in `test`\n" + diffInstructions(instructions, expected))
+
+    }
+  }
+
   @Test def i9456 = {
     val source = """class Foo {
                    |  def test: Int = inline2(inline1(2.+))

Original file line number	Diff line number	Diff line change
`@@ -322,7 +322,7 @@ object PickleQuotes {`
`322`	`322`	`}`
`323`	`323`	`val Block(List(ddef: DefDef), _) = splice: @unchecked`
`324`	`324`	`// TODO: beta reduce inner closure? Or wait until BetaReduce phase?`
`325`		`- BetaReduce(ddef, spliceArgs).select(nme.apply).appliedTo(args(2).asInstance(quotesType))`
	`325`	`+ BetaReduce(ddef, List(spliceArgs)).select(nme.apply).appliedTo(args(2).asInstance(quotesType))`
`326`	`326`	`}`
`327`	`327`	`CaseDef(Literal(Constant(idx)), EmptyTree, rhs)`
`328`	`328`	`}`