Fix #9339: Register doc of extension methods

Author: nicolasstucki

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

@@ -530,7 +530,7 @@ class Namer { typer: Typer =>
    *  enter them into symbol table
    */
   def indexExpanded(origStat: Tree)(using Context): Context = {
-    def recur(stat: Tree): Context = stat match {
+    def recur(stat: Tree, origStat: Tree): Context = stat match {
       case pcl: PackageDef =>
         val pkg = createPackageSymbol(pcl.pid)
         index(pcl.stats)(using ctx.fresh.setOwner(pkg.moduleClass))
@@ -542,16 +542,23 @@ class Namer { typer: Typer =>
       case mdef: DefTree =>
         val sym = createSymbol(mdef)
         enterSymbol(sym)
-        setDocstring(sym, origStat)
+        val origDef = origStat match
+          case origStat: ExtMethods => origStat.methods.head
+          case _: DefTree => origStat
+        setDocstring(sym, origDef)
         addEnumConstants(mdef, sym)
         ctx
       case stats: Thicket =>
-        stats.toList.foreach(recur)
+        val origStats = origStat match
+          case ext: ExtMethods => ext.methods
+          case _ => Iterator.continually(origStat)
+        for (stat, origDef) <- stats.toList.zip(origStats) do
+          recur(stat, origDef)
         ctx
       case _ =>
         ctx
     }
-    recur(expanded(origStat))
+    recur(expanded(origStat), origStat)
   }
 
   /** Determines whether this field holds an enum constant. */

compiler/test/dotty/tools/repl/DocTests.scala

@@ -31,6 +31,23 @@ class DocTests extends ReplTest {
       assertEquals("doc", doc("new Foo"))
     }
 
+  @Test def docOfExtension1 =
+    eval("/** doc */ extension (x: Int) def foo = 0").andThen { implicit s =>
+      assertEquals("doc", doc("extension_foo"))
+    }
+
+  @Test def docOfExtension2 =
+    eval("extension (x: Int) /** doc */ def foo = 0").andThen { implicit s =>
+      assertEquals("doc", doc("extension_foo"))
+    }
+
+  @Test def docOfExtension3 =
+    eval("/** doc0 */ extension (x: Int) { /** doc1 */ def foo = 0; /** doc2 */ def bar = 0; def baz = 0 }").andThen { implicit s =>
+      assertEquals("doc1", doc("extension_foo"))
+      assertEquals("doc2", doc("extension_bar"))
+      assertEquals("doc0", doc("extension_baz"))
+    }
+
   @Test def docOfDefInObject =
     eval("object O { /** doc */ def foo = 0 }").andThen { implicit s =>
       assertEquals("doc", doc("O.foo"))