FIR IDE: Enable ReplaceInfixOrOperatorCallFix for UNSAFE_CALL,

UNSAFE_INFIX_CALL, UNSAFE_OPERATOR_CALL.

idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/MainKtQuickFixRegistrar.kt

@@ -109,6 +109,8 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
         registerPsiQuickFixes(KtFirDiagnostic.UselessCast::class, RemoveUselessCastFix)
         registerPsiQuickFixes(KtFirDiagnostic.UselessIsCheck::class, RemoveUselessIsCheckFix, RemoveUselessIsCheckFixForWhen)
         registerApplicator(ReplaceCallFixFactories.unsafeCallFactory)
+        registerApplicator(ReplaceCallFixFactories.unsafeInfixCallFactory)
+        registerApplicator(ReplaceCallFixFactories.unsafeOperatorCallFactory)
         registerApplicator(AddExclExclCallFixFactories.unsafeCallFactory)
         registerApplicator(AddExclExclCallFixFactories.unsafeInfixCallFactory)
         registerApplicator(AddExclExclCallFixFactories.unsafeOperatorCallFactory)

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

@@ -6,25 +6,22 @@
 package org.jetbrains.kotlin.idea.quickfix.fixes
 
 import org.jetbrains.kotlin.idea.fir.api.fixes.diagnosticFixFactory
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
 import org.jetbrains.kotlin.idea.frontend.api.types.KtTypeNullability
 import org.jetbrains.kotlin.idea.frontend.api.types.KtTypeWithNullability
 import org.jetbrains.kotlin.idea.quickfix.ReplaceImplicitReceiverCallFix
