[FIR IDE] Fix resolving sugared compound operators (like +=, -=)

^KTIJ-21374 Fixed

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirReferenceResolveHelper.kt

@@ -7,7 +7,10 @@ package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.tree.TokenSet
 import org.jetbrains.kotlin.KtFakeSourceElementKind
-import org.jetbrains.kotlin.analysis.api.fir.*
+import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.analysis.api.fir.buildSymbol
+import org.jetbrains.kotlin.analysis.api.fir.getCandidateSymbols
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirPackageSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
@@ -201,7 +204,7 @@ internal object FirReferenceResolveHelper {
             }
             is FirReturnExpression -> getSymbolsByReturnExpression(expression, fir, symbolBuilder)
             is FirErrorNamedReference -> getSymbolsByErrorNamedReference(fir, symbolBuilder)
-            is FirVariableAssignment -> getSymbolsByVariableAssignment(fir, session, symbolBuilder)
+            is FirVariableAssignment -> getSymbolsByVariableAssignment(fir, expression, session, symbolBuilder)
             is FirResolvedNamedReference -> getSymbolByResolvedNameReference(fir, expression, analysisSession, session, symbolBuilder)
             is FirDelegatedConstructorCall ->
                 getSymbolByDelegatedConstructorCall(expression, adjustedResolutionExpression, fir, session, symbolBuilder)
@@ -268,12 +271,29 @@ internal object FirReferenceResolveHelper {
         else -> false
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
     private fun getSymbolsByVariableAssignment(
         fir: FirVariableAssignment,
+        expression: KtSimpleNameExpression,
         session: FirSession,
-        symbolBuilder: KtSymbolByFirBuilder
+        symbolBuilder: KtSymbolByFirBuilder,
-    ): Collection<KtSymbol> = fir.calleeReference.toTargetSymbol(session, symbolBuilder)
+    ): Collection<KtSymbol> {
+        if (expression is KtNameReferenceExpression) {
+            return fir.calleeReference.toTargetSymbol(session, symbolBuilder)
+        }
+
+        val assignmentRValue = fir.rValue
+        if (expression is KtOperationReferenceExpression &&
+            assignmentRValue.source?.kind is KtFakeSourceElementKind.DesugaredCompoundAssignment
+        ) {
+            require(assignmentRValue is FirResolvable) {
+                "Rvalue of desugared compound assignment should be resolvable, but it was ${assignmentRValue::class}"
+            }
+
+            return assignmentRValue.calleeReference.toTargetSymbol(session, symbolBuilder)
+        }
+
+        return emptyList()
+    }
 
     private fun getSymbolsByNameArgumentExpression(
         expression: KtSimpleNameExpression,

analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/FirLibrarySourceReferenceResolveTestGenerated.java

@@ -334,6 +334,24 @@ public class FirLibrarySourceReferenceResolveTestGenerated extends AbstractRefer
         runTest("analysis/analysis-api/testData/referenceResolve/parameterByName.kt");
     }
 
+    @Test
+    @TestMetadata("PlusAssignByHand.kt")
+    public void testPlusAssignByHand() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/PlusAssignByHand.kt");
+    }
+
+    @Test
+    @TestMetadata("PlusAssignOperator.kt")
+    public void testPlusAssignOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/PlusAssignOperator.kt");
+    }
+
+    @Test
+    @TestMetadata("PlusAssignViaPlusOperator.kt")
+    public void testPlusAssignViaPlusOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/PlusAssignViaPlusOperator.kt");
+    }
+
     @Test
     @TestMetadata("ReferenceInClassWhereConstraint.kt")
     public void testReferenceInClassWhereConstraint() throws Exception {

analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/FirSourceReferenceResolveTestGenerated.java

@@ -334,6 +334,24 @@ public class FirSourceReferenceResolveTestGenerated extends AbstractReferenceRes
         runTest("analysis/analysis-api/testData/referenceResolve/parameterByName.kt");
     }
 
+    @Test
+    @TestMetadata("PlusAssignByHand.kt")
+    public void testPlusAssignByHand() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/PlusAssignByHand.kt");
+    }
+
+    @Test
+    @TestMetadata("PlusAssignOperator.kt")
+    public void testPlusAssignOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/PlusAssignOperator.kt");
+    }
+
+    @Test
+    @TestMetadata("PlusAssignViaPlusOperator.kt")
+    public void testPlusAssignViaPlusOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/PlusAssignViaPlusOperator.kt");
+    }
+
     @Test
     @TestMetadata("ReferenceInClassWhereConstraint.kt")
     public void testReferenceInClassWhereConstraint() throws Exception {

analysis/analysis-api/testData/referenceResolve/PlusAssignByHand.kt

@@ -0,0 +1,13 @@
+// UNRESOLVED_REFERENCE
+package test
+
+interface Foo
+
+interface WithOperator {
+    operator fun plus(f: Foo): WithOperator
+}
+
+fun test(withOperator: WithOperator, foo: Foo) {
+    var variable = withOperator
+    variable <caret>= variable + foo
+}

analysis/analysis-api/testData/referenceResolve/PlusAssignOperator.kt

@@ -0,0 +1,11 @@
+package test
+
+interface Foo
+
+interface WithOperator {
+    operator fun plusAssign(f: Foo)
+}
+
+fun test(withOperator: WithOperator, foo: Foo) {
+    withOperator <caret>+= foo
+}

analysis/analysis-api/testData/referenceResolve/PlusAssignOperator.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: (in test.WithOperator) operator fun plusAssign(f: test.Foo)

analysis/analysis-api/testData/referenceResolve/PlusAssignViaPlusOperator.kt

@@ -0,0 +1,12 @@
+package test
+
+interface Foo
+
+interface WithOperator {
+    operator fun plus(f: Foo): WithOperator
+}
+
+fun test(withOperator: WithOperator, foo: Foo) {
+    var variable = withOperator
+    variable <caret>+= foo
+}

analysis/analysis-api/testData/referenceResolve/PlusAssignViaPlusOperator.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: (in test.WithOperator) operator fun plus(f: test.Foo): test.WithOperator