K2: add more tests/fix augmented assignment case around KT-59748

analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayAccessConvention.txt

@@ -20,4 +20,5 @@ FILE: [ResolvedTo(IMPORTS)] compoundAssignWithArrayAccessConvention.kt
             R|<local>/<array>|.R|SubstitutionOverride</MyMap.set: R|kotlin/Unit|>|(R|<local>/<index_0>|, R|<local>/<array>|.R|SubstitutionOverride</MyMap.get: R|kotlin/Int|>|(R|<local>/<index_0>|).R|kotlin/Int.plus|(Int(1)))
         }
 
+        Unit
     }

analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayAccessConvention_lhs.txt

@@ -20,4 +20,5 @@ FILE: [ResolvedTo(IMPORTS)] compoundAssignWithArrayAccessConvention_lhs.kt
             R|<local>/<array>|.R|SubstitutionOverride</MyMap.set: R|kotlin/Unit|>|(R|<local>/<index_0>|, R|<local>/<array>|.R|SubstitutionOverride</MyMap.get: R|kotlin/Int|>|(R|<local>/<index_0>|).R|kotlin/Int.plus|(Int(1)))
         }
 
+        Unit
     }

analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayAccessConvention_propertyAccess.txt

@@ -20,4 +20,5 @@ FILE: [ResolvedTo(IMPORTS)] compoundAssignWithArrayAccessConvention_propertyAcce
             R|<local>/<array>|.R|SubstitutionOverride</MyMap.set: R|kotlin/Unit|>|(R|<local>/<index_0>|, R|<local>/<array>|.R|SubstitutionOverride</MyMap.get: R|kotlin/Int|>|(R|<local>/<index_0>|).R|kotlin/Int.plus|(Int(1)))
         }
 
+        Unit
     }

analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayGetConvention.txt

@@ -17,4 +17,5 @@ FILE: [ResolvedTo(IMPORTS)] compoundAssignWithArrayGetConvention.kt
     }
     public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] m: R|MyMap<kotlin/String, A>|): R|kotlin/Unit| {
         R|<local>/m|.R|SubstitutionOverride</MyMap.get: R|A|>|(String(a)).R|/A.plusAssign|(Int(1))
+        Unit
     }

analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignWithArrayGetConvention_lhs.txt

@@ -17,4 +17,5 @@ FILE: [ResolvedTo(IMPORTS)] compoundAssignWithArrayGetConvention_lhs.kt
     }
     public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] m: R|MyMap<kotlin/String, A>|): R|kotlin/Unit| {
         R|<local>/m|.R|SubstitutionOverride</MyMap.get: R|A|>|(String(a)).R|/A.plusAssign|(Int(1))
+        Unit
     }

compiler/fir/analysis-tests/testData/resolve/arrays/arraySetWithOperation.fir.txt

@@ -46,6 +46,7 @@ FILE: arraySetWithOperation.kt
     }
     public final fun test_1(a: R|A<B>|): R|kotlin/Unit| {
         R|<local>/a|.R|SubstitutionOverride</A.get: R|B|>|(Int(0)).R|/B.plusAssign|(R|/B.B|())
+        Unit
     }
     public final fun test_2(a: R|A<C>|): R|kotlin/Unit| {
          {
@@ -54,10 +55,13 @@ FILE: arraySetWithOperation.kt
             R|<local>/<array>|.R|SubstitutionOverride</A.set: R|kotlin/Unit|>|(R|<local>/<index_0>|, R|<local>/<array>|.R|SubstitutionOverride</A.get: R|C|>|(R|<local>/<index_0>|).R|/C.plus|(R|/C.C|()))
         }
 
+        Unit
     }
     public final fun test_3(a: R|A<D>|): R|kotlin/Unit| {
         R|<local>/a|.R|SubstitutionOverride</A.get: R|D|>|(Int(0)).<Operator overload ambiguity. Compatible candidates: [/D.plusAssign, /A.set]>#(R|/D.D|())
+        Unit
     }
     public final fun test_4(b: R|B|): R|kotlin/Unit| {
         R|<local>/b|.<Unresolved name: get>#(Int(0)).<Unresolved name: +=>#(R|/B.B|())
+        Unit
     }

