Refactor extension desugaring

Author: odersky

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

@@ -859,21 +859,22 @@ object desugar {
    *     given object name extends parents { self => body' }
    *
    *  where every definition in `body` is expanded to an extension method
-   *  taking type parameters `tparams` and a leading parameter `(x: T)`.
-   *  See: makeExtensionDef
+   *  taking type parameters `tparams` and a leading paramter `(x: T)`.
+   *  See: collectiveExtensionBody
    */
   def moduleDef(mdef: ModuleDef)(implicit ctx: Context): Tree = {
     val impl = mdef.impl
     val mods = mdef.mods
     impl.constr match {
-      case DefDef(_, tparams, (vparams @ (vparam :: Nil)) :: givenParamss, _, _) =>
+      case DefDef(_, tparams, vparamss @ (vparam :: Nil) :: givenParamss, _, _) =>
+        // Transform collective extension
         assert(mods.is(Given))
         return moduleDef(
           cpy.ModuleDef(mdef)(
             mdef.name,
             cpy.Template(impl)(
               constr = emptyConstructor,
-              body = impl.body.map(makeExtensionDef(_, tparams, vparams, givenParamss)))))
+              body = collectiveExtensionBody(impl.body, tparams, vparamss))))
       case _ =>
     }
 
@@ -916,43 +917,40 @@ object desugar {
     }
   }
 
-  /** Given tpe parameters `Ts` (possibly empty) and a leading value parameter `(x: T)`,
-   *  map a method definition
+  /** Transform the statements of a collective extension
+   *   @param stats    the original statements as they were parsed
+   *   @param tparams  the collective type parameters
+   *   @param vparamss the collective value parameters, consisting
+   *                   of a single leading value parameter, followed by
+   *                   zero or more context parameter clauses
    *
-   *     def foo [Us] paramss ...
+   *  Note: It is already assured by Parser.checkExtensionMethod that all
+   *  statements conform to requirements.
    *
-   *  to
+   *  Each method in stats is transformed into an extension method. Example:
+   *
+   *    extension on [Ts](x: T)(using C):
+   *      def f(y: T) = ???
+   *      def g(z: T) = f(z)
    *
-   *     <extension> def foo[Ts ++ Us](x: T) parammss ...
+   *  is turned into
    *
-   *  If the given member `mdef` is not of this form, flag it as an error.
+   *    extension:
+   *      <extension> def f[Ts](x: T)(using C)(y: T) = ???
+   *      <extension> def g[Ts](x: T)(using C)(z: T) = f(z)
    */
-
-  def makeExtensionDef(mdef: Tree, tparams: List[TypeDef], leadingParams: List[ValDef],
-                       givenParamss: List[List[ValDef]])(using ctx: Context): Tree = {
-    mdef match {
-      case mdef: DefDef =>
-        if (mdef.mods.is(Extension)) {
-          ctx.error(NoExtensionMethodAllowed(mdef), mdef.sourcePos)
-          mdef
-        } else {
-          if (tparams.nonEmpty && mdef.tparams.nonEmpty) then
-            ctx.error(ExtensionMethodCannotHaveTypeParams(mdef), mdef.tparams.head.sourcePos)
-            mdef
-          else cpy.DefDef(mdef)(
+  def collectiveExtensionBody(stats: List[Tree],
+      tparams: List[TypeDef], vparamss: List[List[ValDef]])(using ctx: Context): List[Tree] =
+    for stat <- stats yield
+      stat match
+        case mdef: DefDef =>
+          cpy.DefDef(mdef)(
             tparams = tparams ++ mdef.tparams,
-            vparamss = leadingParams :: givenParamss ::: mdef.vparamss
+            vparamss = vparamss ::: mdef.vparamss,
           ).withMods(mdef.mods | Extension)
-        }
-      case mdef: Import =>
-        mdef
-      case mdef if !mdef.isEmpty => {
-        ctx.error(ExtensionCanOnlyHaveDefs(mdef), mdef.sourcePos)
-        mdef
-      }
-      case mdef => mdef
-    }
-  }
+        case mdef =>
+          mdef
+  end collectiveExtensionBody
 
   /** Transforms
    *