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

^KTIJ-21374 Fixed
2022-03-21 18:46:50 +04:00
parent 51bb548aef
commit 22830cf159
8 changed files with 101 additions and 5 deletions
@@ -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
-    ): Collection<KtSymbol> = fir.calleeReference.toTargetSymbol(session, symbolBuilder)
+        symbolBuilder: KtSymbolByFirBuilder,
+    ): 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,
@@ -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 {
@@ -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 {
@@ -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
+}
@@ -0,0 +1,11 @@
+package test
+
+interface Foo
+
+interface WithOperator {
+    operator fun plusAssign(f: Foo)
+}
+
+fun test(withOperator: WithOperator, foo: Foo) {
+    withOperator <caret>+= foo
+}
@@ -0,0 +1,2 @@
+Resolved to:
+0: (in test.WithOperator) operator fun plusAssign(f: test.Foo)
@@ -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
+}
@@ -0,0 +1,2 @@
+Resolved to:
+0: (in test.WithOperator) operator fun plus(f: test.Foo): test.WithOperator