bugfix: rename end marker (#18838)

metals backport scalameta/metals#5803

CC: @tgodzik

Author: tgodzik

presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala

@@ -69,7 +69,7 @@ abstract class PcCollector[T](
     case _ => rawPath
   def collect(
       parent: Option[Tree]
-  )(tree: Tree, pos: SourcePosition, symbol: Option[Symbol]): T
+  )(tree: Tree| EndMarker, pos: SourcePosition, symbol: Option[Symbol]): T
 
   /**
    * @return (adjusted position, should strip backticks)
@@ -423,7 +423,7 @@ abstract class PcCollector[T](
         parent: Option[Tree]
     ): Set[T] =
       def collect(
-          tree: Tree,
+          tree: Tree | EndMarker,
           pos: SourcePosition,
           symbol: Option[Symbol] = None
       ) =
@@ -461,6 +461,9 @@ abstract class PcCollector[T](
         case df: NamedDefTree
             if df.span.isCorrect && df.nameSpan.isCorrect &&
               filter(df) && !isGeneratedGiven(df) =>
+          def collectEndMarker =
+            EndMarker.getPosition(df, pos, sourceText).map:
+              collect(EndMarker(df.symbol), _)
           val annots = collectTrees(df.mods.annotations)
           val traverser =
             new PcCollector.DeepFolderWithParent[Set[T]](
@@ -470,7 +473,7 @@ abstract class PcCollector[T](
             occurrences + collect(
               df,
               pos.withSpan(df.nameSpan)
-            )
+            ) ++ collectEndMarker
           ) { case (set, tree) =>
             traverser(set, tree)
           }
@@ -635,3 +638,34 @@ case class ExtensionParamOccurence(
     sym: Symbol,
     methods: List[untpd.Tree]
 )
+
+case class EndMarker(symbol: Symbol)
+
+object EndMarker:
+  /**
+    * Matches end marker line from start to the name's beginning.
+    * E.g.
+    *    end /* some comment */
+    */
+  private val endMarkerRegex = """.*end(/\*.*\*/|\s)+""".r
+  def getPosition(df: NamedDefTree, pos: SourcePosition, sourceText: String)(
+      implicit ct: Context
+  ): Option[SourcePosition] =
+    val name = df.name.toString()
+    val endMarkerLine =
+      sourceText.slice(df.span.start, df.span.end).split('\n').last
+    val index = endMarkerLine.length() - name.length()
+    if index < 0 then None
+    else
+      val (possiblyEndMarker, possiblyEndMarkerName) =
+        endMarkerLine.splitAt(index)
+      Option.when(
+        possiblyEndMarkerName == name &&
+          endMarkerRegex.matches(possiblyEndMarker)
+      )(
+        pos
+          .withStart(df.span.end - name.length())
+          .withEnd(df.span.end)
+      )
+  end getPosition
+end EndMarker

presentation-compiler/src/main/dotty/tools/pc/PcDocumentHighlightProvider.scala

@@ -19,7 +19,7 @@ final class PcDocumentHighlightProvider(
   def collect(
       parent: Option[Tree]
   )(
-      tree: Tree,
+      tree: Tree | EndMarker,
       toAdjust: SourcePosition,
       sym: Option[Symbol]
   ): DocumentHighlight =

presentation-compiler/src/main/dotty/tools/pc/PcInlineValueProviderImpl.scala

@@ -24,24 +24,27 @@ import org.eclipse.lsp4j as l
 final class PcInlineValueProviderImpl(
     val driver: InteractiveDriver,
     val params: OffsetParams
-) extends PcCollector[Occurence](driver, params)
+) extends PcCollector[Option[Occurence]](driver, params)
     with InlineValueProvider:
 
   val text = params.text.toCharArray()
 
   val position: l.Position = pos.toLsp.getStart()
 
   override def collect(parent: Option[Tree])(
-      tree: Tree,
+      tree: Tree | EndMarker,
       pos: SourcePosition,
       sym: Option[Symbol]
-  ): Occurence =
-    val (adjustedPos, _) = adjust(pos)
-    Occurence(tree, parent, adjustedPos)
+  ): Option[Occurence] =
+    tree match
+      case tree: Tree =>
+        val (adjustedPos, _) = adjust(pos)
+        Some(Occurence(tree, parent, adjustedPos))
+      case _ => None
 
   override def defAndRefs(): Either[String, (Definition, List[Reference])] =
     val newctx = driver.currentCtx.fresh.setCompilationUnit(unit)
