Fix #9339: Register doc of extension methods

Author: nicolasstucki

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

@@ -530,28 +530,29 @@ 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, idx: Int): Context = stat match {
       case pcl: PackageDef =>
         val pkg = createPackageSymbol(pcl.pid)
         index(pcl.stats)(using ctx.fresh.setOwner(pkg.moduleClass))
         invalidateCompanions(pkg, Trees.flatten(pcl.stats map expanded))
-        setDocstring(pkg, stat)
+        setDocstring(pkg, stat, idx)
         ctx
       case imp: Import =>
         ctx.importContext(imp, createSymbol(imp))
       case mdef: DefTree =>
         val sym = createSymbol(mdef)
         enterSymbol(sym)
-        setDocstring(sym, origStat)
+        setDocstring(sym, origStat, idx)
         addEnumConstants(mdef, sym)
         ctx
       case stats: Thicket =>
-        stats.toList.foreach(recur)
+        for (stat, i) <- stats.toList.zipWithIndex do
+          recur(stat, i)
         ctx
       case _ =>
         ctx
     }
-    recur(expanded(origStat))
+    recur(expanded(origStat), 0)
   }
 
   /** Determines whether this field holds an enum constant. */
@@ -590,9 +591,11 @@ class Namer { typer: Typer =>
     case _ =>
   }
 
-  def setDocstring(sym: Symbol, tree: Tree)(using Context): Unit = tree match {
+  def setDocstring(sym: Symbol, tree: Tree, idx: Int)(using Context): Unit = tree match {
     case t: MemberDef if t.rawComment.isDefined =>
       ctx.docCtx.foreach(_.addDocstring(sym, t.rawComment))
+    case t: ExtMethods if t.methods(idx).rawComment.isDefined =>
+      ctx.docCtx.foreach(_.addDocstring(sym, t.methods(idx).rawComment))
     case _ => ()
   }
 

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"))