fix: Don't collect map, flatMap, withFilter in for-comprehension (#18430

) backport from metals: scalameta/metals#5552 CC: @tgodzik
scala · Aug 22, 2023 · b823a69 · b823a69
2 parents 6e370a9 + 5554a4b
commit b823a69
Show file tree

Hide file tree

Showing 3 changed files with 111 additions and 1 deletion.
diff --git a/presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala b/presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala
@@ -409,7 +409,9 @@ abstract class PcCollector[T](
          * All select statements such as:
          * val a = hello.<<b>>
          */
-        case sel: Select if sel.span.isCorrect && filter(sel) =>
+        case sel: Select 
+          if sel.span.isCorrect && filter(sel) &&
+            !isForComprehensionMethod(sel) =>
           occurrences + collect(
             sel,
             pos.withSpan(selectNameSpan(sel))
@@ -560,6 +562,17 @@ abstract class PcCollector[T](
         Span(span.start, span.start + realName.length, point)
       else Span(point, span.end, point)
     else span
+
+  private val forCompMethods =
+    Set(nme.map, nme.flatMap, nme.withFilter, nme.foreach)
+
+  // We don't want to collect synthethic `map`, `withFilter`, `foreach` and `flatMap` in for-comprenhensions
+  private def isForComprehensionMethod(sel: Select): Boolean =
+    val syntheticName = sel.name match
+      case name: TermName => forCompMethods(name)
+      case _ => false
+    val wrongSpan = sel.qualifier.span.contains(sel.nameSpan)
+    syntheticName && wrongSpan
 end PcCollector
 
 object PcCollector:

diff --git a/presentation-compiler/test/dotty/tools/pc/tests/edit/PcRenameSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/edit/PcRenameSuite.scala
@@ -471,3 +471,16 @@ class PcRenameSuite extends BasePcRenameSuite:
          |} yield <<a@@bc>>
          |""".stripMargin
     )
+
+  @Test def `map-method` =
+    check(
+      """|case class Bar(x: List[Int]) {
+         |  def <<ma@@p>>(f: Int => Int): Bar = Bar(x.map(f))
+         |}
+         |
+         |val f = 
+         |  for {
+         |    b <- Bar(List(1,2,3))
+         |  } yield b
+         |""".stripMargin,
+    )
diff --git a/presentation-compiler/test/dotty/tools/pc/tests/highlight/DocumentHighlightSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/highlight/DocumentHighlightSuite.scala
@@ -758,6 +758,90 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite:
         |  }
         |}""".stripMargin
     )
+
+  @Test def `for-comp-map` = 
+    check(
+      """|object Main {
+         |  val x = List(1).<<m@@ap>>(_ + 1)
+         |  val y = for {
+         |    a <- List(1)
+         |  } yield a + 1
+         |}
+         |""".stripMargin,
+    )
+
+  @Test def `for-comp-map1` = 
+    check(
+      """|object Main {
+         |  val x = List(1).<<m@@ap>>(_ + 1)
+         |  val y = for {
+         |    a <- List(1)
+         |    if true
+         |  } yield a + 1
+         |}
+         |""".stripMargin,
+    )
+
+  @Test def `for-comp-foreach` = 
+    check(
+      """|object Main {
+         |  val x = List(1).<<for@@each>>(_ => ())
+         |  val y = for {
+         |    a <- List(1)
+         |  } {}
+         |}
+         |""".stripMargin,
+    )
+
+  @Test def `for-comp-withFilter` = 
+    check(
+      """|object Main {
+         |  val x = List(1).<<with@@Filter>>(_ => true)
+         |  val y = for {
+         |    a <- List(1)
+         |    if true
+         |  } {}
+         |}
+         |""".stripMargin,
+    )
+
+  @Test def `for-comp-withFilter1` = 
+    check(
+      """|object Main {
+         |  val x = List(1).withFilter(_ => true).<<m@@ap>>(_ + 1)
+         |  val y = for {
+         |    a <- List(1)
+         |    if true
+         |  } yield a + 1
+         |}
+         |""".stripMargin,
+    )
+
+  @Test def `for-comp-flatMap1` = 
+    check(
+      """|object Main {
+         |  val x = List(1).<<flat@@Map>>(_ => List(1))
+         |  val y = for {
+         |    a <- List(1)
+         |    b <- List(2)
+         |    if true
+         |  } yield a + 1
+         |}
+         |""".stripMargin,
+    )
+
+  @Test def `for-comp-flatMap2` = 
+    check(
+      """|object Main {
+         |  val x = List(1).withFilter(_ => true).<<flat@@Map>>(_ => List(1))
+         |  val y = for {
+         |    a <- List(1)
+         |    if true
+         |    b <- List(2)
+         |  } yield a + 1
+         |}
+         |""".stripMargin,
+    )
 
   @Test def `enum1` =
     check(