From 6f64aedaf039bc790991f3424c109a03170734af Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Fri, 11 Mar 2022 11:38:04 +0300
Subject: [PATCH] [FIR] Fix FindReferencePositioningStrategy for light tree
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

There was a problem which causes to lookup for assign token inside
  functions from arguments of call
---
 .../diagnostics/LightTreePositioningStrategies.kt | 15 ++++++++++++---
 .../inference/plusAssignWithLambda2.fir.kt        |  2 +-
 .../src/org/jetbrains/kotlin/utils/addToStdlib.kt |  1 +
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/compiler/frontend.common-psi/src/org/jetbrains/kotlin/diagnostics/LightTreePositioningStrategies.kt b/compiler/frontend.common-psi/src/org/jetbrains/kotlin/diagnostics/LightTreePositioningStrategies.kt
index 7550e608ab6..4e0cd135636 100644
--- a/compiler/frontend.common-psi/src/org/jetbrains/kotlin/diagnostics/LightTreePositioningStrategies.kt
+++ b/compiler/frontend.common-psi/src/org/jetbrains/kotlin/diagnostics/LightTreePositioningStrategies.kt
@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.psi.KtParameter.VAL_VAR_TOKEN_SET
 import org.jetbrains.kotlin.psi.stubs.elements.KtConstantExpressionElementType
 import org.jetbrains.kotlin.psi.stubs.elements.KtStringTemplateExpressionElementType
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
+import org.jetbrains.kotlin.utils.addToStdlib.runUnless
 
 object LightTreePositioningStrategies {
     val DEFAULT = object : LightTreePositioningStrategy() {
@@ -654,7 +655,7 @@ object LightTreePositioningStrategies {
             tree: FlyweightCapableTreeStructure<LighterASTNode>
         ): List<TextRange> {
             when {
-                node.tokenType == KtNodeTypes.BINARY_EXPRESSION && tree.findDescendantByType(node, KtTokens.EQ) != null -> {
+                node.tokenType == KtNodeTypes.BINARY_EXPRESSION && tree.findDescendantByType(node, KtTokens.EQ, followFunctions = false) != null -> {
                     // Look for reference in LHS of variable assignment.
                     tree.findExpressionDeep(node)?.let {
                         return markElement(it, startOffset, endOffset, tree, node)
@@ -1396,11 +1397,19 @@ fun FlyweightCapableTreeStructure<LighterASTNode>.findLastChildByType(node: Ligh
     return childrenRef.get()?.lastOrNull { it?.tokenType == type }
 }
 
-fun FlyweightCapableTreeStructure<LighterASTNode>.findDescendantByType(node: LighterASTNode, type: IElementType): LighterASTNode? {
+fun FlyweightCapableTreeStructure<LighterASTNode>.findDescendantByType(
+    node: LighterASTNode,
+    type: IElementType,
+    followFunctions: Boolean = true
+): LighterASTNode? {
     val childrenRef = Ref<Array<LighterASTNode?>>()
     getChildren(node, childrenRef)
     return childrenRef.get()?.firstOrNull { it?.tokenType == type } ?: childrenRef.get()
-        ?.firstNotNullOfOrNull { child -> child?.let { findDescendantByType(it, type) } }
+        ?.firstNotNullOfOrNull { child ->
+            runUnless(!followFunctions && child?.tokenType == KtNodeTypes.FUN) {
+                child?.let { findDescendantByType(it, type, followFunctions) }
+            }
+        }
 }
 
 fun FlyweightCapableTreeStructure<LighterASTNode>.findDescendantByTypes(node: LighterASTNode, types: TokenSet): LighterASTNode? {
diff --git a/compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignWithLambda2.fir.kt b/compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignWithLambda2.fir.kt
index 9d98f2352a6..169867e7ea0 100644
--- a/compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignWithLambda2.fir.kt
+++ b/compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignWithLambda2.fir.kt
@@ -14,7 +14,7 @@ fun <T> id(x: T) = x
 fun main() {
     var newValue = A()
     newValue <!ASSIGN_OPERATOR_AMBIGUITY!>+=<!> id { total -> A() }
-    <!ASSIGN_OPERATOR_AMBIGUITY!>newValue<!> += id(fun(total) = A())
+    newValue <!ASSIGN_OPERATOR_AMBIGUITY!>+=<!> id(fun(total) = A())
     newValue <!ASSIGN_OPERATOR_AMBIGUITY!>+=<!> id(fun(total): A { return A() })
     newValue <!ASSIGN_OPERATOR_AMBIGUITY!>+=<!> id(::foo)
 }
diff --git a/core/util.runtime/src/org/jetbrains/kotlin/utils/addToStdlib.kt b/core/util.runtime/src/org/jetbrains/kotlin/utils/addToStdlib.kt
index 7f11f596933..2b4e1713da3 100644
--- a/core/util.runtime/src/org/jetbrains/kotlin/utils/addToStdlib.kt
+++ b/core/util.runtime/src/org/jetbrains/kotlin/utils/addToStdlib.kt
@@ -178,6 +178,7 @@ inline fun <T, R> Iterable<T>.same(extractor: (T) -> R): Boolean {
 }
 
 inline fun <R> runIf(condition: Boolean, block: () -> R): R? = if (condition) block() else null
+inline fun <R> runUnless(condition: Boolean, block: () -> R): R? = if (condition) null else block()
 
 inline fun <T, R> Collection<T>.foldMap(transform: (T) -> R, operation: (R, R) -> R): R {
     val iterator = iterator()