compiler/fir/analysis-tests/testData/resolve/callResolution/arraySetWithPlusAssign.fir.txt

@@ -20,8 +20,10 @@ FILE: arraySetWithPlusAssign.kt
                 Q|kotlin/Unit|
             }
             )
+            Unit
         }
         )
+        Unit
     }
     public final class B : R|kotlin/Any| {
         public constructor(): R|B| {
@@ -47,8 +49,10 @@ FILE: arraySetWithPlusAssign.kt
                 Q|kotlin/Unit|
             }
             )
+            Unit
         }
         )
+        Unit
     }
     public final class C : R|kotlin/Any| {
         public constructor(): R|C| {
@@ -77,8 +81,10 @@ FILE: arraySetWithPlusAssign.kt
                 Q|kotlin/Unit|
             }
             )
+            Unit
         }
         )
+        Unit
     }
     public final class D : R|kotlin/Any| {
         public constructor(): R|D| {
@@ -108,8 +114,10 @@ FILE: arraySetWithPlusAssign.kt
                 Q|kotlin/Unit|
             }
             )
+            Unit
         }
         )
+        Unit
     }
     public final class E : R|kotlin/Any| {
         public constructor(): R|E| {
@@ -118,15 +126,17 @@ FILE: arraySetWithPlusAssign.kt
 
     }
     public final fun test_5(x: R|E|): R|kotlin/Unit| {
-        R|<local>/x|.<Unresolved name: get>#(Int(1)).<Unresolved name: +=>#(fun <anonymous>(): <ERROR TYPE REF: Unresolved name: plusAssign> <inline=Unknown>  {
+        R|<local>/x|.<Unresolved name: get>#(Int(1)).<Unresolved name: +=>#(fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
             R|/someCallInsideLambda|()
-            ^ R|<local>/x|.<Unresolved name: get>#(Int(1)).<Unresolved name: +=>#(fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+            R|<local>/x|.<Unresolved name: get>#(Int(1)).<Unresolved name: +=>#(fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
                 R|/someCallInsideLambda|()
                 Q|kotlin/Unit|
             }
             )
+            Unit
         }
         )
+        Unit
     }
     public final class F : R|kotlin/Any| {
         public constructor(): R|F| {
@@ -154,6 +164,7 @@ FILE: arraySetWithPlusAssign.kt
             ^ String(please choose String)
         }
         )
+        Unit
     }
     public final class G : R|kotlin/Any| {
         public constructor(): R|G| {
@@ -192,10 +203,12 @@ FILE: arraySetWithPlusAssign.kt
                     ))
                 }
 
+                Unit
             }
             ))
         }
 
+        Unit
     }
     public final fun someCallInsideLambda(): R|kotlin/Unit| {
     }

compiler/fir/analysis-tests/testData/resolve/callResolution/exponentialArraySetWithPlusAssign.fir.txt

@@ -46,34 +46,49 @@ FILE: exponentialArraySetWithPlusAssign.kt
                                                                     Q|kotlin/Unit|
                                                                 }
                                                                 )
+                                                                Unit
                                                             }
                                                             )
+                                                            Unit
                                                         }
                                                         )
+                                                        Unit
                                                     }
                                                     )
+                                                    Unit
                                                 }
                                                 )
+                                                Unit
                                             }
                                             )
+                                            Unit
                                         }
                                         )
+                                        Unit
                                     }
                                     )
+                                    Unit
                                 }
                                 )
+                                Unit
                             }
                             )
+                            Unit
                         }
                         )
+                        Unit
                     }
                     )
+                    Unit
                 }
                 )
+                Unit
             }
             )
+            Unit
         }
         )
+        Unit
     }
     public final class D : R|kotlin/Any| {
         public constructor(): R|D| {
@@ -129,34 +144,49 @@ FILE: exponentialArraySetWithPlusAssign.kt
                                                                     Q|kotlin/Unit|
                                                                 }
                                                                 )
+                                                                Unit
                                                             }
                                                             )
+                                                            Unit
                                                         }
                                                         )
+                                                        Unit
                                                     }
                                                     )
+                                                    Unit
                                                 }
                                                 )
