[FIR] Fix augmented array set call on smartcasted receiver

#KT-60917 Fixed

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -1527,15 +1527,17 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         lhsGetCall: FirFunctionCall,
         transformedRhs: FirExpression
     ): AugmentedArraySetAsGetSetCallDesugaringInfo {
+        val initializer = lhsGetCall.explicitReceiver ?: buildErrorExpression {
+            source = augmentedArraySetCall.source
+                ?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
+            diagnostic = ConeSyntaxDiagnostic("No receiver for array access")
+        }
         val arrayVariable = generateTemporaryVariable(
             session.moduleData,
             source = lhsGetCall.explicitReceiver?.source?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment),
             name = SpecialNames.ARRAY,
-            initializer = lhsGetCall.explicitReceiver ?: buildErrorExpression {
-                source = augmentedArraySetCall.source
-                    ?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
-                diagnostic = ConeSyntaxDiagnostic("No receiver for array access")
-            }
+            initializer = initializer,
+            typeRef = initializer.typeRef.copyWithNewSourceKind(KtFakeSourceElementKind.DesugaredCompoundAssignment),
         )
 
         val indexVariables = lhsGetCall.arguments.flatMap {

compiler/testData/diagnostics/tests/AssignToArrayElement.kt

@@ -1,6 +1,9 @@
 // FIR_IDENTICAL
+// !DIAGNOSTICS: -DEBUG_INFO_SMARTCAST
+
 fun getArray(): Array<Int> = throw Exception()
 fun getList(): MutableList<Int> = throw Exception()
+fun getNullableList(): MutableList<Int>? = throw Exception()
 fun fn() {
     getArray()[1] = 2
     getList()[1] = 2
@@ -8,4 +11,9 @@ fun fn() {
     getList()[1]++
     getArray()[1] += 2
     getList()[1] += 2
+
+    val nullable = getNullableList()
+    if (nullable != null) {
+        nullable[1] += 1
+    }
 }

compiler/testData/diagnostics/tests/AssignToArrayElement.txt

@@ -3,3 +3,5 @@ package
 public fun fn(): kotlin.Unit
 public fun getArray(): kotlin.Array<kotlin.Int>
 public fun getList(): kotlin.collections.MutableList<kotlin.Int>
+public fun getNullableList(): kotlin.collections.MutableList<kotlin.Int>?
+

compiler/testData/diagnostics/testsWithStdLib/smartcasts/unsoundSmartcast.fir.kt

@@ -4,7 +4,7 @@
 fun test1() {
     var x: MutableList<Int>? = mutableListOf(1)
     x!!
-    x[if (true) { x = null; 0 } else 0] <!UNRESOLVED_REFERENCE!>+=<!> <!UNSAFE_CALL!>x[0]<!>
+    x[if (true) { x = null; 0 } else 0] += <!UNSAFE_CALL!>x[0]<!>
     <!UNSAFE_CALL!>x[0]<!>.inv()
 }
 

compiler/tests-spec/testData/diagnostics/notLinked/dfa/neg/13.fir.kt

@@ -15,7 +15,7 @@
 // TESTCASE NUMBER: 1
 fun case_1(x: Class?) {
     x!!
-    <!DEBUG_INFO_EXPRESSION_TYPE("Class? & Class")!>x<!>[if (true) {<!VAL_REASSIGNMENT!>x<!>=null;0} else 0] <!UNRESOLVED_REFERENCE!>+=<!> <!DEBUG_INFO_EXPRESSION_TYPE("Class? & Class")!>x<!>[0]
+    <!DEBUG_INFO_EXPRESSION_TYPE("Class? & Class"), DEBUG_INFO_EXPRESSION_TYPE("Class"), DEBUG_INFO_EXPRESSION_TYPE("Class")!>x<!>[if (true) {<!VAL_REASSIGNMENT!>x<!>=null;0} else 0] += <!DEBUG_INFO_EXPRESSION_TYPE("Class? & Class")!>x<!>[0]
     <!DEBUG_INFO_EXPRESSION_TYPE("Class? & Class")!>x<!>
     <!DEBUG_INFO_EXPRESSION_TYPE("Class? & Class")!>x<!>[0].inv()
 }

compiler/tests-spec/testData/diagnostics/notLinked/dfa/neg/26.fir.kt

@@ -18,7 +18,7 @@
 fun case_1() {
     var x: MutableList<Int>? = mutableListOf(1)
     x!!
-    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.MutableList<kotlin.Int>? & kotlin.collections.MutableList<kotlin.Int>")!>x<!>[if (true) {x=null;0} else 0] <!UNRESOLVED_REFERENCE!>+=<!> <!UNSAFE_CALL!><!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.MutableList<kotlin.Int>?")!>x<!>[0]<!>
+    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.MutableList<kotlin.Int>? & kotlin.collections.MutableList<kotlin.Int>"), DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.MutableList<kotlin.Int>"), DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.MutableList<kotlin.Int>")!>x<!>[if (true) {x=null;0} else 0] += <!UNSAFE_CALL!><!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.MutableList<kotlin.Int>?")!>x<!>[0]<!>
     <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.MutableList<kotlin.Int>?")!>x<!>
     <!UNSAFE_CALL!><!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.MutableList<kotlin.Int>?")!>x<!>[0]<!>.inv()
 }