FIR IDE: Use KtFirExpressionTypeProvider.getExpectedType() in

ReplaceCallFix.

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

@@ -18,7 +18,9 @@ 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.*
+import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 
 object ReplaceCallFixFactories {
     private val replaceWithSafeCallFixApplicator =
@@ -42,22 +44,9 @@ object ReplaceCallFixFactories {
         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>)`.
+                // `s.length` in `val x: Int = 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.
+                val expectedType = getExpectedType() as? KtTypeWithNullability
-                // FE1.0 doesn't either (see ReplaceWithSafeCallFixFactory).
+                return expectedType?.nullability == KtTypeNullability.NON_NULLABLE
-                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) {

idea/idea-fir/tests/org/jetbrains/kotlin/idea/quickfix/HighLevelQuickFixTestGenerated.java

@@ -916,6 +916,11 @@ public class HighLevelQuickFixTestGenerated extends AbstractHighLevelQuickFixTes
             runTest("idea/testData/quickfix/replaceWithSafeCall/assignmentToProperty.kt");
         }
 
+        @TestMetadata("assignmentToPropertyWithNoExplicitType.kt")
+        public void testAssignmentToPropertyWithNoExplicitType() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt");
+        }
+
         @TestMetadata("comment.kt")
         public void testComment() throws Exception {
             runTest("idea/testData/quickfix/replaceWithSafeCall/comment.kt");
@@ -936,6 +941,16 @@ public class HighLevelQuickFixTestGenerated extends AbstractHighLevelQuickFixTes
             runTest("idea/testData/quickfix/replaceWithSafeCall/functionCall.kt");
         }
 
+        @TestMetadata("functionExpressionBody.kt")
+        public void testFunctionExpressionBody() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt");
+        }
+
+        @TestMetadata("functionExpressionBodyWithNoExplicitType.kt")
+        public void testFunctionExpressionBodyWithNoExplicitType() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt");
+        }
+
         @TestMetadata("hasElvis.kt")
         public void testHasElvis() throws Exception {
             runTest("idea/testData/quickfix/replaceWithSafeCall/hasElvis.kt");

idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt

@@ -0,0 +1,6 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+class T(s: String?) {
+    var i = s<caret>.length
+}
+/* FIR_COMPARISON */

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

@@ -0,0 +1,6 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+class T(s: String?) {
+    var i = s?.length
+}
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt.fir.after

@@ -0,0 +1,6 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+class T(s: String?) {
+    var i = s?.length ?: <caret>
+}
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/comment.kt

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

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

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

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

@@ -0,0 +1,8 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a // comment1
+        // comment2
+        ?.length ?: <caret>
+}
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt

@@ -0,0 +1,9 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+// ERROR: Type mismatch: inferred type is Int? but Int was expected
+
+// Note: There should be no error in FIR but errors are not currently checked for FIR
+class T {
+    fun foo(s: String?): Int = s<caret>.length
+}
+/* FIR_COMPARISON */

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

@@ -0,0 +1,9 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+// ERROR: Type mismatch: inferred type is Int? but Int was expected
+
+// Note: There should be no error in FIR but errors are not currently checked for FIR
+class T {
+    fun foo(s: String?): Int = s?.length
+}
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt.fir.after

@@ -0,0 +1,9 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+// ERROR: Type mismatch: inferred type is Int? but Int was expected
+
+// Note: There should be no error in FIR but errors are not currently checked for FIR
+class T {
+    fun foo(s: String?): Int = s?.length ?: <caret>
+}
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt

@@ -0,0 +1,6 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+class T {
+    fun foo(s: String?) = s<caret>.length
+}
+/* FIR_COMPARISON */

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

@@ -0,0 +1,6 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+class T {
+    fun foo(s: String?) = s?.length
+}
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt.fir.after

@@ -0,0 +1,6 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+class T {
+    fun foo(s: String?) = s?.length ?: <caret>
+}
+/* FIR_COMPARISON */

idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt

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

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

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

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

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a
+        ?.length ?: <caret>
+}
+/* FIR_COMPARISON */

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -11826,6 +11826,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/replaceWithSafeCall/assignmentToProperty.kt");
         }
 
+        @TestMetadata("assignmentToPropertyWithNoExplicitType.kt")
+        public void testAssignmentToPropertyWithNoExplicitType() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt");
+        }
+
         @TestMetadata("comment.kt")
         public void testComment() throws Exception {
             runTest("idea/testData/quickfix/replaceWithSafeCall/comment.kt");
@@ -11846,6 +11851,16 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/replaceWithSafeCall/functionCall.kt");
         }
 
+        @TestMetadata("functionExpressionBody.kt")
+        public void testFunctionExpressionBody() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt");
+        }
+
+        @TestMetadata("functionExpressionBodyWithNoExplicitType.kt")
+        public void testFunctionExpressionBodyWithNoExplicitType() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt");
+        }
+
         @TestMetadata("hasElvis.kt")
         public void testHasElvis() throws Exception {
             runTest("idea/testData/quickfix/replaceWithSafeCall/hasElvis.kt");