From 63a908cff8bd0aeae9378f8590e53d44c88c35b6 Mon Sep 17 00:00:00 2001
From: Kirill Rakhman <kirill.rakhman@jetbrains.com>
Date: Mon, 16 Jan 2023 09:56:55 +0100
Subject: [PATCH] FIR, Assign plugin: Update for KT-54648

---
 .../k2/FirAssignmentPluginAssignAltererExtension.kt   | 11 +++++------
 .../testData/diagnostics/otherOperators.fir.kt        |  4 ++--
 .../testData/diagnostics/plusAssignPrecedence.fir.kt  |  2 +-
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/plugins/assign-plugin/assign-plugin.k2/src/org/jetbrains/kotlin/assignment/plugin/k2/FirAssignmentPluginAssignAltererExtension.kt b/plugins/assign-plugin/assign-plugin.k2/src/org/jetbrains/kotlin/assignment/plugin/k2/FirAssignmentPluginAssignAltererExtension.kt
index 1f6d17c71e9..89d9728f53f 100644
--- a/plugins/assign-plugin/assign-plugin.k2/src/org/jetbrains/kotlin/assignment/plugin/k2/FirAssignmentPluginAssignAltererExtension.kt
+++ b/plugins/assign-plugin/assign-plugin.k2/src/org/jetbrains/kotlin/assignment/plugin/k2/FirAssignmentPluginAssignAltererExtension.kt
@@ -19,8 +19,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirBackingFieldSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFieldSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
-import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
-import org.jetbrains.kotlin.fir.types.upperBoundIfFlexible
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class FirAssignmentPluginAssignAltererExtension(
@@ -34,7 +33,7 @@ class FirAssignmentPluginAssignAltererExtension(
     }
 
     private fun FirVariableAssignment.supportsTransformVariableAssignment(): Boolean {
-        return when (val lSymbol = lValue.toResolvedVariableSymbol()) {
+        return when (val lSymbol = calleeReference?.toResolvedVariableSymbol()) {
             is FirPropertySymbol -> lSymbol.isVal && !lSymbol.isLocal && lSymbol.hasSpecialAnnotation()
             is FirBackingFieldSymbol -> lSymbol.isVal && lSymbol.hasSpecialAnnotation()
             is FirFieldSymbol -> lSymbol.isVal && lSymbol.hasSpecialAnnotation()
@@ -46,7 +45,7 @@ class FirAssignmentPluginAssignAltererExtension(
         session.annotationMatchingService.isAnnotated(resolvedReturnType.upperBoundIfFlexible().toRegularClassSymbol(session))
 
     private fun buildFunctionCall(variableAssignment: FirVariableAssignment): FirFunctionCall {
-        val leftArgument = variableAssignment.lValue
+        val leftArgument = variableAssignment.calleeReference!!
         val leftSymbol = leftArgument.toResolvedVariableSymbol()!!
         val leftResolvedType = leftSymbol.resolvedReturnTypeRef
         val rightArgument = variableAssignment.rValue
@@ -55,8 +54,8 @@ class FirAssignmentPluginAssignAltererExtension(
             explicitReceiver = buildPropertyAccessExpression {
                 source = leftArgument.source
                 typeRef = leftResolvedType
-                calleeReference = variableAssignment.calleeReference
-                typeArguments += variableAssignment.typeArguments
+                calleeReference = leftArgument
+                (variableAssignment.lValue as? FirQualifiedAccessExpression)?.typeArguments?.let(typeArguments::addAll)
                 annotations += variableAssignment.annotations
                 explicitReceiver = variableAssignment.explicitReceiver
                 dispatchReceiver = variableAssignment.dispatchReceiver
diff --git a/plugins/assign-plugin/testData/diagnostics/otherOperators.fir.kt b/plugins/assign-plugin/testData/diagnostics/otherOperators.fir.kt
index 90429adcce5..16b2667b6c6 100644
--- a/plugins/assign-plugin/testData/diagnostics/otherOperators.fir.kt
+++ b/plugins/assign-plugin/testData/diagnostics/otherOperators.fir.kt
@@ -50,6 +50,6 @@ fun `should not effect error reporting for other operators`() {
     task<!NO_SET_METHOD!>[0]<!> = StringProperty("Fail")
 
     // a.get(i) = b should not be translated to a.get(i).assign(b)
-    task.<!FUNCTION_CALL_EXPECTED!>get<!>(0) = StringProperty("Fail")
-    nullTask?.<!FUNCTION_CALL_EXPECTED!>get<!>(0) = StringProperty("Fail")
+    task.<!VARIABLE_EXPECTED!>get(0)<!> = StringProperty("Fail")
+    nullTask?.<!VARIABLE_EXPECTED!>get(0)<!> = StringProperty("Fail")
 }
diff --git a/plugins/assign-plugin/testData/diagnostics/plusAssignPrecedence.fir.kt b/plugins/assign-plugin/testData/diagnostics/plusAssignPrecedence.fir.kt
index bf436bfb97b..fc856ffa129 100644
--- a/plugins/assign-plugin/testData/diagnostics/plusAssignPrecedence.fir.kt
+++ b/plugins/assign-plugin/testData/diagnostics/plusAssignPrecedence.fir.kt
@@ -96,7 +96,7 @@ fun box(): String {
     task.valInput <!UNRESOLVED_REFERENCE!>+=<!> "K"
     task.varInput <!UNRESOLVED_REFERENCE!>+=<!> "K"
 
-    task.<!VAL_REASSIGNMENT!>valInputWithPlus<!> += "K"
+    task.valInputWithPlus += "K"
 
     task.varInputWithPlusAndPlusAssign <!ASSIGN_OPERATOR_AMBIGUITY!>+=<!> "K"