[FIR] Fix collection of annotations for assigment expressions

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -293,7 +293,8 @@ class ExpressionsConverter(
                     binaryExpression.toFirSourceElement(),
                     rightArg,
                     rightArgAsFir,
-                    firOperation
+                    firOperation,
+                    leftArgAsFir.annotations
                 ) { getAsFirExpression(this) }
             } else {
                 buildEqualityOperatorCall {

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -2162,7 +2162,13 @@ open class RawFirBuilder(
             } else {
                 val firOperation = operationToken.toFirOperation()
                 if (firOperation in FirOperation.ASSIGNMENTS) {
-                    return expression.left.generateAssignment(source, expression.right, rightArgument, firOperation) {
+                    return expression.left.generateAssignment(
+                        source,
+                        expression.right,
+                        rightArgument,
+                        firOperation,
+                        leftArgument.annotations
+                    ) {
                         (this as KtExpression).toFirExpression("Incorrect expression in assignment: ${expression.text}")
                     }
                 } else {

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

@@ -540,7 +540,9 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                     generateResolvedAccessExpression(source, resultVar)
                 else
                     resultInitializer,
-                FirOperation.ASSIGN, convert
+                FirOperation.ASSIGN,
+                resultInitializer.annotations,
+                convert
             )
 
             fun appendAssignment() {
@@ -889,6 +891,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         rhs: T?,
         value: FirExpression, // value is FIR for rhs
         operation: FirOperation,
+        annotations: List<FirAnnotation>,
         convert: T.() -> FirExpression
     ): FirStatement {
         val unwrappedLhs = this.unwrap() ?: return buildErrorExpression {
@@ -901,9 +904,11 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                 context.arraySetArgument[unwrappedLhs] = value
             }
             return if (operation == FirOperation.ASSIGN) {
-                unwrappedLhs.convert()
+                val result = unwrappedLhs.convert()
+                (result.annotations as MutableList<FirAnnotation>) += annotations
+                result
             } else {
-                generateAugmentedArraySetCall(unwrappedLhs, baseSource, operation, rhs, convert)
+                generateAugmentedArraySetCall(unwrappedLhs, baseSource, operation, rhs, annotations, convert)
             }
         }
 
@@ -920,6 +925,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                     )
                 }
                 rightArgument = value
+                this.annotations += annotations
             }
         }
         require(operation == FirOperation.ASSIGN)
@@ -927,7 +933,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         if (this?.elementType == SAFE_ACCESS_EXPRESSION && this != null) {
             val safeCallNonAssignment = convert() as? FirSafeCallExpression
             if (safeCallNonAssignment != null) {
-                return putAssignmentToSafeCall(safeCallNonAssignment, baseSource, value)
+                return putAssignmentToSafeCall(safeCallNonAssignment, baseSource, value, annotations)
             }
         }
 
@@ -935,6 +941,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             source = baseSource
             rValue = value
             calleeReference = initializeLValue(unwrappedLhs) { convert() as? FirQualifiedAccess }
+            this.annotations += annotations
         }
     }
 
@@ -942,7 +949,8 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
     private fun putAssignmentToSafeCall(
         safeCallNonAssignment: FirSafeCallExpression,
         baseSource: KtSourceElement?,
-        value: FirExpression
+        value: FirExpression,
+        annotations: List<FirAnnotation>
     ): FirSafeCallExpression {
         val nestedAccess = safeCallNonAssignment.regularQualifiedAccess
 
@@ -951,6 +959,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             rValue = value
             calleeReference = nestedAccess.calleeReference
             explicitReceiver = safeCallNonAssignment.checkedSubjectRef.value
+            this.annotations += annotations
         }
 
         safeCallNonAssignment.replaceRegularQualifiedAccess(
@@ -965,6 +974,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         baseSource: KtSourceElement?,
         operation: FirOperation,
         rhs: T?,
+        annotations: List<FirAnnotation>,
         convert: T.() -> FirExpression
     ): FirStatement {
         return buildAugmentedArraySetCall {
@@ -972,6 +982,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             this.operation = operation
             assignCall = generateAugmentedCallForAugmentedArraySetCall(unwrappedReceiver, operation, rhs, convert)
             setGetBlock = generateSetGetBlockForAugmentedArraySetCall(unwrappedReceiver, baseSource, operation, rhs, convert)
+            this.annotations += annotations
         }
     }
 

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

@@ -535,6 +535,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                         dispatchReceiver = it.dispatchReceiver
                         extensionReceiver = it.extensionReceiver
                     }
+                    annotations += assignmentOperatorStatement.annotations
                 }
             return assignment.transform(transformer, ResolutionMode.ContextIndependent)
         }
@@ -1038,6 +1039,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         val operatorName = FirOperationNameConventions.ASSIGNMENTS.getValue(augmentedArraySetCall.operation)
 
         val firstCalls = with(augmentedArraySetCall.setGetBlock.statements.last() as FirFunctionCall) setCall@{
+            (annotations as MutableList<FirAnnotation>) += augmentedArraySetCall.annotations
             buildList {
                 add(this@setCall)
                 with(arguments.last() as FirFunctionCall) plusCall@{