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 730a56d046b..8dc417a54ac 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 @@ -223,11 +223,14 @@ abstract class PcCollector[T]( * For comprehensions have two owners, one for the enumerators and one for * yield. This is a heuristic to find that out. */ - def isForComprehensionOwner(named: NameTree) = - soughtNames(named.name) && - named.symbol.owner.isAnonymousFunction && owners.exists(owner => - pos.isDefined && owner.pos.isDefined && owner.pos.point == named.symbol.owner.pos.point - ) + def isForComprehensionOwner(named: NameTree) = { + sought.exists(symbol => + symbol.name == named.name && symbol.pos.start == named.symbol.pos.start + ) && + named.symbol.owner.isAnonymousFunction && owners.exists(owner => + pos.isDefined && owner.pos.isDefined && owner.pos.point == named.symbol.owner.pos.point + ) + } def soughtOrOverride(sym: Symbol) = sought(sym) || sym.allOverriddenSymbols.exists(sought(_)) diff --git a/tests/cross/src/test/scala/tests/highlight/DocumentHighlightSuite.scala b/tests/cross/src/test/scala/tests/highlight/DocumentHighlightSuite.scala index 4e9b11aefd5..975c746cb4e 100644 --- a/tests/cross/src/test/scala/tests/highlight/DocumentHighlightSuite.scala +++ b/tests/cross/src/test/scala/tests/highlight/DocumentHighlightSuite.scala @@ -836,6 +836,21 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite { |}""".stripMargin ) + check( + "for-comp-bind2", + """ + |object Main { + | val abc = for { + | <> <- List(1) + | baz = <> + 1 + | a <- List(<>, 123) + | } yield { + | val x = <> + baz + | x + | } + |}""".stripMargin + ) + check( "for-comp-map", """|object Main { @@ -993,4 +1008,16 @@ class DocumentHighlightSuite extends BaseDocumentHighlightSuite { |}""".stripMargin ) + check( + "map-bind6", + """ + |object Main { + | List("test").map { + | case <>: String if <>.startsWith("a") => <> + "a" + | case stringName: String if stringName.startsWith("b") => stringName + "b" + | case stringName: String => stringName + "c" + | } + |}""".stripMargin + ) + }