Beta-reduce under blocks

`{ ...; (x: T) => f(x) }.apply(y)` is beta-reduced to `{ ...; val x = y; f(x) }`.

Author: nicolasstucki

compiler/src/dotty/tools/dotc/quoted/QuoteContextImpl.scala

@@ -302,9 +302,7 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, scala.intern
       def betaReduce(tree: Term): Option[Term] =
         tree match
           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)
-            if app1 eq app then None
-            else Some(app1.withSpan(tree.span))
+            dotc.transform.BetaReduce.tryReduce(fn, args).map(_.withSpan(tree.span))
           case tpd.Block(Nil, expr) =>
             for e <- betaReduce(expr) yield tpd.cpy.Block(tree)(Nil, e)
           case tpd.Inlined(_, Nil, expr) =>

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

@@ -37,9 +37,12 @@ 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)
-      if app1 ne app then report.log(i"beta reduce $app -> $app1")
-      app1
+      BetaReduce.tryReduce(fn, app.args) match
+        case Some(app1) =>
+          report.log(i"beta reduce $app -> $app1")
+          app1
+        case _ =>
+          app
     case _ =>
       app
 
@@ -48,17 +51,22 @@ object BetaReduce:
   import ast.tpd._
 
   /** Beta-reduces a call to `fn` with arguments `argSyms` or returns `tree` */
-  def apply(tree: Apply, fn: Tree, args: List[Tree])(using Context): Tree =
+  def tryReduce(fn: Tree, args: List[Tree])(using Context): Option[Tree] =
     fn match
-      case Typed(expr, _) => BetaReduce(tree, expr, args)
-      case Block(Nil, expr) => BetaReduce(tree, expr, args)
-      case Inlined(_, Nil, expr) => BetaReduce(tree, expr, args)
-      case Block((anonFun: DefDef) :: Nil, closure: Closure) => BetaReduce(anonFun, args)
-      case _ => tree
-  end apply
+      case Typed(expr, _) =>
+        tryReduce(expr, args)
+      case Block((anonFun: DefDef) :: Nil, closure: Closure) =>
+        Some(reduce(anonFun, args))
+      case Block(stats, expr) =>
+        tryReduce(expr, args).map(newExpr => cpy.Block(fn)(stats, newExpr))
+      case Inlined(call, bindings, expr) =>
+        tryReduce(expr, args).map(newExpr => cpy.Inlined(fn)(call, bindings, newExpr))
+      case _ =>
+        None
+  end tryReduce
 
   /** Beta-reduces a call to `ddef` with arguments `argSyms` */
-  def apply(ddef: DefDef, args: List[Tree])(using Context) =
+  def reduce(ddef: DefDef, args: List[Tree])(using Context) =
     val bindings = List.newBuilder[ValDef]
     val vparams = ddef.vparamss.iterator.flatten.toList
     assert(args.hasSameLengthAs(vparams))
@@ -90,4 +98,4 @@ object BetaReduce:
       bindings.result().filterNot(vdef => vdef.tpt.tpe.isInstanceOf[ConstantType] && isPureExpr(vdef.rhs))
 
     seq(bindings1, expansion1)
-  end apply
+  end reduce

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

@@ -56,6 +56,6 @@ class InlinePatterns extends MiniPhase:
     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.reduce(ddef, args)
           case _ => tree
       case _ => tree

compiler/test/dotty/tools/backend/jvm/InlineBytecodeTests.scala

@@ -368,14 +368,14 @@ class InlineBytecodeTests extends DottyBytecodeTest {
       val instructions = instructionsFromMethod(fun)
       val expected =
         List(
-          IntOp(BIPUSH, 10)
-          , VarOp(ISTORE, 1)
-          , VarOp(ILOAD, 1)
-          , Op(ICONST_1)
-          , Op(IADD)
-          , VarOp(ISTORE, 1)
-          , VarOp(ILOAD, 1)
-          , Op(IRETURN)
+          IntOp(BIPUSH, 10),
+          VarOp(ISTORE, 1),
+          VarOp(ILOAD, 1),
+          Op(ICONST_1),
+          Op(IADD),
+          VarOp(ISTORE, 1),
+          VarOp(ILOAD, 1),
+          Op(IRETURN),
         )
       assert(instructions == expected,
         "`f` was not properly inlined in `fun`\n" + diffInstructions(instructions, expected))
@@ -524,4 +524,35 @@ class InlineBytecodeTests extends DottyBytecodeTest {
 
     }
   }
+
+  @Test def beta_reduce_under_block = {
+    val source = """class Test:
+                   |  def test =
+                   |    {
+                   |      val a = 3
+                   |      (i: Int) => i + a
+                   |    }.apply(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_3),
+          VarOp(ISTORE, 1),
+          Op(ICONST_2),
+          VarOp(ILOAD, 1),
+          Op(IADD),
+          Op(IRETURN),
+        )
+
+      assert(instructions == expected,
+        "`i was not properly beta-reduced in `test`\n" + diffInstructions(instructions, expected))
+
+    }
+  }
 }