Store a list of erased flags in FunctionWithMods, allow parsing more erased function types

`(Int, erased Int) => Int` is now allowed as a function type.

Author: natsukagami

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

@@ -1493,10 +1493,10 @@ object desugar {
       case vd: ValDef => vd
     }
 
-  def makeContextualFunction(formals: List[Tree], body: Tree, isErased: Boolean)(using Context): Function = {
-    val mods = if (isErased) Given | Erased else Given
+  def makeContextualFunction(formals: List[Tree], body: Tree, isErased: List[Boolean])(using Context): Function = {
+    val mods = Given
     val params = makeImplicitParameters(formals, mods)
-    FunctionWithMods(params, body, Modifiers(mods))
+    FunctionWithMods(params, body, Modifiers(mods), isErased)
   }
 
   private def derivedValDef(original: Tree, named: NameTree, tpt: Tree, rhs: Tree, mods: Modifiers)(using Context) = {
@@ -1829,6 +1829,7 @@ object desugar {
             cpy.ByNameTypeTree(parent)(annotate(tpnme.retainsByName, restpt))
           case _ =>
             annotate(tpnme.retains, parent)
+      case f: FunctionWithMods if f.isErased.contains(true) => makeErasedFunctionValDefs(f, pt)
     }
     desugared.withSpan(tree.span)
   }
@@ -1904,6 +1905,22 @@ object desugar {
     TypeDef(tpnme.REFINE_CLASS, impl).withFlags(Trait)
   }
 
+  /** Make erased function definitions use only ValDefs */
+  def makeErasedFunctionValDefs(tree: FunctionWithMods, pt: Type)(using Context): Function = {
+    val isErased = tree.isErased
+    val Function(args, result) = tree
+    args match {
+      case (_ : ValDef) :: _ => tree // ValDef case can be easily handled
+      case _ if !isErased.contains(true) => tree // not erased
+      case _ if !ctx.mode.is(Mode.Type) => tree
+      case _ =>
+        val applyVParams = args.zipWithIndex.map {
+          case (p, n) => makeSyntheticParameter(n + 1, p).withAddedFlags(tree.mods.flags)
+        }
+        untpd.FunctionWithMods(applyVParams, tree.body, tree.mods, isErased)
+    }
+  }
+
   /** Returns list of all pattern variables, possibly with their types,
    *  without duplicates
    */

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

@@ -76,8 +76,10 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   }
 
   /** A function type or closure with `implicit`, `erased`, or `given` modifiers */
