bugfix: highlight for enum type params

Author: kasiaMarek

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

@@ -111,6 +111,33 @@ abstract class PcCollector[T](
   end adjust
 
   def symbolAlternatives(sym: Symbol) =
+    def member(parent: Symbol) = parent.info.member(sym.name).symbol
+    def primaryConstructorTypeParam(owner: Symbol) =
+      for
+        typeParams <- owner.primaryConstructor.paramSymss.headOption
+        param <- typeParams.find(_.name == sym.name)
+        if (param.isType)
+      yield param
+    def additionalForEnumTypeParam(enumClass: Symbol) =
+      if enumClass.is(Flags.Enum) then
+        val enumOwner =
+          if enumClass.is(Flags.Case)
+          then
+            Option.when(member(enumClass).is(Flags.Synthetic))(
+              enumClass.maybeOwner.companionClass
+            )
+          else Some(enumClass)
+        enumOwner.toSet.flatMap { enumOwner =>
+          val symsInEnumCases = enumOwner.children.toSet.flatMap(enumCase =>
+            if member(enumCase).is(Flags.Synthetic)
+            then primaryConstructorTypeParam(enumCase)
+            else None
+          )
+          val symsInEnumOwner =
+            primaryConstructorTypeParam(enumOwner).toSet + member(enumOwner)
+          symsInEnumCases ++ symsInEnumOwner
+        }
+      else Set.empty
     val all =
       if sym.is(Flags.ModuleClass) then
         Set(sym, sym.companionModule, sym.companionModule.companion)
@@ -129,7 +156,11 @@ abstract class PcCollector[T](
         ) ++ sym.allOverriddenSymbols.toSet
       // type used in primary constructor will not match the one used in the class
       else if sym.isTypeParam && sym.owner.isPrimaryConstructor then
-        Set(sym, sym.owner.owner.info.member(sym.name).symbol)
+        Set(sym, member(sym.maybeOwner.maybeOwner))
+          ++ additionalForEnumTypeParam(sym.maybeOwner.maybeOwner)
+      else if sym.isTypeParam then
+        primaryConstructorTypeParam(sym.maybeOwner).toSet
+          ++ additionalForEnumTypeParam(sym.maybeOwner) + sym
       else Set(sym)
     all.filter(s => s != NoSymbol && !s.isError)
   end symbolAlternatives
@@ -409,7 +440,7 @@ abstract class PcCollector[T](
          * All select statements such as:
          * val a = hello.<<b>>
          */
-        case sel: Select 
+        case sel: Select
           if sel.span.isCorrect && filter(sel) &&
             !isForComprehensionMethod(sel) =>
           occurrences + collect(

presentation-compiler/test/dotty/tools/pc/tests/highlight/DocumentHighlightSuite.scala

@@ -758,8 +758,8 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
         |  }
         |}""".stripMargin
     )
-  
-  @Test def `for-comp-map` = 
+
+  @Test def `for-comp-map` =
     check(
       """|object Main {
          |  val x = List(1).<<m@@ap>>(_ + 1)
@@ -770,7 +770,7 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
          |""".stripMargin,
     )
 
-  @Test def `for-comp-map1` = 
+  @Test def `for-comp-map1` =
     check(
       """|object Main {
          |  val x = List(1).<<m@@ap>>(_ + 1)
@@ -782,7 +782,7 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
          |""".stripMargin,
     )
 
-  @Test def `for-comp-foreach` = 
+  @Test def `for-comp-foreach` =
     check(
       """|object Main {
          |  val x = List(1).<<for@@each>>(_ => ())
@@ -793,7 +793,7 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
          |""".stripMargin,
     )
 
-  @Test def `for-comp-withFilter` = 
+  @Test def `for-comp-withFilter` =
     check(
       """|object Main {
          |  val x = List(1).<<with@@Filter>>(_ => true)
@@ -805,7 +805,7 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
          |""".stripMargin,
     )
 
-  @Test def `for-comp-withFilter1` = 
+  @Test def `for-comp-withFilter1` =
     check(
       """|object Main {
          |  val x = List(1).withFilter(_ => true).<<m@@ap>>(_ + 1)
@@ -817,7 +817,7 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
          |""".stripMargin,
     )
 
-  @Test def `for-comp-flatMap1` = 
+  @Test def `for-comp-flatMap1` =
     check(
       """|object Main {
          |  val x = List(1).<<flat@@Map>>(_ => List(1))
@@ -830,7 +830,7 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
          |""".stripMargin,
     )
 
-  @Test def `for-comp-flatMap2` = 
+  @Test def `for-comp-flatMap2` =
     check(
       """|object Main {
          |  val x = List(1).withFilter(_ => true).<<flat@@Map>>(_ => List(1))
@@ -1102,3 +1102,44 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
          |val alpha = MyOption.<<MySome>>(1)
          |""".stripMargin,
     )
+
+  @Test def `type-params-in-enum` =
+    check(
+      """|enum MyOption[+<<A@@A>>]:
+         |  case MySome(value: <<AA>>)
+         |  case MyNone
+         |""".stripMargin,
+    )
+
+  @Test def `type-params-in-enum2` =
+    check(
+      """|enum MyOption[+<<AA>>]:
+         |  case MySome(value: <<A@@A>>)
+         |  case MyNone
+         |""".stripMargin,
+    )
+
+  @Test def `type-params-in-enum3` =
+    check(
+      """|enum MyOption[<<AA>>](v: <<AA>>):
+         |  def get: <<A@@A>> = ???
+         |  case MySome[AA](value: AA) extends MyOption[Int](1)
+         |""".stripMargin,
+    )
+
+  @Test def `type-params-in-enum4` =
+    check(
+      """|enum MyOption[+<<AA>>]:
+         |  def get: <<A@@A>> = ???
+         |  case MySome(value: <<AA>>)
+         |  case MyNone
+         |""".stripMargin,
+    )
+
+  @Test def `type-params-in-enum5` =
+    check(
+      """|enum MyOption[AA]:
+         |  def get: AA = ???
+         |  case MySome[<<AA>>](value: <<A@@A>>) extends MyOption[Int]
+         |""".stripMargin,
+    )