-    val allOccurences = result()
+    val allOccurences = result().flatten
     for
       definition <- allOccurences
         .collectFirst { case Occurence(defn: ValDef, _, pos) =>

presentation-compiler/src/main/dotty/tools/pc/PcRenameProvider.scala

@@ -34,7 +34,7 @@ final class PcRenameProvider(
 
   def collect(
       parent: Option[Tree]
-  )(tree: Tree, toAdjust: SourcePosition, sym: Option[Symbol]): l.TextEdit =
+  )(tree: Tree | EndMarker, toAdjust: SourcePosition, sym: Option[Symbol]): l.TextEdit =
     val (pos, stripBackticks) = adjust(toAdjust, forRename = true)
     l.TextEdit(
       pos.toLsp,

presentation-compiler/src/main/dotty/tools/pc/PcSemanticTokensProvider.scala

@@ -32,7 +32,7 @@ final class PcSemanticTokensProvider(
    * 3. type parameters,
    * In all those cases we don't have a specific value for sure.
    */
-  private def isDeclaration(tree: Tree) = tree match
+  private def isDeclaration(tree: Tree | EndMarker) = tree match
     case df: ValOrDefDef => df.rhs.isEmpty
     case df: TypeDef =>
       df.rhs match
@@ -49,7 +49,8 @@ final class PcSemanticTokensProvider(
    * that the compiler sees them as vals, as it's not clear
    * if they should be declaration/definition at all.
    */
-  private def isDefinition(tree: Tree) = tree match
+  private def isDefinition(tree: Tree | EndMarker) = tree match
+    case _: EndMarker => true
     case df: Bind => true
     case df: ValOrDefDef =>
       !df.rhs.isEmpty && !df.symbol.isAllOf(Flags.EnumCase)
@@ -62,8 +63,12 @@ final class PcSemanticTokensProvider(
   object Collector extends PcCollector[Option[Node]](driver, params):
     override def collect(
         parent: Option[Tree]
-    )(tree: Tree, pos: SourcePosition, symbol: Option[Symbol]): Option[Node] =
-      val sym = symbol.fold(tree.symbol)(identity)
+    )(tree: Tree | EndMarker, pos: SourcePosition, symbol: Option[Symbol]): Option[Node] =
+      val sym =
+         tree match
+           case tree: Tree =>
+             symbol.fold(tree.symbol)(identity)
+           case EndMarker(sym) => sym
       if !pos.exists || sym == null || sym == NoSymbol then None
       else
         Some(

presentation-compiler/test/dotty/tools/pc/tests/edit/PcRenameSuite.scala

@@ -478,9 +478,33 @@ class PcRenameSuite extends BasePcRenameSuite:
          |  def <<ma@@p>>(f: Int => Int): Bar = Bar(x.map(f))
          |}
          |
-         |val f = 
+         |val f =
          |  for {
          |    b <- Bar(List(1,2,3))
          |  } yield b
-         |""".stripMargin,
+         |""".stripMargin
+    )
+
+  @Test def `end-marker` =
+    check(
+      """|def <<he@@llo>>(a: Int) =
+        |  ???
+        |end <<hello>>
+        |""".stripMargin
+    )
+
+  @Test def `end-marker-with-comment` =
+    check(
+      """|def <<he@@llo>>(a: Int) =
+        |  ???
+        |end /* a comment */ <<hello>> /* a comment */
+        |""".stripMargin
+    )
+
+  @Test def `end-marker-wrong` =
+    check(
+      """|def <<f@@oo>> =
+        |  def bar =
+        |    ???
+        |  end bar""".stripMargin
     )

presentation-compiler/test/dotty/tools/pc/tests/tokens/SemanticTokensSuite.scala

@@ -343,3 +343,11 @@ class SemanticTokensSuite extends BaseSemanticTokensSuite:
            |}
           |""".stripMargin,
     )
+
+  @Test def `end-marker` =
+    check(
+        """|def <<foo>>/*method,definition*/ =
+           |  1
+           |end <<foo>>/*method,definition*/
+           |""".stripMargin,
+    )

presentation-compiler/test/dotty/tools/pc/utils/DefSymbolCollector.scala

@@ -8,14 +8,15 @@ import dotty.tools.dotc.core.Symbols.*
 import dotty.tools.dotc.interactive.InteractiveDriver
 import dotty.tools.dotc.util.SourcePosition
 import dotty.tools.pc.PcCollector
+import dotty.tools.pc.EndMarker
 
 final class DefSymbolCollector(
     driver: InteractiveDriver,
     params: VirtualFileParams
 ) extends PcCollector[Option[Symbol]](driver, params):
 
   def collect(parent: Option[Tree])(
-      tree: Tree,
+      tree: Tree | EndMarker,
       toAdjust: SourcePosition,
       sym: Option[Symbol]
   ): Option[Symbol] =