KT-60954 [Analysis API] Unwrap variable assignments in KtFirReferenceShortener

To get to the proper qualified expression, we need to unwrap the outer
`FirVariableAssignment` if we deal with the property access inside of
assignment expressions

^KT-60954 Fixed

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt

@@ -973,7 +973,7 @@ private class ElementsToShortenCollector(
         val scopes = shorteningContext.findScopesAtPosition(qualifiedProperty, getNamesToImport(), towerContextProvider) ?: return
         val availableCallables = shorteningContext.findPropertiesInScopes(scopes, callableSymbol.name)
 
-        val firPropertyAccess = qualifiedProperty.getOrBuildFir(firResolveSession) as? FirQualifiedAccessExpression ?: return
+        val firPropertyAccess = qualifiedProperty.getCorrespondingPropertyAccessExpression() ?: return
 
         // if explicit receiver is a property access or a function call, we cannot shorten it
         if (firPropertyAccess.explicitReceiver !is FirResolvedQualifier) return
@@ -992,6 +992,15 @@ private class ElementsToShortenCollector(
         )
     }
 
+    private fun KtDotQualifiedExpression.getCorrespondingPropertyAccessExpression(): FirPropertyAccessExpression? {
+        val accessExpression = when (val fir = getOrBuildFir(firResolveSession)) {
+            is FirVariableAssignment -> fir.unwrapLValue()
+            else -> fir
+        }
+
+        return accessExpression as? FirPropertyAccessExpression
+    }
+
     private fun processFunctionCall(functionCall: FirFunctionCall) {
         if (!canBePossibleToDropReceiver(functionCall)) return
 

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated.java

@@ -430,6 +430,24 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variable2.kt");
     }
 
+    @Test
+    @TestMetadata("variableAssignment.kt")
+    public void testVariableAssignment() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment.kt");
+    }
+
+    @Test
+    @TestMetadata("variableAssignment_plusAssignOperator.kt")
+    public void testVariableAssignment_plusAssignOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusAssignOperator.kt");
+    }
+
+    @Test
+    @TestMetadata("variableAssignment_plusOperator.kt")
+    public void testVariableAssignment_plusOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusOperator.kt");
+    }
+
     @Nested
     @TestMetadata("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses")
     @TestDataPath("$PROJECT_ROOT")

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/references/FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenerated.java

@@ -430,6 +430,24 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenera
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variable2.kt");
     }
 
+    @Test
+    @TestMetadata("variableAssignment.kt")
+    public void testVariableAssignment() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment.kt");
+    }
+
+    @Test
+    @TestMetadata("variableAssignment_plusAssignOperator.kt")
+    public void testVariableAssignment_plusAssignOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusAssignOperator.kt");
+    }
+
+    @Test
+    @TestMetadata("variableAssignment_plusOperator.kt")
+    public void testVariableAssignment_plusOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusOperator.kt");
+    }
+
     @Nested
     @TestMetadata("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses")
     @TestDataPath("$PROJECT_ROOT")

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment.kt

@@ -0,0 +1,11 @@
+// FILE: main.kt
+package test
+
+fun usage() {
+    <expr>dependency.foo = 20</expr>
+}
+
+// FILE: dependency.kt
+package dependency
+
+var foo: Int = 10

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment.txt

@@ -0,0 +1,7 @@
+Before shortening: dependency.foo = 20
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[qualifier] dependency.foo
+with SHORTEN_AND_STAR_IMPORT:
+[qualifier] dependency.foo

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusAssignOperator.kt

@@ -0,0 +1,15 @@
+// FILE: main.kt
+package test
+
+fun usage() {
+    <expr>dependency.foo += 20</expr>
+}
+
+// FILE: dependency.kt
+package dependency
+
+val foo: Foo = Foo()
+
+class Foo {
+    operator fun plusAssign(i: Int) {}
+}

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusAssignOperator.txt

@@ -0,0 +1,7 @@
+Before shortening: dependency.foo += 20
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[qualifier] dependency.foo
+with SHORTEN_AND_STAR_IMPORT:
+[qualifier] dependency.foo

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusOperator.kt

@@ -0,0 +1,15 @@
+// FILE: main.kt
+package test
+
+fun usage() {
+    <expr>dependency.foo += 20</expr>
+}
+
+// FILE: dependency.kt
+package dependency
+
+var foo: Foo = Foo()
+
+class Foo {
+    operator fun plus(i: Int): Foo = this
+}

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusOperator.txt

@@ -0,0 +1,7 @@
+Before shortening: dependency.foo += 20
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[qualifier] dependency.foo
+with SHORTEN_AND_STAR_IMPORT:
+[qualifier] dependency.foo