FIR IDE: Enable ReplaceCallFix for UNSAFE_CALL.

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.idea.fir.api.fixes.KtQuickFixesList
 import org.jetbrains.kotlin.idea.fir.api.fixes.KtQuickFixesListBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
 import org.jetbrains.kotlin.idea.quickfix.fixes.ChangeTypeQuickFix
+import org.jetbrains.kotlin.idea.quickfix.fixes.ReplaceCallFixFactories
 
 class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
     private val modifiers = KtQuickFixesListBuilder.registerPsiQuickFix {
@@ -69,9 +70,14 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
         registerPsiQuickFixes(KtFirDiagnostic.MustBeInitialized::class, ChangeVariableMutabilityFix.MUST_BE_INITIALIZED_FACTORY)
     }
 
+    private val expressions = KtQuickFixesListBuilder.registerPsiQuickFix {
+        registerApplicator(ReplaceCallFixFactories.unsafeCallFactory)
+    }
+
     override val list: KtQuickFixesList = KtQuickFixesList.createCombined(
         modifiers,
         overrides,
         mutability,
+        expressions,
     )
 }

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

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.quickfix.fixes
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.KotlinBundle
+import org.jetbrains.kotlin.idea.fir.api.applicator.HLApplicatorInput
+import org.jetbrains.kotlin.idea.fir.api.applicator.applicatorByQuickFix
+import org.jetbrains.kotlin.idea.fir.api.fixes.HLApplicatorWithTargetAndInput
+import org.jetbrains.kotlin.idea.fir.api.fixes.diagnosticFixFactory
+import org.jetbrains.kotlin.idea.fir.api.fixes.withInput
+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.ReplaceCallFix
+import org.jetbrains.kotlin.idea.quickfix.ReplaceImplicitReceiverCallFix
+import org.jetbrains.kotlin.idea.quickfix.ReplaceWithSafeCallFix
+import org.jetbrains.kotlin.psi.*
+
+object ReplaceCallFixFactories {
+    private val replaceWithSafeCallFixApplicator =
+        applicatorByQuickFix<KtExpression, Input, ReplaceCallFix>(
+            getFamilyName = KotlinBundle.lazyMessage("replace.with.safe.call"),
+            isApplicableByPsi = { psi -> psi is KtDotQualifiedExpression }
+        ) { psi, input ->
+            ReplaceWithSafeCallFix(psi as KtDotQualifiedExpression, input.notNullNeeded)
+        }
+
+    private val replaceImplicitReceiverCallFixApplicator =
+        applicatorByQuickFix<KtExpression, Input, ReplaceImplicitReceiverCallFix>(
+            getFamilyName = KotlinBundle.lazyMessage("replace.with.safe.this.call")
+        ) { psi, input ->
+            ReplaceImplicitReceiverCallFix(psi, input.notNullNeeded)
+        }
+
+    class Input(val notNullNeeded: Boolean) : HLApplicatorInput
+
+    val unsafeCallFactory =
+        diagnosticFixFactory<PsiElement, KtFirDiagnostic.UnsafeCall, KtExpression, Input> { 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 `1 + s.length` with a safe call, it should be replaced with `(s.length ?: <caret>)`.
+                // NOTE: Even though we could, we don't check to see if the call is an argument and if the parameter type is nullable.
+                // FE1.0 doesn't either (see ReplaceWithSafeCallFixFactory).
+                val type = when (val parent = parent) {
+                    is KtBinaryExpression -> parent.left?.getKtType()
+                    is KtProperty -> {
+                        // Case: `val i = s.length`. `parent.getReturnKtType()` would return Int and therefore an elvis would be added,
+                        // which does not match FE1.0 behavior. We want to see if there is an explicit type, e.g., `val i: Int = s.length`
+                        // and default to the replacement being nullable if there is no explicit type.
+                        if (parent.typeReference != null) {
+                            parent.getReturnKtType()
+                        } else null
+                    }
+                    else -> null
+                }
+                return (type as? KtTypeWithNullability)?.nullability == KtTypeNullability.NON_NULLABLE
+            }
+
+            when (val psi = diagnostic.psi) {
+                is KtDotQualifiedExpression -> listOf(
+                    HLApplicatorWithTargetAndInput(replaceWithSafeCallFixApplicator, psi withInput Input(psi.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(
+                        HLApplicatorWithTargetAndInput(
+                            replaceImplicitReceiverCallFixApplicator,
+                            psi withInput Input(psi.shouldHaveNotNullType())
+                        )
+                    )
+                }
+                else -> emptyList()
+            }
+        }
+}

idea/testData/quickfix/replaceWithSafeCall/apply.kt

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         this<caret>.length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/apply.kt.after

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         this?.length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/applyWithImplicitParameter.kt

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         <caret>length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/applyWithImplicitParameter.kt.after

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         this?.length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignment.kt

@@ -5,3 +5,4 @@ var i = 0
 fun foo(s: String?) {
     i = s<caret>.length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignment.kt.after

@@ -5,3 +5,4 @@ var i = 0
 fun foo(s: String?) {
     i = s?.length ?: <caret>
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignmentFromImplicitParameter.kt

@@ -7,3 +7,4 @@ fun foo(a: String?) {
         i = <caret>length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignmentFromImplicitParameter.kt.after

@@ -7,3 +7,4 @@ fun foo(a: String?) {
         i = this?.length ?: <caret>
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignmentToNullable.kt

@@ -5,3 +5,4 @@ var i: Int? = 0
 fun foo(s: String?) {
     i = s<caret>.length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignmentToNullable.kt.after

@@ -5,3 +5,4 @@ var i: Int? = 0
 fun foo(s: String?) {
     i = s?.length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignmentToProperty.kt

@@ -3,3 +3,4 @@
 class T(s: String?) {
     var i: Int = s<caret>.length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignmentToProperty.kt.after

@@ -3,3 +3,4 @@
 class T(s: String?) {
     var i: Int = s?.length ?: <caret>
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/comment.kt

@@ -5,3 +5,4 @@ fun foo(a: String?) {
             // comment2
             .<caret>length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/comment.kt.after

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         // comment2
         ?.length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/expression.kt

@@ -3,3 +3,4 @@
 fun foo(s: String?) {
     1 + s<caret>.length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/expression.kt.after

@@ -3,3 +3,4 @@
 fun foo(s: String?) {
     1 + (s?.length ?: <caret>)
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/hasElvis.kt

@@ -8,3 +8,4 @@ fun main() {
 fun foo(): String? {
     return ""
 }
+/* FIR_COMPARISON */

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

@@ -8,3 +8,4 @@ fun main() {
 fun foo(): String? {
     return ""
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/hasElvis2.kt

@@ -7,3 +7,4 @@ fun foo(a: String?) {
         i = <caret>length ?: 0
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/hasElvis2.kt.after

@@ -7,3 +7,4 @@ fun foo(a: String?) {
         i = this?.length ?: 0
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/let.kt

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         it<caret>.length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/let.kt.after

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         it?.length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/letWithParameter.kt

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         b<caret>.length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/letWithParameter.kt.after

@@ -5,3 +5,4 @@ fun foo(a: String?) {
         b?.length
     }
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt

@@ -4,3 +4,4 @@ fun foo(a: String?) {
     val b = a
             .<caret>length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt.after

@@ -4,3 +4,4 @@ fun foo(a: String?) {
     val b = a
         ?.length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/normal.kt

@@ -3,3 +3,4 @@
 fun foo(a: String?) {
     a<caret>.length
 }
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/normal.kt.after

@@ -3,3 +3,4 @@
 fun foo(a: String?) {
     a?.length
 }
+/* FIR_COMPARISON */