Make implicit function variables safe as arguments (fixes #10889)

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

@@ -73,6 +73,7 @@ class Compiler {
          new ProtectedAccessors,     // Add accessors for protected members
          new ExtensionMethods,       // Expand methods of value classes with extension methods
          new UncacheGivenAliases,    // Avoid caching RHS of simple parameterless given aliases
+         new ElimRedundantContextualClosure, // Unwrapped context closures if the contents is already a context closure of compatible type
          new ByNameClosures,         // Expand arguments to by-name parameters to closures
          new HoistSuperArgs,         // Hoist complex arguments of supercalls to enclosing scope
          new SpecializeApplyMethods, // Adds specialized methods to FunctionN

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

@@ -0,0 +1,64 @@
+package dotty.tools
+package dotc
+package transform
+
+import MegaPhase._
+import core._
+import Symbols._
+import SymDenotations._
+import Contexts._
+import Types._
+import Flags._
+import Decorators._
+import DenotTransformers._
+import core.StdNames.nme
+import core.StdNames
+import ast.Trees._
+import reporting.trace
+import config.Printers.{ transforms => debug }
+
+/** Abstract base class of ByNameClosures and ElimByName, factoring out the
+ *  common functionality to transform arguments of by-name parameters.
+ */
+class ElimRedundantContextualClosure extends MiniPhase with IdentityDenotTransformer { thisPhase: DenotTransformer =>
+  import ast.tpd._
+  import ast.untpd
+
+  override def phaseName:String = ElimRedundantContextualClosure.name
+
+  /** The info of the tree's symbol before it is potentially transformed in this phase */
+  private def originalDenotation(tree: Tree)(using Context) =
+    atPhase(thisPhase)(tree.symbol.denot)
+
+  override def transformApply(tree: Apply)(using Context): Tree =
+    trace(s"transforming ${tree.show} at phase ${ctx.phase}", show = true) {
+
+      def transformArg(arg: Tree, formal: Type): Tree = {
+        val veryFormal = formal.widenDealias
+        if(defn.isContextFunctionType(veryFormal) && untpd.isContextualClosure(arg)) {
+          val body = unsplice(closureBody(arg)) match {
+            case Apply(Select(fn, nme.apply), _) => fn
+            case other => other
+          } // no-op if not a nested closure of some kind
+          val underlyingBodyType = body.tpe.widenDealias
+          val bodyIsContextual = defn.isContextFunctionType(underlyingBodyType)
+          val bodyTypeMatches = TypeComparer.isSubType(underlyingBodyType, veryFormal)
+          if(bodyIsContextual && bodyTypeMatches) {
+            body
+          } else {
+            arg
+          }
+        } else {
+          arg
+        }
+      }
+
+      val mt @ MethodType(_) = tree.fun.tpe.widen
+      val args1 = tree.args.zipWithConserve(mt.paramInfos)(transformArg)
+      cpy.Apply(tree)(tree.fun, args1)
+    }
+}
+
+object ElimRedundantContextualClosure {
+  val name: String = "elimRedundantContextualClosure"
+}

tests/run/contextual-closure-unwrapping.check

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

tests/run/contextual-closure-unwrapping.scala

@@ -0,0 +1,84 @@
+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
+    })
+  }
+}