From e6e471383aa89fb9d70eef2ee7cd0f30eb577875 Mon Sep 17 00:00:00 2001
From: Katarzyna Marek <kasia@marek.net>
Date: Tue, 12 Sep 2023 13:48:41 +0200
Subject: [PATCH] bugfix: semantic tokens and highlight for extension method
 defined in companion object

---
 .../scala/meta/internal/pc/PcCollector.scala  | 42 ++++++++++++++++++-
 .../Scala3DocumentHighlightSuite.scala        | 11 +++++
 2 files changed, 52 insertions(+), 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
index adc1fffc540..fa84e74bdb9 100644
--- 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
index e9681f5d7a4..cd43049ade5 100644
--- 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,
+  )
+
 }