FIR: Report UNSAFE_OPERATOR_CALL for augmented assignments (was

reporting UNSAFE_CALL).

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -19087,6 +19087,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/AssertNotNull.kt");
             }
 
+            @Test
+            @TestMetadata("augmentedAssignment.kt")
+            public void testAugmentedAssignment() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/augmentedAssignment.kt");
+            }
+
             @Test
             @TestMetadata("dataFlowInfoAfterExclExcl.kt")
             public void testDataFlowInfoAfterExclExcl() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java

@@ -19087,6 +19087,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/AssertNotNull.kt");
             }
 
+            @Test
+            @TestMetadata("augmentedAssignment.kt")
+            public void testAugmentedAssignment() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/augmentedAssignment.kt");
+            }
+
             @Test
             @TestMetadata("dataFlowInfoAfterExclExcl.kt")
             public void testDataFlowInfoAfterExclExcl() throws Exception {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/coneDiagnosticToFirDiagnostic.kt

@@ -100,11 +100,18 @@ private fun mapUnsafeCallError(
     val receiverExpression = candidate.callInfo.explicitReceiver
     val singleArgument = candidate.callInfo.argumentList.arguments.singleOrNull()
     if (receiverExpression != null && singleArgument != null &&
-        source.elementType == KtNodeTypes.OPERATION_REFERENCE &&
+        (source.elementType == KtNodeTypes.OPERATION_REFERENCE || source.elementType == KtNodeTypes.BINARY_EXPRESSION) &&
         (candidateFunction?.isOperator == true || candidateFunction?.isInfix == true)
     ) {
-        val operationToken = source.getChild(KtTokens.IDENTIFIER)
-        return if (operationToken != null) {
+        // For augmented assignment operations (e.g., `a += b`), the source is the entire binary expression (BINARY_EXPRESSION).
+        // TODO: No need to check for source.elementType == BINARY_EXPRESSION if we use operator as callee reference source
+        //  (see FirExpressionsResolveTransformer.transformAssignmentOperatorStatement)
+        val operationSource = if (source.elementType == KtNodeTypes.BINARY_EXPRESSION) {
+            source.getChild(KtNodeTypes.OPERATION_REFERENCE)
+        } else {
+            source
+        }
+        return if (operationSource?.getChild(KtTokens.IDENTIFIER) != null) {
             FirErrors.UNSAFE_INFIX_CALL.on(
                 source,
                 receiverExpression,

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

@@ -26,8 +26,6 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.*
 import org.jetbrains.kotlin.fir.resolve.inference.FirStubInferenceSession
 import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
-import org.jetbrains.kotlin.fir.resolve.substitution.substituteOrNull
-import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.*
@@ -41,6 +39,7 @@ import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
 open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) : FirPartialBodyResolveTransformer(transformer) {
     private inline val builtinTypes: BuiltinTypes get() = session.builtinTypes
@@ -372,6 +371,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             explicitReceiver = leftArgument
             argumentList = buildUnaryArgumentList(rightArgument)
             calleeReference = buildSimpleNamedReference {
+                // TODO: Use source of operator for callee reference source
                 source = assignmentOperatorStatement.source?.fakeElement(FirFakeSourceElementKind.DesugaredCompoundAssignment)
                 this.name = name
                 candidateSymbol = null

compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/augmentedAssignment.fir.kt

@@ -0,0 +1,17 @@
+class A {
+    operator fun plusAssign(s: String) {}
+}
+
+fun test() {
+    var a: A? = A()
+    a <!UNSAFE_OPERATOR_CALL!>+=<!> ""
+}
+
+class B {
+    operator fun plus(other: B) = this
+}
+
+fun test2() {
+    var b: B? = B()
+    b <!UNRESOLVED_REFERENCE!>+=<!> B()
+}

compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/augmentedAssignment.kt

@@ -0,0 +1,17 @@
+class A {
+    operator fun plusAssign(s: String) {}
+}
+
+fun test() {
+    var a: A? = A()
+    a <!UNSAFE_OPERATOR_CALL!>+=<!> ""
+}
+
+class B {
+    operator fun plus(other: B) = this
+}
+
+fun test2() {
+    var b: B? = B()
+    b <!UNSAFE_OPERATOR_CALL!>+=<!> B()
+}

compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/augmentedAssignment.txt

@@ -0,0 +1,20 @@
+package
+
+public fun test(): kotlin.Unit
+public fun test2(): kotlin.Unit
+
+public final class A {
+    public constructor A()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final operator fun plusAssign(/*0*/ s: kotlin.String): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class B {
+    public constructor B()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final operator fun plus(/*0*/ other: B): B
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -19093,6 +19093,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/AssertNotNull.kt");
             }
 
+            @Test
+            @TestMetadata("augmentedAssignment.kt")
+            public void testAugmentedAssignment() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/augmentedAssignment.kt");
+            }
+
             @Test
             @TestMetadata("dataFlowInfoAfterExclExcl.kt")
             public void testDataFlowInfoAfterExclExcl() throws Exception {

idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/ReplaceCallFixFactories.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.idea.quickfix.ReplaceInfixOrOperatorCallFix
 import org.jetbrains.kotlin.idea.quickfix.ReplaceWithSafeCallFix
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.unwrapParenthesesLabelsAndAnnotations
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
@@ -62,7 +63,7 @@ object ReplaceCallFixFactories {
                 return@diagnosticFixFactory emptyList()
             }
 
-            val target = psi.parent as? KtBinaryExpression ?: return@diagnosticFixFactory emptyList()
+            val target = psi.getNonStrictParentOfType<KtBinaryExpression>() ?: return@diagnosticFixFactory emptyList()
             listOf(ReplaceInfixOrOperatorCallFix(target, shouldHaveNotNullType(target), diagnostic.operator))
         }
 

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -16741,6 +16741,11 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/AssertNotNull.kt");
             }
 
+            @TestMetadata("augmentedAssignment.kt")
+            public void testAugmentedAssignment() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/augmentedAssignment.kt");
+            }
+
             @TestMetadata("dataFlowInfoAfterExclExcl.kt")
             public void testDataFlowInfoAfterExclExcl() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/dataFlowInfoAfterExclExcl.kt");

idea/testData/quickfix/replaceInfixOrOperatorCall/augmentedAssignmentAvailable.kt

@@ -6,6 +6,4 @@ class A {
 fun foo(b: A) {
     var a: A? = A()
     a <caret>+= b
-}
-
-/* IGNORE_FIR */
+}

idea/testData/quickfix/replaceInfixOrOperatorCall/augmentedAssignmentAvailable.kt.after

@@ -6,6 +6,4 @@ class A {
 fun foo(b: A) {
     var a: A? = A()
     a?.plusAssign(b)
-}
-
-/* IGNORE_FIR */
+}