diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/ReplaceCallFixFactories.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/ReplaceCallFixFactories.kt
index 91e3d0598f4..4ebf78f6316 100644
--- a/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/ReplaceCallFixFactories.kt
+++ b/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>)`.
-                // 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
+                // `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
             }
 
             when (val psi = diagnostic.psi) {
diff --git a/idea/idea-fir/tests/org/jetbrains/kotlin/idea/quickfix/HighLevelQuickFixTestGenerated.java b/idea/idea-fir/tests/org/jetbrains/kotlin/idea/quickfix/HighLevelQuickFixTestGenerated.java
index c8aacb10955..c61e1d6eca9 100644
--- a/idea/idea-fir/tests/org/jetbrains/kotlin/idea/quickfix/HighLevelQuickFixTestGenerated.java
+++ b/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");
diff --git a/idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt b/idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt
new file mode 100644
index 00000000000..3ffa75e66b1
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt.after b/idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt.after
new file mode 100644
index 00000000000..35fa4f738e5
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt.fir.after b/idea/testData/quickfix/replaceWithSafeCall/assignmentToPropertyWithNoExplicitType.kt.fir.after
new file mode 100644
index 00000000000..33fb3a59a34
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/comment.kt b/idea/testData/quickfix/replaceWithSafeCall/comment.kt
index 0c277614b1e..8f48b917d5d 100644
--- a/idea/testData/quickfix/replaceWithSafeCall/comment.kt
+++ b/idea/testData/quickfix/replaceWithSafeCall/comment.kt
@@ -5,5 +5,4 @@ fun foo(a: String?) {
             // comment2
             .<caret>length
 }
-// FIR_IDENTICAL
-/* FIR_COMPARISON */
+/* FIR_COMPARISON */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/comment.kt.after b/idea/testData/quickfix/replaceWithSafeCall/comment.kt.after
index 1ba6dfa318f..2b50aa5091c 100644
--- a/idea/testData/quickfix/replaceWithSafeCall/comment.kt.after
+++ b/idea/testData/quickfix/replaceWithSafeCall/comment.kt.after
@@ -5,5 +5,4 @@ fun foo(a: String?) {
         // comment2
         ?.length
 }
-// FIR_IDENTICAL
-/* FIR_COMPARISON */
+/* FIR_COMPARISON */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/comment.kt.fir.after b/idea/testData/quickfix/replaceWithSafeCall/comment.kt.fir.after
new file mode 100644
index 00000000000..f4a8b8f67a7
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt b/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt
new file mode 100644
index 00000000000..00e6be5e594
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt.after b/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt.after
new file mode 100644
index 00000000000..117f087173e
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt.fir.after b/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBody.kt.fir.after
new file mode 100644
index 00000000000..aea00874402
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt b/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt
new file mode 100644
index 00000000000..18f796f9ef8
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt.after b/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt.after
new file mode 100644
index 00000000000..dad32eb52e0
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt.fir.after b/idea/testData/quickfix/replaceWithSafeCall/functionExpressionBodyWithNoExplicitType.kt.fir.after
new file mode 100644
index 00000000000..774e80998b1
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt b/idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt
index 8170f5a2f99..dde03f04af6 100644
--- a/idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt.after b/idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt.after
index faafd1d4c6e..19f89e8c0e3 100644
--- a/idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt.after
+++ b/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 */
\ No newline at end of file
diff --git a/idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt.fir.after b/idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt.fir.after
new file mode 100644
index 00000000000..953fcb4017e
--- /dev/null
+++ b/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 */
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
index 7cfd0610865..3831e252bd1 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
+++ b/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");