Unified extension methods

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

@@ -559,7 +559,7 @@ object desugar {
     val copiedAccessFlags = if migrateTo3 then EmptyFlags else AccessFlags
 
     // Methods to add to a case class C[..](p1: T1, ..., pN: Tn)(moreParams)
-    //     def _1: T1 = this.p1 
+    //     def _1: T1 = this.p1
     //     ...
     //     def _N: TN = this.pN (unless already given as valdef or parameterless defdef)
     //     def copy(p1: T1 = p1: @uncheckedVariance, ...,
@@ -572,7 +572,7 @@ object desugar {
     val caseClassMeths = {
       def syntheticProperty(name: TermName, tpt: Tree, rhs: Tree) =
         DefDef(name, Nil, Nil, tpt, rhs).withMods(synthetic)
-        
+
       def productElemMeths =
         val caseParams = derivedVparamss.head.toArray
         val selectorNamesInBody = normalizedBody.collect {
@@ -850,6 +850,7 @@ object desugar {
    *  where every definition in `body` is expanded to an extension method
    *  taking type parameters `tparams` and a leading paramter `(x: T)`.
    *  See: collectiveExtensionBody
+   *  TODO: drop this part
    */
   def moduleDef(mdef: ModuleDef)(implicit ctx: Context): Tree = {
     val impl = mdef.impl
@@ -906,6 +907,25 @@ object desugar {
     }
   }
 
+  /** Transform extension construct to list of extension methods */
+  def extMethods(ext: ExtMethods)(using Context): Tree = flatTree {
+    for mdef <- ext.methods yield
+      if mdef.tparams.nonEmpty then
+        ctx.error("extension method cannot have type parameters here, all type parameters go after `extension`",
+          mdef.tparams.head.sourcePos)
+      defDef(
+        cpy.DefDef(mdef)(
+          name = mdef.name.toExtensionName,
+          tparams = ext.tparams ++ mdef.tparams,
+          vparamss = mdef.vparamss match
+            case vparams1 :: vparamss1 if !isLeftAssoc(mdef.name) =>
+              vparams1 :: ext.vparamss ::: vparamss1
+            case _ =>
+              ext.vparamss ++ mdef.vparamss
+        ).withMods(mdef.mods | Extension)
+      )
+  }
+
   /** Transform the statements of a collective extension
    *   @param stats    the original statements as they were parsed
    *   @param tparams  the collective type parameters
@@ -934,6 +954,7 @@ object desugar {
       stat match
         case mdef: DefDef =>
           cpy.DefDef(mdef)(
+            name = mdef.name.toExtensionName,
             tparams = tparams ++ mdef.tparams,
             vparamss = vparamss ::: mdef.vparamss,
           ).withMods(mdef.mods | Extension)
@@ -964,16 +985,21 @@ object desugar {
   /** The normalized name of `mdef`. This means
    *   1. Check that the name does not redefine a Scala core class.
    *      If it does redefine, issue an error and return a mangled name instead of the original one.
-   *   2. If the name is missing (this can be the case for instance definitions), invent one instead.
+   *   2. Check that the name does not start with `extension_` unless the
+   *      method is an extension method.
+   *   3. If the name is missing (this can be the case for instance definitions), invent one instead.
    */
   def normalizeName(mdef: MemberDef, impl: Tree)(implicit ctx: Context): Name = {
     var name = mdef.name
     if (name.isEmpty) name = name.likeSpaced(inventGivenOrExtensionName(impl))
+    def errPos = mdef.source.atSpan(mdef.nameSpan)
     if (ctx.owner == defn.ScalaPackageClass && defn.reservedScalaClassNames.contains(name.toTypeName)) {
       val kind = if (name.isTypeName) "class" else "object"
-      ctx.error(IllegalRedefinitionOfStandardKind(kind, name), mdef.sourcePos)
+      ctx.error(IllegalRedefinitionOfStandardKind(kind, name), errPos)
       name = name.errorName
     }
+    if name.isExtensionName && !mdef.mods.is(Extension) then
+      ctx.error(em"illegal method name: $name may not start with `extension_`", errPos)
     name
   }
 
@@ -1247,7 +1273,7 @@ object desugar {
   }
 
   private def isTopLevelDef(stat: Tree)(using Context): Boolean = stat match
-    case _: ValDef | _: PatDef | _: DefDef | _: Export => true
+    case _: ValDef | _: PatDef | _: DefDef | _: Export | _: ExtMethods => true
     case stat: ModuleDef =>
       stat.mods.isOneOf(GivenOrImplicit)
     case stat: TypeDef =>

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

@@ -114,6 +114,7 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   case class ContextBounds(bounds: TypeBoundsTree, cxBounds: List[Tree])(implicit @constructorOnly src: SourceFile) extends TypTree
   case class PatDef(mods: Modifiers, pats: List[Tree], tpt: Tree, rhs: Tree)(implicit @constructorOnly src: SourceFile) extends DefTree
   case class Export(expr: Tree, selectors: List[ImportSelector])(implicit @constructorOnly src: SourceFile) extends Tree
+  case class ExtMethods(tparams: List[TypeDef], vparamss: List[List[ValDef]], methods: List[DefDef])(implicit @constructorOnly src: SourceFile) extends Tree
   case class MacroTree(expr: Tree)(implicit @constructorOnly src: SourceFile) extends Tree
 
   case class ImportSelector(imported: Ident, renamed: Tree = EmptyTree, bound: Tree = EmptyTree)(implicit @constructorOnly src: SourceFile) extends Tree {
@@ -617,6 +618,9 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
       case tree: Export if (expr eq tree.expr) && (selectors eq tree.selectors) => tree
       case _ => finalize(tree, untpd.Export(expr, selectors)(tree.source))
     }
+    def ExtMethods(tree: Tree)(tparams: List[TypeDef], vparamss: List[List[ValDef]], methods: List[DefDef])(using Context): Tree = tree match
+      case tree: ExtMethods if (tparams eq tree.tparams) && (vparamss eq tree.vparamss) && (methods == tree.methods) => tree
+      case _ => finalize(tree, untpd.ExtMethods(tparams, vparamss, methods)(tree.source))
     def ImportSelector(tree: Tree)(imported: Ident, renamed: Tree, bound: Tree)(implicit ctx: Context): Tree = tree match {
       case tree: ImportSelector if (imported eq tree.imported) && (renamed eq tree.renamed) && (bound eq tree.bound) => tree
       case _ => finalize(tree, untpd.ImportSelector(imported, renamed, bound)(tree.source))
@@ -683,6 +687,8 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
         cpy.PatDef(tree)(mods, transform(pats), transform(tpt), transform(rhs))
       case Export(expr, selectors) =>
         cpy.Export(tree)(transform(expr), selectors)
+      case ExtMethods(tparams, vparamss, methods) =>
+        cpy.ExtMethods(tree)(transformSub(tparams), vparamss.mapConserve(transformSub(_)), transformSub(methods))
       case ImportSelector(imported, renamed, bound) =>
         cpy.ImportSelector(tree)(transformSub(imported), transform(renamed), transform(bound))
       case Number(_, _) | TypedSplice(_) =>
@@ -742,6 +748,8 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
         this(this(this(x, pats), tpt), rhs)
       case Export(expr, _) =>
         this(x, expr)
+      case ExtMethods(tparams, vparamss, methods) =>
+        this(vparamss.foldLeft(this(x, tparams))(apply), methods)
       case ImportSelector(imported, renamed, bound) =>
         this(this(this(x, imported), renamed), bound)
       case Number(_, _) =>

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

@@ -131,6 +131,21 @@ object NameOps {
       else name.toTermName
     }
 
+    /** Does this name start with `extension_`? */
+    def isExtensionName: Boolean = name match
+      case name: SimpleName => name.startsWith("extension_")
+      case _ => false
+
+    /** Add an `extension_` in front of this name */
+    def toExtensionName = termName("extension_" ++ name.toString)
+
+    /** Drop `extension_` in front of this name, if it has this prefix */
+    def dropExtension = name match
+      case name: SimpleName if name.startsWith("extension_") =>
+        name.drop("extension_".length)
+      case _ =>
+        name
+
     /** The expanded name.
      *  This is the fully qualified name of `base` with `ExpandPrefixName` as separator,
      *  followed by `kind` and the name.

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

@@ -1190,6 +1190,15 @@ object SymDenotations {
       else if (this.exists) owner.enclosingMethod
       else NoSymbol
 
+    /** The closest enclosing extension method containing this definition,
+     *  provided the extension method appears in the same class.
+     */
+    final def enclosingExtensionMethod(using Context): Symbol =
+      if this.isAllOf(ExtensionMethod) then symbol
+      else if this.isClass then NoSymbol
+      else if this.exists then owner.enclosingExtensionMethod
+      else NoSymbol
+
     /** The top-level class containing this denotation,
      *  except for a toplevel module, where its module class is returned.
      */