+                                                Unit
                                             }
                                             )
+                                            Unit
                                         }
                                         )
+                                        Unit
                                     }
                                     )
+                                    Unit
                                 }
                                 )
+                                Unit
                             }
                             )
+                            Unit
                         }
                         )
+                        Unit
                     }
                     )
+                    Unit
                 }
                 )
+                Unit
             }
             )
+            Unit
         }
         )
+        Unit
     }
     public final fun someCallInsideLambda(): R|kotlin/Unit| {
     }

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/safeCallsWithAugmentedAssignment.txt

@@ -4,10 +4,16 @@ FILE: safeCallsWithAugmentedAssignment.kt
         a#?.{ $subj$.b# }?.{ +=($subj$.c#, IntegerLiteral(1)) }
         +=(a#?.{ $subj$.b# }.c#, IntegerLiteral(1))
         a#?.{ ArraySet:[$subj$.b#.get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
+        Unit
         a#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
+        Unit
         ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)) += IntegerLiteral(1)]
+        Unit
         a#?.{ ArraySet:[$subj$.b#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
+        Unit
         a#?.{ $subj$.b# }?.{ ArraySet:[$subj$.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)) += IntegerLiteral(1)] }
+        Unit
         ArraySet:[a#?.{ $subj$.b# }.c#.get#(IntegerLiteral(0)).get#(IntegerLiteral(0)) += IntegerLiteral(1)]
+        Unit
         +=(a#?.{ $subj$.b# }.d#(), IntegerLiteral(1))
     }

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilder.kt

@@ -746,28 +746,33 @@ abstract class AbstractRawFirBuilder<T>(val baseSession: FirSession, val context
             if (operation == FirOperation.ASSIGN) {
                 context.arraySetArgument[unwrappedLhs] = rhsExpression
             }
-            return if (operation == FirOperation.ASSIGN) {
+            return buildBlock {
-                val result = unwrappedLhs.convert()
+                if (operation == FirOperation.ASSIGN) {
-                result.replaceAnnotations(result.annotations.smartPlus(annotations))
+                    val result = unwrappedLhs.convert()
-                buildBlock {
+                    result.replaceAnnotations(result.annotations.smartPlus(annotations))
                     source = result.source
                     statements += result.pullUpSafeCallIfNecessary()
-                    statements += buildUnitExpression {
+                } else {
-                        source = result.source?.fakeElement(KtFakeSourceElementKind.ImplicitUnit.IndexedAssignmentCoercion)
+                    val receiver = unwrappedLhs.convert()
+
+                    if (receiver is FirSafeCallExpression) {
+                        receiver.replaceSelector(
+                            generateAugmentedArraySetCall(
+                                receiver.selector as FirExpression, baseSource, arrayAccessSource, operation, annotations, rhsAST, convert
+                            )
+                        )
+                        source = receiver.source
+                        statements += receiver
+                    } else {
+                        val augmentedArraySetCall = generateAugmentedArraySetCall(
+                            receiver, baseSource, arrayAccessSource, operation, annotations, rhsAST, convert
+                        )
+                        source = augmentedArraySetCall.source
+                        statements += augmentedArraySetCall
                     }
                 }
-            } else {
+                statements += buildUnitExpression {
-                val receiver = unwrappedLhs.convert()
+                    source = this@buildBlock.source?.fakeElement(KtFakeSourceElementKind.ImplicitUnit.IndexedAssignmentCoercion)
-
-                if (receiver is FirSafeCallExpression) {
-                    receiver.replaceSelector(
-                        generateAugmentedArraySetCall(
-                            receiver.selector as FirExpression, baseSource, arrayAccessSource, operation, annotations, rhsAST, convert
-                        )
-                    )
-                    receiver
-                } else {
-                    generateAugmentedArraySetCall(receiver, baseSource, arrayAccessSource, operation, annotations, rhsAST, convert)
                 }
             }
         }

compiler/frontend.common/src/org/jetbrains/kotlin/KtSourceElement.kt

@@ -95,7 +95,7 @@ sealed class KtFakeSourceElementKind(final override val shouldSkipErrorTypeRepor
         // fake source refers to the return statement
         object Return : ImplicitUnit()
 
-        // this source is used for 'a[i] = b' converted to { a[i] = b; Unit }
+        // this source is used for 'a[i] = b' or 'a[i] += b' converted to { a[i] = b; Unit }
         // fake source refers to the assignment statement
         object IndexedAssignmentCoercion : ImplicitUnit()
     }

compiler/testData/diagnostics/testsWithStdLib/listAssignmentInWhen.fir.kt

@@ -10,3 +10,20 @@ fun bar(list: MutableList<Any?>, condition: Boolean): Unit = <!RETURN_TYPE_MISMA
     condition -> list.set(0, "")
     else -> Unit
 }<!>
+
+fun plusFoo(list: MutableList<String>, condition: Boolean): Unit = when {
+    condition -> list[0] += ""
+    else -> Unit
+}
+
+var x: String = ""
+
+fun assign(condition: Boolean): Unit = when {
+    condition -> x = ""
+    else -> Unit
+}
+
+fun plugAssign(condition: Boolean): Unit = when {
+    condition -> x += ""
+    else -> Unit
+}

compiler/testData/diagnostics/testsWithStdLib/listAssignmentInWhen.fir.txt

@@ -25,3 +25,48 @@ FILE: listAssignmentInWhen.fir.kt
         }
 
     }
+    public final fun plusFoo(list: R|kotlin/collections/MutableList<kotlin/String>|, condition: R|kotlin/Boolean|): R|kotlin/Unit| {
+        ^plusFoo when () {
+            R|<local>/condition| ->  {
+                 {
+                     {
+                        lval <array>: R|kotlin/collections/MutableList<kotlin/String>| = R|<local>/list|
+                        lval <index_0>: R|kotlin/Int| = Int(0)
+                        R|<local>/<array>|.R|SubstitutionOverride<kotlin/collections/MutableList.set: R|kotlin/String|>|(R|<local>/<index_0>|, R|<local>/<array>|.R|SubstitutionOverride<kotlin/collections/MutableList.get: R|kotlin/String|>|(R|<local>/<index_0>|).R|kotlin/String.plus|(String()))
+                    }
+
+                    Unit
+                }
+
+            }
+            else ->  {
+                Q|kotlin/Unit|
+            }
+        }
+
+    }
+    public final var x: R|kotlin/String| = String()
+        public get(): R|kotlin/String|
+        public set(value: R|kotlin/String|): R|kotlin/Unit|
+    public final fun assign(condition: R|kotlin/Boolean|): R|kotlin/Unit| {
+        ^assign when () {
+            R|<local>/condition| ->  {
+                R|/x| = String()
+            }
+            else ->  {
+                Q|kotlin/Unit|
+            }
+        }
+
+    }
+    public final fun plugAssign(condition: R|kotlin/Boolean|): R|kotlin/Unit| {
+        ^plugAssign when () {
+            R|<local>/condition| ->  {
+                R|/x| = R|/x|.R|kotlin/String.plus|(String())
+            }
+            else ->  {
+                Q|kotlin/Unit|
+            }
+        }
+
+    }

compiler/testData/diagnostics/testsWithStdLib/listAssignmentInWhen.kt

@@ -10,3 +10,20 @@ fun bar(list: MutableList<Any?>, condition: Boolean): Unit = <!TYPE_MISMATCH!>wh
     condition -> list.set(0, "")
     else -> Unit
 }<!>
+
+fun plusFoo(list: MutableList<String>, condition: Boolean): Unit = when {
+    condition -> list[0] += ""
+    else -> Unit
+}
+
+var x: String = ""
+
+fun assign(condition: Boolean): Unit = when {
+    condition -> x = ""
+    else -> Unit
+}
+
+fun plugAssign(condition: Boolean): Unit = when {
+    condition -> x += ""
+    else -> Unit
+}

compiler/testData/ir/irText/expressions/callableReferences/caoWithAdaptationForSam.fir.ir.txt

@@ -174,19 +174,20 @@ FILE fqName:<root> fileName:/caoWithAdaptationForSam.kt
           if: TYPE_OP type=kotlin.Boolean origin=INSTANCEOF typeOperand=<root>.IFoo
             GET_VAR 'fn: kotlin.Function1<kotlin.Int, kotlin.Unit> declared in <root>.test4' type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=null
           then: BLOCK type=kotlin.Unit origin=null
-            VAR IR_TEMPORARY_VARIABLE name:tmp_6 type:<root>.A [val]
+            BLOCK type=kotlin.Unit origin=null
-              GET_OBJECT 'CLASS OBJECT name:A modality:FINAL visibility:public superTypes:[kotlin.Any]' type=<root>.A
+              VAR IR_TEMPORARY_VARIABLE name:tmp_6 type:<root>.A [val]
-            VAR IR_TEMPORARY_VARIABLE name:tmp_7 type:<root>.IFoo [val]
+                GET_OBJECT 'CLASS OBJECT name:A modality:FINAL visibility:public superTypes:[kotlin.Any]' type=<root>.A
-              TYPE_OP type=<root>.IFoo origin=IMPLICIT_CAST typeOperand=<root>.IFoo
+              VAR IR_TEMPORARY_VARIABLE name:tmp_7 type:<root>.IFoo [val]
-                GET_VAR 'fn: kotlin.Function1<kotlin.Int, kotlin.Unit> declared in <root>.test4' type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=null
+                TYPE_OP type=<root>.IFoo origin=IMPLICIT_CAST typeOperand=<root>.IFoo
-            CALL 'public final fun set (i: <root>.IFoo, newValue: kotlin.Int): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                  GET_VAR 'fn: kotlin.Function1<kotlin.Int, kotlin.Unit> declared in <root>.test4' type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=null
-              $receiver: GET_VAR 'val tmp_6: <root>.A declared in <root>.test4' type=<root>.A origin=null
+              CALL 'public final fun set (i: <root>.IFoo, newValue: kotlin.Int): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-              i: GET_VAR 'val tmp_7: <root>.IFoo declared in <root>.test4' type=<root>.IFoo origin=null
+                $receiver: GET_VAR 'val tmp_6: <root>.A declared in <root>.test4' type=<root>.A origin=null
-              newValue: CALL 'public final fun plus (other: kotlin.Int): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
+                i: GET_VAR 'val tmp_7: <root>.IFoo declared in <root>.test4' type=<root>.IFoo origin=null
-                $this: CALL 'public final fun get (i: <root>.IFoo): kotlin.Int declared in <root>' type=kotlin.Int origin=null
+                newValue: CALL 'public final fun plus (other: kotlin.Int): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
-                  $receiver: GET_VAR 'val tmp_6: <root>.A declared in <root>.test4' type=<root>.A origin=null
+                  $this: CALL 'public final fun get (i: <root>.IFoo): kotlin.Int declared in <root>' type=kotlin.Int origin=null
-                  i: GET_VAR 'val tmp_7: <root>.IFoo declared in <root>.test4' type=<root>.IFoo origin=null
+                    $receiver: GET_VAR 'val tmp_6: <root>.A declared in <root>.test4' type=<root>.A origin=null
-                other: CONST Int type=kotlin.Int value=1
+                    i: GET_VAR 'val tmp_7: <root>.IFoo declared in <root>.test4' type=<root>.IFoo origin=null
+                  other: CONST Int type=kotlin.Int value=1
   FUN name:test5 visibility:public modality:FINAL <> (a:kotlin.Any) returnType:kotlin.Unit
     VALUE_PARAMETER name:a index:0 type:kotlin.Any
     BLOCK_BODY

compiler/testData/ir/irText/expressions/callableReferences/caoWithAdaptationForSam.fir.kt.txt

@@ -76,9 +76,11 @@ fun test3(fn: Function1<Int, Unit>) {
 fun test4(fn: Function1<Int, Unit>) {
   when {
     fn is IFoo -> { // BLOCK
-      val <array>: A = A
+      { // BLOCK
-      val <index_0>: IFoo = fn /*as IFoo */
+        val <array>: A = A
-      <array>.set(i = <index_0>, newValue = <array>.get(i = <index_0>).plus(other = 1))
+        val <index_0>: IFoo = fn /*as IFoo */
+        <array>.set(i = <index_0>, newValue = <array>.get(i = <index_0>).plus(other = 1))
+      }
     }
   }
 }