Fix erased context function types, 2nd attempt

compiler/src/dotty/tools/dotc/ast/tpd.scala

@@ -121,9 +121,9 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
       Closure(Nil, call, targetTpt))
   }
 
-  /** A closure whole anonymous function has the given method type */
+  /** A closure whose anonymous function has the given method type */
   def Lambda(tpe: MethodType, rhsFn: List[Tree] => Tree)(using Context): Block = {
-    val meth = newSymbol(ctx.owner, nme.ANON_FUN, Synthetic | Method, tpe)
+    val meth = newAnonFun(ctx.owner, tpe)
     Closure(meth, tss => rhsFn(tss.head).changeOwner(ctx.owner, meth))
   }
 
@@ -1104,6 +1104,21 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
       if (sym.exists) sym.defTree = tree
       tree
     }
+
+    def etaExpandCFT(using Context): Tree =
+      def expand(target: Tree, tp: Type)(using Context): Tree = tp match
+        case defn.ContextFunctionType(argTypes, resType, isErased) =>
+          val anonFun = newAnonFun(
+            ctx.owner,
+            MethodType.companion(isContextual = true, isErased = isErased)(argTypes, resType),
+            coord = ctx.owner.coord)
+          def lambdaBody(refss: List[List[Tree]]) =
+            expand(target.select(nme.apply).appliedToArgss(refss), resType)(
+              using ctx.withOwner(anonFun))
+          Closure(anonFun, lambdaBody)
+        case _ =>
+          target
+      expand(tree, tree.tpe.widen)
   }
 
   inline val MapRecursionLimit = 10

compiler/src/dotty/tools/dotc/core/Definitions.scala

