bugfix: semantic tokens and highlight for extension method defined in…

… companion object
scalameta · Sep 12, 2023 · e6e4713 · e6e4713
1 parent 3818161
commit e6e4713
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 1 deletion.
diff --git a/mtags/src/main/scala-3/scala/meta/internal/pc/PcCollector.scala b/mtags/src/main/scala-3/scala/meta/internal/pc/PcCollector.scala
@@ -258,6 +258,22 @@ abstract class PcCollector[T](
           .selector(pos.span)
           .map(sym => (symbolAlternatives(sym), sym.sourcePos))
 
+      /* Workaround for missing span in:
+       * class MyIntOut(val value: Int)
+       * object MyIntOut:
+       *   extension (i: MyIntOut) def <<uneven>> = i.value % 2 == 1
+       *
+       * val a = MyIntOut(1).un@@even
+       */
+      case Apply(sel: Select, _) :: _
+          if sel.span.isZeroExtent && sel.symbol.is(Flags.ExtensionMethod) =>
+        untypedPath match
+          case (untypedSel: untpd.Select) :: _ =>
+            Some(
+              symbolAlternatives(sel.symbol),
+              pos.withSpan(untypedSel.nameSpan),
+            )
+          case _ => None
       case _ => None
 
     sought match
@@ -266,9 +282,11 @@ abstract class PcCollector[T](
 
   end soughtSymbols
 
-  lazy val extensionMethods =
+  lazy val untypedPath =
     NavigateAST
       .untypedPath(pos.span)(using compilatonUnitContext)
+  lazy val extensionMethods =
+    untypedPath
       .collectFirst { case em @ ExtMethods(_, _) => em }
 
   private def findAllExtensionParamSymbols(
@@ -448,6 +466,28 @@ abstract class PcCollector[T](
             sel,
             pos.withSpan(selectNameSpan(sel)),
           )
+        /* Workaround for missing span in:
+         * class MyIntOut(val value: Int)
+         * object MyIntOut:
+         *   extension (i: MyIntOut) def <<uneven>> = i.value % 2 == 1
+         *
+         * val a = MyIntOut(1).un@@even
+         */
+        case sel: Select
+            if sel.span.isZeroExtent && sel.symbol.is(Flags.ExtensionMethod) =>
+          scala.util
+            .Try(NavigateAST.untypedPath(sel.span)(using compilatonUnitContext))
+            .toOption match
+            case Some((_: Ident) :: (untypedSel: untpd.Select) :: _)
+                if untypedSel.span.isCorrect =>
+              val amendedSelect = sel.withSpan(untypedSel.span)
+              if filter(amendedSelect) then
+                occurences + collect(
+                  amendedSelect,
+                  pos.withSpan(untypedSel.nameSpan),
+                )
+              else occurences
+            case _ => occurences
         /* all definitions:
          * def <<foo>> = ???
          * class <<Foo>> = ???

diff --git a/tests/cross/src/test/scala/tests/highlight/Scala3DocumentHighlightSuite.scala b/tests/cross/src/test/scala/tests/highlight/Scala3DocumentHighlightSuite.scala
@@ -308,4 +308,15 @@ class Scala3DocumentHighlightSuite extends BaseDocumentHighlightSuite {
        |""".stripMargin,
   )
 
+  check(
+    "i5630".tag(IgnoreForScala3CompilerPC),
+    """|class MyIntOut(val value: Int)
+       |object MyIntOut:
+       |  extension (i: MyIntOut) def <<uneven>> = i.value % 2 == 1
+       |
+       |val a = MyIntOut(1)
+       |val m = a.<<un@@even>>
+       |""".stripMargin,
+  )
+
 }