diff --git a/presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala b/presentation-compiler/src/main/dotty/tools/pc/PcCollector.scala index 1700c4e16197..274b0704674d 100644 --- 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.<> */ - 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 index a6254ef3ac09..256b0cb1075a 100644 --- 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 <> |""".stripMargin ) + + @Test def `map-method` = + check( + """|case class Bar(x: List[Int]) { + | def <>(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 index 47a07c873a61..4bf04ed1d0af 100644 --- 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).<>(_ + 1) + | val y = for { + | a <- List(1) + | } yield a + 1 + |} + |""".stripMargin, + ) + + @Test def `for-comp-map1` = + check( + """|object Main { + | val x = List(1).<>(_ + 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).<>(_ => ()) + | val y = for { + | a <- List(1) + | } {} + |} + |""".stripMargin, + ) + + @Test def `for-comp-withFilter` = + check( + """|object Main { + | val x = List(1).<>(_ => true) + | val y = for { + | a <- List(1) + | if true + | } {} + |} + |""".stripMargin, + ) + + @Test def `for-comp-withFilter1` = + check( + """|object Main { + | val x = List(1).withFilter(_ => true).<>(_ + 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).<>(_ => 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).<>(_ => List(1)) + | val y = for { + | a <- List(1) + | if true + | b <- List(2) + | } yield a + 1 + |} + |""".stripMargin, + ) @Test def `enum1` = check(