bugfix: highlight for enum type params

scalameta · Aug 25, 2023 · ba67d09 · ba67d09
1 parent 4b401c9
commit ba67d09
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 3 deletions.
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
@@ -112,6 +112,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.owner.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)
@@ -130,7 +157,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.owner.owner))
+          ++ additionalForEnumTypeParam(sym.owner.owner)
+      else if sym.isTypeParam then
+        primaryConstructorTypeParam(sym.owner).toSet
+          ++ additionalForEnumTypeParam(sym.owner) + sym
       else Set(sym)
     all.filter(s => s != NoSymbol && !s.isError)
   end symbolAlternatives

diff --git a/tests/cross/src/test/scala/tests/highlight/Scala3DocumentHighlightSuite.scala b/tests/cross/src/test/scala/tests/highlight/Scala3DocumentHighlightSuite.scala
@@ -248,7 +248,7 @@ class Scala3DocumentHighlightSuite extends BaseDocumentHighlightSuite {
   )
 
   check(
-    "enum-cases".tag(IgnoreScala2),
+    "enum-cases",
     """|enum MyOption:
        |  case <<My@@Some>>(value: Int)
        |  case MyNone
@@ -258,7 +258,7 @@ class Scala3DocumentHighlightSuite extends BaseDocumentHighlightSuite {
   )
 
   check(
-    "enum-cases2".tag(IgnoreScala2),
+    "enum-cases2",
     """|enum MyOption:
        |  case <<My@@Some>>[U](value: U)
        |  case MyNone
@@ -267,4 +267,45 @@ class Scala3DocumentHighlightSuite extends BaseDocumentHighlightSuite {
        |""".stripMargin,
   )
 
+  check(
+    "type-params-in-enum",
+    """|enum MyOption[+<<A@@A>>]:
+       |  case MySome(value: <<AA>>)
+       |  case MyNone
+       |""".stripMargin,
+  )
+
+  check(
+    "type-params-in-enum2",
+    """|enum MyOption[+<<AA>>]:
+       |  case MySome(value: <<A@@A>>)
+       |  case MyNone
+       |""".stripMargin,
+  )
+
+  check(
+    "type-params-in-enum3",
+    """|enum MyOption[<<AA>>](v: <<AA>>):
+       |  def get: <<A@@A>> = ???
+       |  case MySome[AA](value: AA) extends MyOption[Int](1)
+       |""".stripMargin,
+  )
+
+  check(
+    "type-params-in-enum4",
+    """|enum MyOption[+<<AA>>]:
+       |  def get: <<A@@A>> = ??
+       |  case MySome(value: <<AA>>)
+       |  case MyNone
+       |""".stripMargin,
+  )
+
+  check(
+    "type-params-in-enum5",
+    """|enum MyOption[AA]:
+       |  def get: AA = ???
+       |  case MySome[<<AA>>](value: <<A@@A>>) extends MyOption[Int]
+       |""".stripMargin,
+  )
+
 }