Allow cross references in collective extensions

Author: odersky

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

@@ -499,7 +499,7 @@ object Flags {
 
   /** Flags that can apply to a module val */
   val RetainedModuleValFlags: FlagSet = RetainedModuleValAndClassFlags |
-    Override | Final | Method | Implicit | Given | Lazy |
+    Override | Final | Method | Implicit | Given | Lazy | Extension |
     Accessor | AbsOverride | StableRealizable | Captured | Synchronized | Erased
 
   /** Flags that can apply to a module class */

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

@@ -220,4 +220,7 @@ class SymUtils(val self: Symbol) extends AnyVal {
   /** Is symbol a splice operation? */
   def isSplice(implicit ctx: Context): Boolean =
     self == defn.InternalQuoted_exprSplice || self == defn.QuotedType_splice
+
+  def isCollectiveExtensionClass(using Context): Boolean =
+    self.is(ModuleClass) && self.sourceModule.is(Extension, butNot = Method)
 }

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

@@ -401,6 +401,19 @@ class Typer extends Namer
     if (name == nme.ROOTPKG)
       return tree.withType(defn.RootPackage.termRef)
 
+    /** Convert a reference `f` to an extension method in a collective extension
+     *  on parameter `x` to `x.f`
+     */
+    def extensionMethodSelect(xmethod: Symbol): untpd.Tree =
+      val leadParamName = xmethod.info.paramNamess.head.head
+      def isLeadParam(sym: Symbol) =
+        sym.is(Param) && sym.owner.owner == xmethod.owner && sym.name == leadParamName
+      def leadParam(ctx: Context): Symbol =
+        ctx.scope.lookupAll(leadParamName).find(isLeadParam) match
+          case Some(param) => param
+          case None => leadParam(ctx.outersIterator.dropWhile(_.scope eq ctx.scope).next)
+      untpd.cpy.Select(tree)(untpd.ref(leadParam(ctx).termRef), name)
+
     val rawType = {
       val saved1 = unimported
       val saved2 = foundUnderScala2
@@ -441,8 +454,14 @@ class Typer extends Namer
         errorType(new MissingIdent(tree, kind, name.show), tree.sourcePos)
 
     val tree1 = ownType match {
-      case ownType: NamedType if !prefixIsElidable(ownType) =>
-        ref(ownType).withSpan(tree.span)
+      case ownType: NamedType =>
+        val sym = ownType.symbol
+        if sym.isAllOf(ExtensionMethod)
+            && sym.owner.isCollectiveExtensionClass
+            && ctx.owner.isContainedIn(sym.owner)
+        then typed(extensionMethodSelect(sym), pt)
+        else if prefixIsElidable(ownType) then tree.withType(ownType)
+        else ref(ownType).withSpan(tree.span)
       case _ =>
         tree.withType(ownType)
     }

tests/run/collective-extensions.scala

@@ -0,0 +1,23 @@
+extension on (x: String):
+  def foo(y: String): String = x ++ y
+  def bar(y: String): String = foo(y)
+  def baz(y: String): String =
+    val x = y
+    bar(x)
+  def bam(y: String): String = this.baz(x)(y)
+  def ban(foo: String): String = x + foo
+  def bao(y: String): String =
+    val bam = "ABC"
+    x ++ y ++ bam
+
+  def app(n: Int, suffix: String): String =
+    if n == 0 then x ++ suffix
+    else app(n - 1, suffix ++ suffix)
+
+@main def Test =
+  assert("abc".bar("def") == "abcdef")
+  assert("abc".baz("def") == "abcdef")
+  assert("abc".bam("def") == "abcdef")
+  assert("abc".ban("def") == "abcdef")
+  assert("abc".bao("def") == "abcdefABC")
+  assert("abc".app(3, "!") == "abc!!!!!!!!")