"Wrap with let": fix for nullable extension function call #KT-6364 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/WrapWithSafeLetCallFix.kt

@@ -28,11 +28,12 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer
 import org.jetbrains.kotlin.psi.psiUtil.getLastParentOfTypeInRow
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
 import org.jetbrains.kotlin.types.typeUtil.isNullabilityMismatch
 
 class WrapWithSafeLetCallFix(
-        expression: KtExpression,
+    expression: KtExpression,
-        nullableExpression: KtExpression
+    nullableExpression: KtExpression
 ) : KotlinQuickFixAction<KtExpression>(expression) {
     private val nullableExpressionPointer = nullableExpression.createSmartPointer()
 
@@ -43,16 +44,19 @@ class WrapWithSafeLetCallFix(
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
         val element = element ?: return
         val nullableExpression = nullableExpressionPointer.element ?: return
+        val qualifiedExpression = element.getQualifiedExpressionForSelector()
+        val receiverExpression = qualifiedExpression?.receiverExpression
         val factory = KtPsiFactory(element)
         val nullableText = nullableExpression.text
         val validator = NewDeclarationNameValidator(element, nullableExpression, NewDeclarationNameValidator.Target.VARIABLES)
         val name = KotlinNameSuggester.suggestNameByName("it", validator)
         nullableExpression.replace(factory.createExpression(name))
+        val newExpression: Any = if (receiverExpression != null) "${receiverExpression.text}.${element.text}" else element
         val wrapped = when (name) {
-            "it" -> factory.createExpressionByPattern("$0?.let { $1 }", nullableText, element)
+            "it" -> factory.createExpressionByPattern("$0?.let { $1 }", nullableText, newExpression)
-            else -> factory.createExpressionByPattern("$0?.let { $1 -> $2 }", nullableText, name, element)
+            else -> factory.createExpressionByPattern("$0?.let { $1 -> $2 }", nullableText, name, newExpression)
         }
-        element.replace(wrapped)
+        (qualifiedExpression ?: element).replace(wrapped)
     }
 
     object UnsafeFactory : KotlinSingleIntentionActionFactory() {

idea/testData/quickfix/wrapWithSafeLetCall/extentionFunctionCall.kt

@@ -0,0 +1,6 @@
+// "Wrap with '?.let { ... }' call" "true"
+// WITH_RUNTIME
+
+fun f(s: String, action: (String.() -> Unit)?) {
+    s.action<caret>()
+}

idea/testData/quickfix/wrapWithSafeLetCall/extentionFunctionCall.kt.after

@@ -0,0 +1,6 @@
+// "Wrap with '?.let { ... }' call" "true"
+// WITH_RUNTIME
+
+fun f(s: String, action: (String.() -> Unit)?) {
+    action?.let { s.it() }
+}

idea/testData/quickfix/wrapWithSafeLetCall/extentionFunctionCall2.kt

@@ -0,0 +1,10 @@
+// "Wrap with '?.let { ... }' call" "true"
+// WITH_RUNTIME
+
+fun f(s: String, action: (String.() -> Unit)?) {
+    s.foo().bar().action<caret>()
+}
+
+fun String.foo() = ""
+
+fun String.bar() = ""

idea/testData/quickfix/wrapWithSafeLetCall/extentionFunctionCall2.kt.after

@@ -0,0 +1,10 @@
+// "Wrap with '?.let { ... }' call" "true"
+// WITH_RUNTIME
+
+fun f(s: String, action: (String.() -> Unit)?) {
+    action?.let { s.foo().bar().it() }
+}
+
+fun String.foo() = ""
+
+fun String.bar() = ""

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

@@ -12059,6 +12059,16 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/wrapWithSafeLetCall/extensionMethod.kt");
         }
 
+        @TestMetadata("extentionFunctionCall.kt")
+        public void testExtentionFunctionCall() throws Exception {
+            runTest("idea/testData/quickfix/wrapWithSafeLetCall/extentionFunctionCall.kt");
+        }
+
+        @TestMetadata("extentionFunctionCall2.kt")
+        public void testExtentionFunctionCall2() throws Exception {
+            runTest("idea/testData/quickfix/wrapWithSafeLetCall/extentionFunctionCall2.kt");
+        }
+
         @TestMetadata("insideLet.kt")
         public void testInsideLet() throws Exception {
             runTest("idea/testData/quickfix/wrapWithSafeLetCall/insideLet.kt");