-  class FunctionWithMods(args: List[Tree], body: Tree, val mods: Modifiers)(implicit @constructorOnly src: SourceFile)
-    extends Function(args, body)
+  class FunctionWithMods(args: List[Tree], body: Tree, val mods: Modifiers, val isErased: List[Boolean])(implicit @constructorOnly src: SourceFile)
+    extends Function(args, body) {
+      assert(args.length == isErased.length)
+    }
 
   /** A polymorphic function type */
   case class PolyFunction(targs: List[Tree], body: Tree)(implicit @constructorOnly src: SourceFile) extends Tree {

compiler/src/dotty/tools/dotc/parsing/Parsers.scala

@@ -1464,6 +1464,7 @@ object Parsers {
     def typ(): Tree =
       val start = in.offset
       var imods = Modifiers()
+      var erasedArgs: ListBuffer[Boolean] = ListBuffer()
       def functionRest(params: List[Tree]): Tree =
         val paramSpan = Span(start, in.lastOffset)
         atSpan(start, in.offset) {
@@ -1495,10 +1496,10 @@ object Parsers {
               if isByNameType(tpt) then
                 syntaxError(em"parameter of type lambda may not be call-by-name", tpt.span)
             TermLambdaTypeTree(params.asInstanceOf[List[ValDef]], resultType)
-          else if imods.isOneOf(Given | Impure) then
+          else if imods.isOneOf(Given | Impure) || erasedArgs.exists(_ == true) then
             if imods.is(Given) && params.isEmpty then
               syntaxError(em"context function types require at least one parameter", paramSpan)
-            FunctionWithMods(params, resultType, imods)
+            FunctionWithMods(params, resultType, imods, erasedArgs.toList)
           else if !ctx.settings.YkindProjector.isDefault then
             val (newParams :+ newResultType, tparams) = replaceKindProjectorPlaceholders(params :+ resultType): @unchecked
             lambdaAbstract(tparams, Function(newParams, newResultType))
@@ -1517,25 +1518,29 @@ object Parsers {
           }
           else {
             val paramStart = in.offset
-            val firstMods = if isErased then addModifier(imods) else imods
+            def addErased() =
+              erasedArgs.addOne(isErased)
+              if isErased then { in.skipToken(); }
+            addErased()
             val ts = in.currentRegion.withCommasExpected {
               funArgType() match
                 case Ident(name) if name != tpnme.WILDCARD && in.isColon =>
                   isValParamList = true
                   def funParam(start: Offset, mods: Modifiers) = {
                     atSpan(start) {
-                      val mods1 = if isErased then addModifier(mods) else mods
-                      typedFunParam(in.offset, ident(), mods1)
+                      addErased()
+                      typedFunParam(in.offset, ident(), imods)
                     }
                   }
                   commaSeparatedRest(
-                    typedFunParam(paramStart, name.toTermName, firstMods),
+                    typedFunParam(paramStart, name.toTermName, imods),
                     () => funParam(in.offset, imods))
                 case t =>
-                  // For now we just reject `erased` in (T, U) => V definitions (i.e. you need parameter names)
-                  if firstMods.is(Erased) then
-                    syntaxError(em"Implementation restriction: erased parameters are not allowed without parameter names", paramStart)
-                  commaSeparatedRest(t, funArgType)
+                  def funParam() = {
+                      addErased()
+                      funArgType()
+                  }
+                  commaSeparatedRest(t, funParam)
             }
             accept(RPAREN)
             if isValParamList || in.isArrow || isPureArrow then
@@ -1589,7 +1594,7 @@ object Parsers {
           if isPureArrow then
             functionRest(t :: Nil)
           else
-            if (imods.is(Erased) && !t.isInstanceOf[FunctionWithMods])
+            if (erasedArgs.exists(_ == true) && !t.isInstanceOf[FunctionWithMods])
               syntaxError(ErasedTypesCanOnlyBeFunctionTypes(), implicitKwPos(start))
             t
     end typ

compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala

@@ -647,27 +647,29 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
           case str: Literal => strText(str)
         }
         toText(id) ~ "\"" ~ Text(segments map segmentText, "") ~ "\""
-      case Function(args, body) =>
+      case fn @ Function(args, body) =>
         var implicitSeen: Boolean = false
         var isGiven: Boolean = false
-        var isErased: Boolean = false
-        def argToText(arg: Tree) = arg match {
+        val isErased = fn match {
+          case fn: FunctionWithMods => fn.isErased
+          case _ => fn.args.map(_ => false)
+        }
+        def argToText(arg: Tree, isErased: Boolean) = arg match {
           case arg @ ValDef(name, tpt, _) =>
             val implicitText =
               if ((arg.mods.is(Given))) { isGiven = true; "" }
-              else if ((arg.mods.is(Erased))) { isErased = true; "" }
               else if ((arg.mods.is(Implicit)) && !implicitSeen) { implicitSeen = true; keywordStr("implicit ") }
               else ""
-            implicitText ~ toText(name) ~ optAscription(tpt)
+            val erasedText = if isErased then keywordStr("erased ") else ""
+            implicitText ~ erasedText ~ toText(name) ~ optAscription(tpt)
           case _ =>
             toText(arg)
         }
         val argsText = args match {
-          case (arg @ ValDef(_, tpt, _)) :: Nil if tpt.isEmpty => argToText(arg)
+          case (arg @ ValDef(_, tpt, _)) :: Nil if tpt.isEmpty => argToText(arg, isErased(0))
           case _ =>
             "("
-            ~ keywordText("erased ").provided(isErased)
-            ~ Text(args.map(argToText), ", ")
+            ~ Text(args.zip(isErased).map(argToText), ", ")
             ~ ")"
         }
         val isPure =

compiler/src/dotty/tools/dotc/typer/EtaExpansion.scala

@@ -285,8 +285,8 @@ object EtaExpansion extends LiftImpure {
     val body = Apply(lifted, ids)
     if (mt.isContextualMethod) body.setApplyKind(ApplyKind.Using)
     val fn =
-      if (mt.isContextualMethod) new untpd.FunctionWithMods(params, body, Modifiers(Given))
-      else if (mt.isImplicitMethod) new untpd.FunctionWithMods(params, body, Modifiers(Implicit))
+      if (mt.isContextualMethod) new untpd.FunctionWithMods(params, body, Modifiers(Given), params.map(_ => false))
+      else if (mt.isImplicitMethod) new untpd.FunctionWithMods(params, body, Modifiers(Implicit), params.map(_ => false))
       else untpd.Function(params, body)
     if (defs.nonEmpty) untpd.Block(defs.toList map (untpd.TypedSplice(_)), fn) else fn
   }

compiler/src/dotty/tools/dotc/typer/Typer.scala

@@ -1290,32 +1290,15 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
 
   def typedFunctionType(tree: untpd.Function, pt: Type)(using Context): Tree = {
     val untpd.Function(args, body) = tree
-    var funFlags = tree match {
-      case tree: untpd.FunctionWithMods => tree.mods.flags
-      case _ => EmptyFlags
+    var (funFlags, isErased) = tree match {
+      case tree: untpd.FunctionWithMods => (tree.mods.flags, tree.isErased)
+      case _ => (EmptyFlags, args.map(_ => false))
     }
 
-    assert(!funFlags.is(Erased) || !args.isEmpty, "An empty function cannot not be erased")
-
     val numArgs = args.length
     val isContextual = funFlags.is(Given)
-    val isErased = args.collect({ case x: untpd.ValDef => x }).exists(_.mods.flags.is(Erased))
     val isImpure = funFlags.is(Impure)
-    val funSym = defn.FunctionSymbol(numArgs, isContextual, isErased, isImpure)
-
-    /** If `app` is a function type with arguments that are all erased classes,
-     *  turn it into an erased function type.
-     */
-    def propagateErased(app: Tree): Tree = app match
-      case AppliedTypeTree(tycon: TypeTree, args)
-      if !isErased
-         && numArgs > 0
-         && args.indexWhere(!_.tpe.isErasedClass) == numArgs =>
-        val tycon1 = TypeTree(defn.FunctionSymbol(numArgs, isContextual, true, isImpure).typeRef)
-          .withSpan(tycon.span)
-        assignType(cpy.AppliedTypeTree(app)(tycon1, args), tycon1, args)
-      case _ =>
-        app
+    val funSym = defn.FunctionSymbol(numArgs, isContextual, isErased.contains(true), isImpure)
 
     /** Typechecks dependent function type with given parameters `params` */
     def typedDependent(params: List[untpd.ValDef])(using Context): Tree =
@@ -1330,7 +1313,10 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
         if funFlags.is(Given) then params.map(_.withAddedFlags(Given))
         else params
       val params2 = params1.map(fixThis.transformSub)
-      val appDef0 = untpd.DefDef(nme.apply, List(params2), body, EmptyTree).withSpan(tree.span)
+      val params3 = params2.zipWithConserve(isErased) { (arg, isErased) =>
+        if isErased then arg.withAddedFlags(Erased) else arg
+      }
+      val appDef0 = untpd.DefDef(nme.apply, List(params3), body, EmptyTree).withSpan(tree.span)
       index(appDef0 :: Nil)
       val appDef = typed(appDef0).asInstanceOf[DefDef]
       val mt = appDef.symbol.info.asInstanceOf[MethodType]
@@ -1339,7 +1325,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
       val resTpt = TypeTree(mt.nonDependentResultApprox).withSpan(body.span)
       val typeArgs = appDef.termParamss.head.map(_.tpt) :+ resTpt
       val tycon = TypeTree(funSym.typeRef)
-      val core = propagateErased(AppliedTypeTree(tycon, typeArgs))
+      val core = AppliedTypeTree(tycon, typeArgs)
       RefinedTypeTree(core, List(appDef), ctx.owner.asClass)
     end typedDependent
 
@@ -1348,17 +1334,22 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
         typedDependent(args.asInstanceOf[List[untpd.ValDef]])(
           using ctx.fresh.setOwner(newRefinedClassSymbol(tree.span)).setNewScope)
       case _ =>
-        propagateErased(
-          typed(cpy.AppliedTypeTree(tree)(untpd.TypeTree(funSym.typeRef), args :+ body), pt))
+        if isErased.contains(true) then
+          typedFunctionType(
+            desugar.makeErasedFunctionValDefs(tree.asInstanceOf[untpd.FunctionWithMods], pt),
+            pt
+          )
+        else
+          typed(cpy.AppliedTypeTree(tree)(untpd.TypeTree(funSym.typeRef), args :+ body), pt)
     }
   }
 
   def typedFunctionValue(tree: untpd.Function, pt: Type)(using Context): Tree = {
     val untpd.Function(params: List[untpd.ValDef] @unchecked, _) = tree: @unchecked
 
-    val isContextual = tree match {
-      case tree: untpd.FunctionWithMods => tree.mods.is(Given)
-      case _ => false
+    val (isContextual, isDefinedErased) = tree match {
+      case tree: untpd.FunctionWithMods => (tree.mods.is(Given), tree.isErased)
+      case _ => (false, tree.args.map(_ => false))
     }
 
     /** The function body to be returned in the closure. Can become a TypedSplice
@@ -1467,8 +1458,8 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
 
     /** Returns the type and whether the parameter is erased */
     def protoFormal(i: Int): (Type, Boolean) =
-      if (protoFormals.length == params.length) (protoFormals(i), protoIsErased(i))
-      else (errorType(WrongNumberOfParameters(protoFormals.length), tree.srcPos), false)
+      if (protoFormals.length == params.length) (protoFormals(i), protoIsErased(i) || isDefinedErased(i))
+      else (errorType(WrongNumberOfParameters(protoFormals.length), tree.srcPos), isDefinedErased(i))
 
     /** Is `formal` a product type which is elementwise compatible with `params`? */
     def ptIsCorrectProduct(formal: Type) =
@@ -3034,7 +3025,14 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
       else formals.map(untpd.TypeTree)
     }
 
-    val ifun = desugar.makeContextualFunction(paramTypes, tree, defn.isErasedFunctionType(pt))
+    val erasedParams = pt match {
+      case RefinedType(parent, nme.apply, mt) if defn.isErasedFunctionType(parent) =>
+        val companion = mt.asInstanceOf[MethodType].companion.asInstanceOf[ErasedMethodCompanion]
+        companion.isErased
+      case _ => paramTypes.map(_ => false)
+    }
+
+    val ifun = desugar.makeContextualFunction(paramTypes, tree, erasedParams)
     typr.println(i"make contextual function $tree / $pt ---> $ifun")
     typedFunctionValue(ifun, pt)
   }