From e71c171473ace59a8075313b8633e39b10b8e215 Mon Sep 17 00:00:00 2001
From: Jakub Ciesluk <323892@uwr.edu.pl>
Date: Thu, 16 Nov 2023 15:24:09 +0100
Subject: [PATCH] bugfix: Fix PcCollector for Scala 2 in anon-fun parameters

Anonymous function parameters with type annotation have incorrect name position on definition,
so we need to handle them separately.
---
 .../scala/meta/internal/pc/PcCollector.scala  | 35 +++++++++++
 .../highlight/DocumentHighlightSuite.scala    | 63 +++++++++++++++++++
 .../tests/tokens/SemanticTokensSuite.scala    | 10 +++
 3 files changed, 108 insertions(+)

diff --git a/mtags/src/main/scala-2/scala/meta/internal/pc/PcCollector.scala b/mtags/src/main/scala-2/scala/meta/internal/pc/PcCollector.scala
index 60f2f3fbd9b..730a56d046b 100644
--- a/mtags/src/main/scala-2/scala/meta/internal/pc/PcCollector.scala
+++ b/mtags/src/main/scala-2/scala/meta/internal/pc/PcCollector.scala
@@ -96,6 +96,14 @@ abstract class PcCollector[T](
       else {
         Some(symbolAlternatives(id.symbol), id.pos)
       }
+    /* Anonynous function parameters such as:
+     * List(1).map{ <<abc>>: Int => abc}
+     * In this case, parameter has incorrect namePosition, so we need to handle it separately
+     */
+    case (vd: ValDef) if isAnonFunctionParam(vd) =>
+      val namePos = vd.pos.withEnd(vd.pos.start + vd.name.length)
+      if (namePos.includes(pos)) Some(symbolAlternatives(vd.symbol), namePos)
+      else None
 
     /* all definitions:
      * def fo@@o = ???
@@ -322,6 +330,30 @@ abstract class PcCollector[T](
             newAcc,
             sel.qualifier
           )
+
+        /**
+         * Anonynous function parameters such as:
+         * List(1).map{ <<abc>>: Int => abc}
+         * In this case, parameter has incorrect namePosition, so we need to handle it separately
+         */
+        case Function(params, body) =>
+          val newAcc = params
+            .filter(vd => vd.pos.isRange && filter(vd))
+            .foldLeft(
+              acc
+            ) { case (acc, vd) =>
+              val pos = vd.pos.withEnd(vd.pos.start + vd.name.length)
+              annotationChildren(vd).foldLeft({
+                acc + collect(
+                  vd,
+                  pos
+                )
+              })(traverse(_, _))
+            }
+          traverse(
+            params.map(_.tpt).foldLeft(newAcc)(traverse(_, _)),
+            body
+          )
         /* all definitions:
          * def <<foo>> = ???
          * class <<Foo>> = ???
@@ -495,4 +527,7 @@ abstract class PcCollector[T](
     }
   }
 
+  private def isAnonFunctionParam(vd: ValDef): Boolean =
+    vd.symbol != null && vd.symbol.owner.isAnonymousFunction && vd.rhs.isEmpty
+
 }
diff --git a/tests/cross/src/test/scala/tests/highlight/DocumentHighlightSuite.scala b/tests/cross/src/test/scala/tests/highlight/DocumentHighlightSuite.scala
index 15f1a14f56d..4e9b11aefd5 100644
--- a/tests/cross/src/test/scala/tests/highlight/DocumentHighlightSuite.scala
+++ b/tests/cross/src/test/scala/tests/highlight/DocumentHighlightSuite.scala
@@ -930,4 +930,67 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite {
        |}
        |""".stripMargin
   )
+
+  check(
+    "map-bind".tag(IgnoreScala3),
+    """
+      |object Main {
+      |  List(1).map {
+      |    <<abc>>: Int => <<ab@@c>> + 1
+      |  }
+      |}""".stripMargin
+  )
+
+  check(
+    "map-bind1".tag(IgnoreScala3),
+    """
+      |object Main {
+      |  List(1).foldLeft(0){
+      |    (<<abc>>: Int, bde: Int) => <<ab@@c>> + bde
+      |  }
+      |}""".stripMargin
+  )
+
+  check(
+    "map-bind2".tag(IgnoreScala3),
+    """
+      |object Main {
+      |  List(1).map {
+      |    someLongName:@@ Int => someLongName + 1
+      |  }
+      |}""".stripMargin
+  )
+
+  check(
+    "map-bind3".tag(IgnoreScala3),
+    """
+      |object Main {
+      |  List(1).map {
+      |    someVeryLongName: Int =@@> someVeryLongName + 1
+      |  }
+      |}""".stripMargin
+  )
+
+  check(
+    "map-bind4".tag(IgnoreScala3),
+    """
+      |object Main {
+      |  List(1).map {
+      |    someLongName: <<I@@nt>> => someLongName + 1
+      |  }
+      |}""".stripMargin
+  )
+
+  check(
+    "map-bind5".tag(IgnoreScala3),
+    """
+      |object Main {
+      |  List(1).map {
+      |    someLongName: Int => 
+      |      val <<ab@@c>> = 2
+      |      someLongName + <<abc>>
+      |  }
+      |}""".stripMargin
+  )
+
 }
diff --git a/tests/cross/src/test/scala/tests/tokens/SemanticTokensSuite.scala b/tests/cross/src/test/scala/tests/tokens/SemanticTokensSuite.scala
index 0461cbfaffd..45401f74576 100644
--- a/tests/cross/src/test/scala/tests/tokens/SemanticTokensSuite.scala
+++ b/tests/cross/src/test/scala/tests/tokens/SemanticTokensSuite.scala
@@ -332,4 +332,14 @@ class SemanticTokensSuite extends BaseSemanticTokensSuite {
        |""".stripMargin
   )
 
+  check(
+    "map-bind".tag(IgnoreScala3),
+    """
+      |object <<Main>>/*class*/ {
+      |  <<List>>/*class*/(1).<<foldLeft>>/*method*/(0){
+      |    (<<abc>>/*parameter,declaration,readonly*/: <<Int>>/*class,abstract*/, <<bde>>/*parameter,declaration,readonly*/: <<Int>>/*class,abstract*/) => <<abc>>/*parameter,readonly*/ <<+>>/*method,abstract*/ <<bde>>/*parameter,readonly*/
+      |  }
+      |}""".stripMargin
+  )
+
 }