Add eta reduction miniphase

Fixes #10889. Alternative to #13750, from which the test is taken.

Author: odersky

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -106,6 +106,7 @@ class Compiler {
     List(new ElimErasedValueType,    // Expand erased value types to their underlying implmementation types
          new PureStats,              // Remove pure stats from blocks
          new VCElideAllocations,     // Peep-hole optimization to eliminate unnecessary value class allocations
+         new EtaReduce,
          new ArrayApply,             // Optimize `scala.Array.apply([....])` and `scala.Array.apply(..., [....])` into `[...]`
          new sjs.AddLocalJSFakeNews, // Adds fake new invocations to local JS classes in calls to `createLocalJSClass`
          new ElimPolyFunction,       // Rewrite PolyFunction subclasses to FunctionN subclasses

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

@@ -0,0 +1,46 @@
+package dotty.tools
+package dotc
+package transform
+
+import MegaPhase.MiniPhase
+import core.*
+import Symbols.*, Contexts.*, Types.*, Decorators.*
+import StdNames.nme
+import ast.Trees.*
+
+/** Rewrite `(x1, ... xN) => f(x1, ... xN)` for N >= 0 to `f`,
+ *  provided `f` is a pure path of function type.
+ *
+ *  This optimization is crucial for context functions. The compiler
+ *  produces a contextual closure around values passed as arguments
+ *  where a context function is expected, unless that value has the
+ *  syntactic form of a context function literal.
+ *
+ *  Without this phase, when a contextual function is passed as an argument to a
+ *  recursive function, that would have the unfortunate effect of a linear growth
+ *  in transient thunks of identical type wrapped around each other, leading
+ *  to performance degradation, and in some cases, stack overflows.
+ */
+class EtaReduce extends MiniPhase:
+  import ast.tpd._
+
+  override def phaseName: String = "etaReduce"
+
+  override def transformBlock(tree: Block)(using Context): Tree = tree match
+    case Block((meth : DefDef) :: Nil, closure: Closure)
+    if meth.symbol == closure.meth.symbol =>
+      meth.rhs match
+        case Apply(Select(fn, nme.apply), args)
+        if meth.paramss.head.corresponds(args)((param, arg) =>
+              arg.isInstanceOf[Ident] && arg.symbol == param.symbol)
+            && isPurePath(fn) =>
+          val treeSym = tree.tpe.widen.typeSymbol
+          val fnSym = fn.tpe.widen.typeSymbol
+          if treeSym == fnSym && defn.isFunctionClass(fnSym) then
+            report.log(i"eta reducing $tree --> $fn")
+            fn
+          else tree
+        case _ => tree
+    case _ => tree
+
+end EtaReduce

tests/run/contextual-closure-unwrapping.check

@@ -0,0 +1 @@
+true

tests/run/contextual-closure-unwrapping.scala

@@ -0,0 +1,83 @@
+import scala.annotation.tailrec
+
+import java.lang.management.ManagementFactory
+import java.lang.management.RuntimeMXBean
+
+import scala.jdk.CollectionConverters._
+
+object Test {
+  trait Txn {}
+  type AtomicOp[Z] = Txn ?=> Z
+  type VanillaAtomicOp[Z] = Txn => Z
+
+  object AbortAndRetry extends scala.util.control.ControlThrowable("abort and retry") {}
+
+  def beginTxn:Txn = new Txn {}
+
+  @tailrec def retryN[Z](n:Int)(txn:AtomicOp[Z], i:Int = 0):Z = {
+    try {
+      given Txn = beginTxn
+      val ret:Z = txn
+      if(i < n) { throw AbortAndRetry }
+      ret
+    } catch {
+      case AbortAndRetry => retryN(n)(txn, i + 1)
+    }
+  }
+
+  @tailrec def safeRetryN[Z](n:Int)(txn:VanillaAtomicOp[Z], i:Int = 0):Z = {
+    try {
+      given Txn = beginTxn
+      val ret:Z = txn.asInstanceOf[AtomicOp[Z]]
+      if(i < n) { throw AbortAndRetry }
+      ret
+    } catch {
+      case AbortAndRetry => safeRetryN(n)(txn, i + 1)
+    }
+  }
+
+  object StackSize {
+    def unapply(arg:String):Option[Int] = {
+      Option(arg match {
+        case s"-Xss$rest" => rest
+        case s"-XX:ThreadStackSize=$rest" => rest
+        case _ => null
+      }).map{ rest =>
+        val shift = (rest.toLowerCase.last match {
+          case 'k' => 10
+          case 'm' => 20
+          case 'g' => 30
+          case _ => 0
+        })
+        (if(shift > 0) {
+          rest.dropRight(1)
+        } else {
+          rest
+        }).toInt << shift
+      }
+    }
+  }
+  def main(args:Array[String]) = {
+    val arguments = ManagementFactory.getRuntimeMXBean.getInputArguments
+    // 64bit VM defaults to 1024K / 1M stack size, 32bit defaults to 320k
+    // Use 1024 as upper bound.
+    val maxStackSize:Int = arguments.asScala.reverseIterator.collectFirst{case StackSize(stackSize) => stackSize}.getOrElse(1 << 20)
+
+    val testTxn:VanillaAtomicOp[Boolean] = {
+      (txn:Txn) =>
+        given Txn = txn
+        true
+    }
+
+    Console.println(try {
+      // maxStackSize is a good upper bound on linear stack growth
+      // without assuming too much about frame size
+      // (1 byte / frame is conservative even for a z80)
+      retryN(maxStackSize)(testTxn.asInstanceOf[AtomicOp[Boolean]]) == safeRetryN(maxStackSize)(testTxn)
+    } catch {
+      case e:StackOverflowError =>
+        Console.println(s"Exploded after ${e.getStackTrace.length} frames")
+        false
+    })
+  }
+}