From 5ded2e3ecb3a8ffb88c9893f5bfe818666bc89ea Mon Sep 17 00:00:00 2001
From: Jakub Ciesluk <323892@uwr.edu.pl>
Date: Wed, 6 Dec 2023 09:16:55 +0100
Subject: [PATCH] bugfix: Document highlight on class constructors

[Cherry-picked 33011d97025c18cb7d0f5bb18ba213790256008e]
---
 .../src/main/dotty/tools/pc/PcCollector.scala | 37 ++++++++
 .../highlight/DocumentHighlightSuite.scala    | 94 +++++++++++++++++++
 .../pc/tests/tokens/SemanticTokensSuite.scala | 37 ++++++++
 3 files changed, 168 insertions(+)

diff --git a/presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala b/presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala
index 89bfd05a7386..310edd60d87e 100644
--- a/presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala
+++ b/presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala
@@ -144,6 +144,22 @@ abstract class PcCollector[T](
           if id.symbol
             .is(Flags.Param) && id.symbol.owner.is(Flags.ExtensionMethod) =>
         Some(findAllExtensionParamSymbols(id.sourcePos, id.name, id.symbol))
+      /**
+       * Workaround for missing symbol in:
+       * class A[T](a: T)
+       * val x = new <<A>>(1)
+       */
+      case t :: (n: New) :: (sel: Select) :: _
+          if t.symbol == NoSymbol && sel.symbol.isConstructor =>
+        Some(symbolAlternatives(sel.symbol.owner), namePos(t))
+      /**
+       * Workaround for missing symbol in:
+       * class A[T](a: T)
+       * val x = <<A>>[Int](1)
+       */
+      case (sel @ Select(New(t), _)) :: (_: TypeApply) :: _
+          if sel.symbol.isConstructor =>
+        Some(symbolAlternatives(sel.symbol.owner), namePos(t))
       /* simple identifier:
        * val a = val@@ue + value
        */
@@ -403,6 +419,22 @@ abstract class PcCollector[T](
               ident.sourcePos
             )
           else occurrences
+        /**
+         * Workaround for missing symbol in:
+         * class A[T](a: T)
+         * val x = new <<A>>(1)
+         */
+        case sel @ Select(New(t), _)
+            if sel.span.isCorrect &&
+              sel.symbol.isConstructor &&
+              t.symbol == NoSymbol =>
+          if soughtFilter(_ == sel.symbol.owner) then
+            occurrences + collect(
+              sel,
+              namePos(t),
+              Some(sel.symbol.owner),
+            )
+          else occurrences
         /**
          * All select statements such as:
          * val a = hello.<<b>>
@@ -564,6 +596,11 @@ abstract class PcCollector[T](
         Span(span.start, span.start + realName.length, point)
       else Span(point, span.end, point)
     else span
+
+  private def namePos(tree: Tree): SourcePosition =
+    tree match
+      case sel: Select => sel.sourcePos.withSpan(selectNameSpan(sel))
+      case _ => tree.sourcePos
 end PcCollector
 
 object PcCollector:
diff --git a/presentation-compiler/test/dotty/tools/pc/tests/highlight/DocumentHighlightSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/highlight/DocumentHighlightSuite.scala
index 80451fc67562..1dc2a7b156db 100644
--- a/presentation-compiler/test/dotty/tools/pc/tests/highlight/DocumentHighlightSuite.scala
+++ b/presentation-compiler/test/dotty/tools/pc/tests/highlight/DocumentHighlightSuite.scala
@@ -1185,3 +1185,97 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
          |val a = MyIntOut(1) <<+@@+>> 3
          |""".stripMargin,
     )
+
+  @Test def `constructor` = 
+    check(
+      """
+        |object Main {
+        |  class <<A@@bc>>[T](abc: T)
+        |  val x = new <<Abc>>(123)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor1` = 
+    check(
+      """
+        |object Main {
+        |  case class <<Abc>>[T](abc: T)
+        |  val x = <<A@@bc>>(123)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor2` = 
+    check(
+      """
+        |object Main {
+        |  class <<A@@bc>>[T](abc: T)
+        |  object <<Abc>>
+        |  val x = new <<Abc>>(123)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor3` = 
+    check(
+      """
+        |object Main {
+        |  class <<Abc>>[T](abc: T)
+        |  object <<Abc>>
+        |  val x = new <<A@@bc>>(123)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor4` = 
+    check(
+      """
+        |object Main {
+        |  class <<Abc>>[T](abc: T)
+        |  object <<Ab@@c>>
+        |  val x = new <<Abc>>(123)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor5` = 
+    check(
+      """
+        |object Main {
+        |  class <<Abc>>[T](abc: T)
+        |  object <<Abc>> {
+        |    def apply(abc: Int, bde: Int) = new <<Abc>>(abc + bde)
+        |  }
+        |  val x = <<Ab@@c>>(123, 456)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor6` = 
+    check(
+      """
+        |class <<Abc>>[T](a: T)
+        |object O {
+        |  def foo(a: Int) = new <<Abc>>[Int](a)
+        |  val x = <<Ab@@c>>[Int](2)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor7` = 
+    check(
+      """
+        |object Bar {
+        |class <<Abc>>[T](a: T)
+        |}
+        |
+        |object O {
+        |  val x = new Bar.<<Ab@@c>>(2)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor8` = 
+    check(
+      """
+        |object Bar {
+        |class <<Abc>>[T](a: T)
+        |}
+        |
+        |object O {
+        |  val x = Bar.<<Ab@@c>>[Int](2)
+        |}""".stripMargin
+    )
diff --git a/presentation-compiler/test/dotty/tools/pc/tests/tokens/SemanticTokensSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/tokens/SemanticTokensSuite.scala
index d4ddefcc11d0..c8cfac044732 100644
--- a/presentation-compiler/test/dotty/tools/pc/tests/tokens/SemanticTokensSuite.scala
+++ b/presentation-compiler/test/dotty/tools/pc/tests/tokens/SemanticTokensSuite.scala
@@ -375,3 +375,40 @@ class SemanticTokensSuite extends BaseSemanticTokensSuite:
            |end <<foo>>/*method,definition*/
            |""".stripMargin,
     )
+
+  @Test def `constructor` =
+    check(
+      """
+        |object <<Bar>>/*class*/ {
+        |  class <<Abc>>/*class*/[<<T>>/*typeParameter,definition,abstract*/](<<a>>/*variable,declaration,readonly*/: <<T>>/*typeParameter,abstract*/)
+        |}
+        |
+        |object <<O>>/*class*/ {
+        |  val <<x>>/*variable,definition,readonly*/ = new <<Bar>>/*class*/.<<Abc>>/*class*/(2)
+        |  val <<y>>/*variable,definition,readonly*/ = new <<Bar>>/*class*/.<<Abc>>/*class*/[<<Int>>/*class,abstract*/](2)
+        |  val <<z>>/*variable,definition,readonly*/ = <<Bar>>/*class*/.<<Abc>>/*class*/(2)
+        |  val <<w>>/*variable,definition,readonly*/ = <<Bar>>/*class*/.<<Abc>>/*class*/[<<Int>>/*class,abstract*/](2)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor1` =
+    check(
+      """
+        |object <<Main>>/*class*/ {
+        |  class <<Abc>>/*class*/[<<T>>/*typeParameter,definition,abstract*/](<<abc>>/*variable,declaration,readonly*/: <<T>>/*typeParameter,abstract*/)
+        |  object <<Abc>>/*class*/
+        |  val <<x>>/*variable,definition,readonly*/ = new <<Abc>>/*class*/(123)
+        |}""".stripMargin
+    )
+
+  @Test def `constructor2` =
+    check(
+      """
+        |object <<Main>>/*class*/ {
+        |  class <<Abc>>/*class*/[<<T>>/*typeParameter,definition,abstract*/](<<abc>>/*variable,declaration,readonly*/: <<T>>/*typeParameter,abstract*/)
+        |  object <<Abc>>/*class*/ {
+        |    def <<apply>>/*method,definition*/[<<T>>/*typeParameter,definition,abstract*/](<<abc>>/*parameter,declaration,readonly*/: <<T>>/*typeParameter,abstract*/, <<bde>>/*parameter,declaration,readonly*/: <<T>>/*typeParameter,abstract*/) = new <<Abc>>/*class*/(<<abc>>/*parameter,readonly*/)
+        |  }
+        |  val <<x>>/*variable,definition,readonly*/ = <<Abc>>/*class*/(123, 456)
+        |}""".stripMargin
+    )
\ No newline at end of file