+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.unwrapParenthesesLabelsAndAnnotations
+import org.jetbrains.kotlin.types.expressions.OperatorConventions
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 object ReplaceCallFixFactories {
     val unsafeCallFactory =
         diagnosticFixFactory<KtFirDiagnostic.UnsafeCall> { diagnostic ->
-            fun KtExpression.shouldHaveNotNullType(): Boolean {
-                // This function is used to determine if we may need to add an elvis operator after the safe call. For example, to replace
-                // `s.length` in `val x: Int = s.length` with a safe call, it should be replaced with `s.length ?: <caret>`.
-                val expectedType = getExpectedType() as? KtTypeWithNullability
-                return expectedType?.nullability == KtTypeNullability.NON_NULLABLE
-            }
-
             val psi = diagnostic.psi
             val target = if (psi is KtBinaryExpression && psi.operationToken in KtTokens.ALL_ASSIGNMENTS) {
                 // UNSAFE_CALL for assignments (e.g., `foo.bar = value`) is reported on the entire statement (KtBinaryExpression).
@@ -34,16 +31,45 @@ object ReplaceCallFixFactories {
                 psi
             }.unwrapParenthesesLabelsAndAnnotations()
 
+            val shouldHaveNotNullType = target.safeAs<KtExpression>()?.let { shouldHaveNotNullType(it) } ?: false
             when (target) {
-                is KtDotQualifiedExpression -> listOf(ReplaceWithSafeCallFix(target, target.shouldHaveNotNullType()))
+                is KtDotQualifiedExpression -> listOf(ReplaceWithSafeCallFix(target, shouldHaveNotNullType))
                 is KtNameReferenceExpression -> {
                     // TODO: As a safety precaution, resolve the expression to determine if it is a call with an implicit receiver.
                     // This is a defensive check to ensure that the diagnostic was reported on such a call and not some other name reference.
                     // This isn't strictly needed because FIR checkers aren't reporting on wrong elements, but ReplaceWithSafeCallFixFactory
                     // in FE1.0 does so.
-                    listOf(ReplaceImplicitReceiverCallFix(target, target.shouldHaveNotNullType()))
+                    listOf(ReplaceImplicitReceiverCallFix(target, shouldHaveNotNullType))
                 }
+                is KtArrayAccessExpression -> listOf(ReplaceInfixOrOperatorCallFix(target, shouldHaveNotNullType))
                 else -> emptyList()
             }
         }
+
+    val unsafeInfixCallFactory =
+        diagnosticFixFactory<KtFirDiagnostic.UnsafeInfixCall> { diagnostic ->
+            val psi = diagnostic.psi
+            val target = psi.parent as? KtBinaryExpression ?: return@diagnosticFixFactory emptyList()
+            listOf(ReplaceInfixOrOperatorCallFix(target, shouldHaveNotNullType(target), diagnostic.operator))
+        }
+
+    val unsafeOperatorCallFactory =
+        diagnosticFixFactory<KtFirDiagnostic.UnsafeOperatorCall> { diagnostic ->
+            val psi = diagnostic.psi
+            val operationToken = psi.safeAs<KtOperationReferenceExpression>()?.getReferencedNameElementType()
+            if (operationToken == KtTokens.EQ || operationToken in OperatorConventions.COMPARISON_OPERATIONS) {
+                // This matches FE1.0 behavior; see ReplaceInfixOrOperatorCallFixFactory.kt
+                return@diagnosticFixFactory emptyList()
+            }
+
+            val target = psi.parent as? KtBinaryExpression ?: return@diagnosticFixFactory emptyList()
+            listOf(ReplaceInfixOrOperatorCallFix(target, shouldHaveNotNullType(target), diagnostic.operator))
+        }
+
+    private fun KtAnalysisSession.shouldHaveNotNullType(expression: KtExpression): Boolean {
+        // This function is used to determine if we may need to add an elvis operator after the safe call. For example, to replace
+        // `s.length` in `val x: Int = s.length` with a safe call, it should be replaced with `s.length ?: <caret>`.
+        val expectedType = expression.getExpectedType() as? KtTypeWithNullability
+        return expectedType?.nullability == KtTypeNullability.NON_NULLABLE
+    }
 }

idea/testData/quickfix/nullables/unsafeInfixCall/unsafeGet.fir.kt.after

@@ -0,0 +1,4 @@
+// "Replace with safe (?.) call" "true"
+
+operator fun Int.get(row: Int, column: Int) = this
+fun foo(arg: Int?) = arg?.get(42, 13) ?: <caret>

idea/testData/quickfix/nullables/unsafeInfixCall/unsafeGet.kt

@@ -2,5 +2,3 @@
 
 operator fun Int.get(row: Int, column: Int) = this
 fun foo(arg: Int?) = arg<caret>[42, 13]
-
-/* IGNORE_FIR */

idea/testData/quickfix/nullables/unsafeInfixCall/unsafeGet.kt.after

@@ -2,5 +2,3 @@
 
 operator fun Int.get(row: Int, column: Int) = this
 fun foo(arg: Int?) = arg?.get(42, 13)
-
-/* IGNORE_FIR */

idea/testData/quickfix/nullables/unsafeInfixCall/unsafePlus.kt

@@ -2,5 +2,3 @@
 
 operator fun Int.plus(index: Int) = this
 fun fox(arg: Int?) = arg <caret>+ 42
-
-/* IGNORE_FIR */

idea/testData/quickfix/nullables/unsafeInfixCall/unsafePlus.kt.after

@@ -2,5 +2,3 @@
 
 operator fun Int.plus(index: Int) = this
 fun fox(arg: Int?) = arg?.plus(42)
-
-/* IGNORE_FIR */

idea/testData/quickfix/nullables/unsafeInfixCall/unsafeSet.kt

@@ -4,5 +4,3 @@ operator fun Int.set(row: Int, column: Int, value: Int) {}
 fun foo(arg: Int?) {
     arg<caret>[42, 13] = 0
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/nullables/unsafeInfixCall/unsafeSet.kt.after

@@ -4,5 +4,3 @@ operator fun Int.set(row: Int, column: Int, value: Int) {}
 fun foo(arg: Int?) {
     arg?.set(42, 13, 0)
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/replaceInfixOrOperatorCall/array.kt

@@ -4,5 +4,3 @@
 fun foo(array: Array<String>?) {
     array<caret>[0]
 }
-
-/* IGNORE_FIR */

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

@@ -4,5 +4,3 @@
 fun foo(array: Array<String>?) {
     array?.get(0)
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/replaceInfixOrOperatorCall/arraySet.kt

@@ -4,5 +4,3 @@
 fun foo(array: Array<String>?) {
     array<caret>[0] = ""
 }
-
-/* IGNORE_FIR */

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

@@ -4,5 +4,3 @@
 fun foo(array: Array<String>?) {
     array?.set(0, "")
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentArray.kt

@@ -5,5 +5,3 @@ fun foo(array: Array<String>?) {
     var s = ""
     s = array[0]<caret>
 }
-
-/* IGNORE_FIR */

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

@@ -5,5 +5,3 @@ fun foo(array: Array<String>?) {
     var s = ""
     s = array?.get(0) ?: <caret>
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentBinaryOperator.kt

@@ -5,5 +5,3 @@ fun foo(bar: Int?) {
     var i: Int = 1
     i = bar +<caret> 1
 }
-
-/* IGNORE_FIR */

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

@@ -5,5 +5,3 @@ fun foo(bar: Int?) {
     var i: Int = 1
     i = bar?.plus(1) ?: <caret>
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentList.kt

@@ -5,5 +5,3 @@ fun foo(list: List<String>?) {
     var s = ""
     s = list[0]<caret>
 }
-
-/* IGNORE_FIR */

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

@@ -5,5 +5,3 @@ fun foo(list: List<String>?) {
     var s = ""
     s = list?.get(0) ?: <caret>
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/replaceInfixOrOperatorCall/binaryOperator.kt

@@ -4,5 +4,3 @@
 fun foo(bar: Int?) {
     bar +<caret> 1
 }
-
-/* IGNORE_FIR */

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

@@ -4,5 +4,3 @@
 fun foo(bar: Int?) {
     bar?.plus(1)
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/replaceInfixOrOperatorCall/hasElvis.kt

@@ -5,5 +5,3 @@ fun foo(list: List<String>?) {
     var s = ""
     s = list[0]<caret> ?: ""
 }
-
-/* IGNORE_FIR */

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

@@ -5,5 +5,3 @@ fun foo(list: List<String>?) {
     var s = ""
     s = list?.get(0) ?: ""
 }
-
-/* IGNORE_FIR */

idea/testData/quickfix/replaceInfixOrOperatorCall/list.kt

@@ -4,5 +4,3 @@
 fun foo(list: List<String>?) {
     list<caret>[0]
 }
-
-/* IGNORE_FIR */

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

@@ -4,5 +4,3 @@
 fun foo(list: List<String>?) {
     list?.get(0)
 }
-
-/* IGNORE_FIR */