@@ -1353,23 +1353,6 @@ class Definitions {
   def isBoxedUnitClass(cls: Symbol): Boolean =
     cls.isClass && (cls.owner eq ScalaRuntimePackageClass) && cls.name == tpnme.BoxedUnit
 
-  /** Returns the erased class of the function class `cls`
-   *    - FunctionN for N > 22 becomes FunctionXXL
-   *    - FunctionN for 22 > N >= 0 remains as FunctionN
-   *    - ContextFunctionN for N > 22 becomes FunctionXXL
-   *    - ContextFunctionN for N <= 22 becomes FunctionN
-   *    - ErasedFunctionN becomes Function0
-   *    - ImplicitErasedFunctionN becomes Function0
-   *    - anything else becomes a NoSymbol
-   */
-  def erasedFunctionClass(cls: Symbol): Symbol = {
-    val arity = scalaClassName(cls).functionArity
-    if (cls.name.isErasedFunction) FunctionClass(0)
-    else if (arity > 22) FunctionXXLClass
-    else if (arity >= 0) FunctionClass(arity)
-    else NoSymbol
-  }
-
   /** Returns the erased type of the function class `cls`
    *    - FunctionN for N > 22 becomes FunctionXXL
    *    - FunctionN for 22 > N >= 0 remains as FunctionN
@@ -1379,13 +1362,12 @@ class Definitions {
    *    - ImplicitErasedFunctionN becomes Function0
    *    - anything else becomes a NoType
    */
-  def erasedFunctionType(cls: Symbol): Type = {
+  def functionTypeErasure(cls: Symbol): Type =
     val arity = scalaClassName(cls).functionArity
-    if (cls.name.isErasedFunction) FunctionType(0)
-    else if (arity > 22) FunctionXXLClass.typeRef
-    else if (arity >= 0) FunctionType(arity)
+    if cls.name.isErasedFunction then FunctionType(0)
+    else if arity > 22 then FunctionXXLClass.typeRef
+    else if arity >= 0 then FunctionType(arity)
     else NoType
-  }
 
   val predefClassNames: Set[Name] =
     Set("Predef$", "DeprecatedPredef", "LowPriorityImplicits").map(_.toTypeName.unmangleClassName)

compiler/src/dotty/tools/dotc/core/NameKinds.scala

@@ -365,7 +365,7 @@ object NameKinds {
   val ExtMethName: SuffixNameKind = new SuffixNameKind(EXTMETH, "$extension")
   val ParamAccessorName: SuffixNameKind = new SuffixNameKind(PARAMACC, "$accessor")
   val ModuleClassName: SuffixNameKind = new SuffixNameKind(OBJECTCLASS, "$", optInfoString = "ModuleClass")
-  val ImplMethName: SuffixNameKind = new SuffixNameKind(IMPLMETH, "$")
+  val DirectMethName: SuffixNameKind = new SuffixNameKind(DIRECT, "$direct")
   val AdaptedClosureName: SuffixNameKind = new SuffixNameKind(ADAPTEDCLOSURE, "$adapted") { override def definesNewName = true }
   val SyntheticSetterName: SuffixNameKind = new SuffixNameKind(SETTER, "_$eq")
 

compiler/src/dotty/tools/dotc/core/NameTags.scala

@@ -24,8 +24,9 @@ object NameTags extends TastyFormat.NameTags {
 
   final val ADAPTEDCLOSURE = 31 // Used in Erasure to adapt closures over primitive types.
 
-  final val IMPLMETH = 32       // Used to define methods in implementation classes
-                                // (can probably be removed).
+  final val DIRECT = 32         // Used to define implementations of methods with
+                                // erased context function results that can override some
+                                // other method.
 
   final val PARAMACC = 33       // Used for a private parameter alias
 
@@ -48,7 +49,7 @@ object NameTags extends TastyFormat.NameTags {
     case INITIALIZER => "INITIALIZER"
     case FIELD => "FIELD"
     case EXTMETH => "EXTMETH"
-    case IMPLMETH => "IMPLMETH"
+    case DIRECT => "DIRECT"
     case PARAMACC => "PARAMACC"
     case ADAPTEDCLOSURE => "ADAPTEDCLOSURE"
     case OBJECTCLASS => "OBJECTCLASS"

compiler/src/dotty/tools/dotc/core/Symbols.scala

@@ -712,6 +712,10 @@ object Symbols {
       coord: Coord = NoCoord)(using Context): TermSymbol =
     newSymbol(cls, nme.CONSTRUCTOR, flags | Method, MethodType(paramNames, paramTypes, cls.typeRef), privateWithin, coord)
 
+  /** Create an anonymous function symbol */
+  def newAnonFun(owner: Symbol, info: Type, coord: Coord = NoCoord)(using Context): TermSymbol =
+    newSymbol(owner, nme.ANON_FUN, Synthetic | Method, info, coord = coord)
+
   /** Create an empty default constructor symbol for given class `cls`. */
   def newDefaultConstructor(cls: ClassSymbol)(using Context): TermSymbol =
     newConstructor(cls, EmptyFlags, Nil, Nil)

compiler/src/dotty/tools/dotc/core/TypeErasure.scala

@@ -581,7 +581,7 @@ class TypeErasure(sourceLanguage: SourceLanguage, semiEraseVCs: Boolean, isConst
       val sym = tp.symbol
       if (!sym.isClass) this(tp.translucentSuperType)
       else if (semiEraseVCs && isDerivedValueClass(sym)) eraseDerivedValueClass(tp)
-      else if (defn.isSyntheticFunctionClass(sym)) defn.erasedFunctionType(sym)
+      else if (defn.isSyntheticFunctionClass(sym)) defn.functionTypeErasure(sym)
       else eraseNormalClassRef(tp)
     case tp: AppliedType =>
       val tycon = tp.tycon
@@ -791,7 +791,7 @@ class TypeErasure(sourceLanguage: SourceLanguage, semiEraseVCs: Boolean, isConst
           if (erasedVCRef.exists) return sigName(erasedVCRef)
         }
         if (defn.isSyntheticFunctionClass(sym))
-          sigName(defn.erasedFunctionType(sym))
+          sigName(defn.functionTypeErasure(sym))
         else
           val cls = normalizeClass(sym.asClass)
           val fullName =

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

@@ -51,7 +51,10 @@ abstract class AccessProxies {
               forwardedArgss.nonEmpty && forwardedArgss.head.nonEmpty) // defensive conditions
             accessRef.becomes(forwardedArgss.head.head)
           else
-            accessRef.appliedToTypeTrees(forwardedTpts).appliedToArgss(forwardedArgss)
+            accessRef
+              .appliedToTypeTrees(forwardedTpts)
+              .appliedToArgss(forwardedArgss)
+              .etaExpandCFT(using ctx.withOwner(accessor))
         rhs.withSpan(accessed.span)
       })
 

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

@@ -9,6 +9,11 @@ import ast.untpd
 import collection.{mutable, immutable}
 import util.Spans.Span
 import util.SrcPos
+import ContextFunctionResults.{contextResultCount, contextFunctionResultTypeAfter}
+import StdNames.nme
+import Constants.Constant
+import TypeErasure.transformInfo
+import Erasure.Boxing.adaptClosure
 
 /** A helper class for generating bridge methods in class `root`. */
 class Bridges(root: ClassSymbol, thisPhase: DenotTransformer)(using Context) {
@@ -112,12 +117,52 @@ class Bridges(root: ClassSymbol, thisPhase: DenotTransformer)(using Context) {
       toBeRemoved += other
     }
 
-    def bridgeRhs(argss: List[List[Tree]]) = {
+    val memberCount = contextResultCount(member)
+
+    /** Eta expand application `ref(args)` as needed.
+     *  To do this correctly, we have to look at the member's original pre-erasure
+     *  type and figure out which context function types in its result are
+     *  not yet instantiated.
+     */
+    def etaExpand(ref: Tree, args: List[Tree])(using Context): Tree =
+      def expand(args: List[Tree], tp: Type, n: Int)(using Context): Tree =
+        if n <= 0 then
+          assert(ctx.typer.isInstanceOf[Erasure.Typer])
+          ctx.typer.typed(untpd.cpy.Apply(ref)(ref, args), member.info.finalResultType)
+        else
+          val defn.ContextFunctionType(argTypes, resType, isErased) = tp: @unchecked
+          val anonFun = newAnonFun(ctx.owner,
+            MethodType(if isErased then Nil else argTypes, resType),
+            coord = ctx.owner.coord)
+          anonFun.info = transformInfo(anonFun, anonFun.info)
+
+          def lambdaBody(refss: List[List[Tree]]) =
+            val refs :: Nil = refss: @unchecked
+            val expandedRefs = refs.map(_.withSpan(ctx.owner.span.endPos)) match
+              case (bunchedParam @ Ident(nme.ALLARGS)) :: Nil =>
+                argTypes.indices.toList.map(n =>
+                  bunchedParam
+                    .select(nme.primitive.arrayApply)
+                    .appliedTo(Literal(Constant(n))))
+              case refs1 => refs1
+            expand(args ::: expandedRefs, resType, n - 1)(using ctx.withOwner(anonFun))
+
+          val unadapted = Closure(anonFun, lambdaBody)
+          cpy.Block(unadapted)(unadapted.stats,
+            adaptClosure(unadapted.expr.asInstanceOf[Closure]))
+      end expand
+
+      val otherCount = contextResultCount(other)
+      val start = contextFunctionResultTypeAfter(member, otherCount)(using preErasureCtx)
+      expand(args, start, memberCount - otherCount)(using ctx.withOwner(bridge))
+    end etaExpand
+
+    def bridgeRhs(argss: List[List[Tree]]) =
       assert(argss.tail.isEmpty)
       val ref = This(root).select(member)
-      if (member.info.isParameterless) ref // can happen if `member` is a module
-      else Erasure.partialApply(ref, argss.head)
-    }
+      if member.info.isParameterless then ref // can happen if `member` is a module
+      else if memberCount == 0 then ref.appliedToTermArgs(argss.head)
+      else etaExpand(ref, argss.head)
 
     bridges += DefDef(bridge, bridgeRhs(_).withSpan(bridge.span))
   }

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

@@ -30,11 +30,9 @@ class ByNameClosures extends TransformByNameApply with IdentityDenotTransformer
     // ExpanSAMs applied to partial functions creates methods that need
     // to be fully defined before converting. Test case is pos/i9391.scala.
 
-  override def mkByNameClosure(arg: Tree, argType: Type)(using Context): Tree = {
-    val meth = newSymbol(
-      ctx.owner, nme.ANON_FUN, Synthetic | Method, MethodType(Nil, Nil, argType))
+  override def mkByNameClosure(arg: Tree, argType: Type)(using Context): Tree =
+    val meth = newAnonFun(ctx.owner, MethodType(Nil, argType))
     Closure(meth, _ => arg.changeOwnerAfter(ctx.owner, meth, thisPhase)).withSpan(arg.span)
-  }
 }
 
 object ByNameClosures {

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

@@ -8,12 +8,16 @@ import StdNames.nme
 import ast.untpd
 import ast.tpd._
 import config.Config
+import Decorators.*
 
 object ContextFunctionResults:
 
   /** Annotate methods that have context function result types directly matched by context
    *  closures on their right-hand side. Parameters to such closures will be integrated
    *  as additional method parameters in erasure.
+   *
+   *  A @ContextResultCount(n) annotation means that the method's result type
+   *  consists of a string of `n` nested context closures.
    */
   def annotateContextResults(mdef: DefDef)(using Context): Unit =
     def contextResultCount(rhs: Tree, tp: Type): Int = tp match
@@ -50,6 +54,15 @@ object ContextFunctionResults:
         crCount
       case none => 0
 
+  /** True iff `ContextResultCount` is not zero and all context functions in the result
+   *  type are erased.
+   */
+  def contextResultsAreErased(sym: Symbol)(using Context): Boolean =
+    def allErased(tp: Type): Boolean = tp.dealias match
+      case defn.ContextFunctionType(_, resTpe, isErased) => isErased && allErased(resTpe)
+      case _ => true
+    contextResultCount(sym) > 0 && allErased(sym.info.finalResultType)
+
   /** Turn the first `crCount` context function types in the result type of `tp`
    *  into the curried method types.
    */
@@ -86,33 +99,13 @@ object ContextFunctionResults:
     normalParamCount(sym.info)
   end totalParamCount
 
-  /** The rightmost context function type in the result type of `meth`
-   *  that represents `paramCount` curried, non-erased parameters that
-   *  are included in the `contextResultCount` of `meth`.
-   *  Example:
-   *
-   *  Say we have `def m(x: A): B ?=> (C1, C2, C3) ?=> D ?=> E ?=> F`,
-   *  paramCount == 4, and the contextResultCount of `m` is 3.
-   *  Then we return the type `(C1, C2, C3) ?=> D ?=> E ?=> F`, since this
-   *  type covers the 4 rightmost parameters C1, C2, C3 and D before the
-   *  contextResultCount runs out at E ?=> F.
-   *  Erased parameters are ignored; they contribute nothing to the
-   *  parameter count.
-   */
-  def contextFunctionResultTypeCovering(meth: Symbol, paramCount: Int)(using Context) =
-    atPhase(erasurePhase) {
-      // Recursive instances return pairs of context types and the
-      // # of parameters they represent.
-      def missingCR(tp: Type, crCount: Int): (Type, Int) =
-        if crCount == 0 then (tp, 0)
-        else
-          val defn.ContextFunctionType(formals, resTpe, isErased) = tp: @unchecked
-          val result @ (rt, nparams) = missingCR(resTpe, crCount - 1)
-          assert(nparams <= paramCount)
-          if nparams == paramCount || isErased then result
-          else (tp, nparams + formals.length)
-      missingCR(meth.info.finalResultType, contextResultCount(meth))._1
-    }
+  /** The `depth` levels nested context function type in the result type of `meth` */
+  def contextFunctionResultTypeAfter(meth: Symbol, depth: Int)(using Context) =
+    def recur(tp: Type, n: Int): Type =
+      if n == 0 then tp
+      else tp match
+        case defn.ContextFunctionType(_, resTpe, _) => recur(resTpe, n - 1)
+    recur(meth.info.finalResultType, depth)
 
   /** Should selection `tree` be eliminated since it refers to an `apply`
    *  node of a context function type whose parameters will end up being