"Add not-null asserted (!!) call": add '!!' to receiver of function reference

#KT-37841 Fixed

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

@@ -48,6 +48,7 @@ import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.util.isValidOperator
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 
 abstract class ExclExclCallFix(psiElement: PsiElement) : KotlinQuickFixAction<PsiElement>(psiElement) {
@@ -96,10 +97,15 @@ class AddExclExclCallFix(psiElement: PsiElement, val checkImplicitReceivers: Boo
 
         val expr = getExpressionForIntroduceCall() ?: return
         val modifiedExpression = expr.expression
+        val psiFactory = KtPsiFactory(project)
         val exclExclExpression = if (expr.implicitReceiver) {
-            KtPsiFactory(project).createExpressionByPattern("this!!.$0", modifiedExpression)
+            if (modifiedExpression is KtCallableReferenceExpression) {
+                psiFactory.createExpressionByPattern("this!!::$0", modifiedExpression.callableReference)
+            } else {
+                psiFactory.createExpressionByPattern("this!!.$0", modifiedExpression)
+            }
         } else {
-            KtPsiFactory(project).createExpressionByPattern("$0!!", modifiedExpression)
+            psiFactory.createExpressionByPattern("$0!!", modifiedExpression)
         }
         modifiedExpression.replace(exclExclExpression)
     }
@@ -129,24 +135,26 @@ class AddExclExclCallFix(psiElement: PsiElement, val checkImplicitReceivers: Boo
                 }
             }
             is KtExpression -> {
+                val parent = psiElement.parent
                 val context = psiElement.analyze()
                 if (checkImplicitReceivers && psiElement.getResolvedCall(context)?.getImplicitReceiverValue() is ExtensionReceiver) {
-                    val expressionToReplace = psiElement.parent as? KtCallExpression ?: psiElement
+                    val expressionToReplace = parent as? KtCallExpression ?: parent as? KtCallableReferenceExpression ?: psiElement
                     expressionToReplace.expressionForCall(implicitReceiver = true)
                 } else {
-                    context[BindingContext.EXPRESSION_TYPE_INFO, psiElement]?.let {
+                    val targetElement = parent.safeAs<KtCallableReferenceExpression>()?.receiverExpression ?: psiElement
+                    context[BindingContext.EXPRESSION_TYPE_INFO, targetElement]?.let {
                         val type = it.type
 
-                        val dataFlowValueFactory = psiElement.getResolutionFacade().getDataFlowValueFactory()
+                        val dataFlowValueFactory = targetElement.getResolutionFacade().getDataFlowValueFactory()
 
                         if (type != null) {
                             val nullability = it.dataFlowInfo.getStableNullability(
-                                dataFlowValueFactory.createDataFlowValue(psiElement, type, context, psiElement.findModuleDescriptor())
+                                dataFlowValueFactory.createDataFlowValue(targetElement, type, context, targetElement.findModuleDescriptor())
                             )
                             if (!nullability.canBeNonNull()) return null
                         }
                     }
-                    psiElement.expressionForCall()
+                    targetElement.expressionForCall()
                 }
             }
             else -> null

idea/testData/quickfix/addExclExclCall/functionReference.kt

@@ -0,0 +1,8 @@
+// "Add non-null asserted (!!) call" "true"
+class Foo {
+    fun f() = 1
+}
+
+fun test(foo: Foo?) {
+    val f = foo::f<caret>
+}

idea/testData/quickfix/addExclExclCall/functionReference.kt.after

@@ -0,0 +1,8 @@
+// "Add non-null asserted (!!) call" "true"
+class Foo {
+    fun f() = 1
+}
+
+fun test(foo: Foo?) {
+    val f = foo!!::f
+}

idea/testData/quickfix/addExclExclCall/functionReference2.kt

@@ -0,0 +1,12 @@
+// "Add non-null asserted (!!) call" "true"
+class Foo {
+    val bar = Bar()
+}
+
+class Bar {
+    fun f() = 1
+}
+
+fun test(foo: Foo?) {
+    val f = foo?.bar::f<caret>
+}

idea/testData/quickfix/addExclExclCall/functionReference2.kt.after

@@ -0,0 +1,12 @@
+// "Add non-null asserted (!!) call" "true"
+class Foo {
+    val bar = Bar()
+}
+
+class Bar {
+    fun f() = 1
+}
+
+fun test(foo: Foo?) {
+    val f = foo?.bar!!::f
+}

idea/testData/quickfix/addExclExclCall/functionReference3.kt

@@ -0,0 +1,8 @@
+// "Add non-null asserted (!!) call" "true"
+class Foo {
+    fun f() = 1
+}
+
+fun Foo?.test() {
+    val f = ::f<caret>
+}

idea/testData/quickfix/addExclExclCall/functionReference3.kt.after

@@ -0,0 +1,8 @@
+// "Add non-null asserted (!!) call" "true"
+class Foo {
+    fun f() = 1
+}
+
+fun Foo?.test() {
+    val f = this!!::f
+}

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

@@ -641,6 +641,21 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/addExclExclCall/array4.kt");
         }
 
+        @TestMetadata("functionReference.kt")
+        public void testFunctionReference() throws Exception {
+            runTest("idea/testData/quickfix/addExclExclCall/functionReference.kt");
+        }
+
+        @TestMetadata("functionReference2.kt")
+        public void testFunctionReference2() throws Exception {
+            runTest("idea/testData/quickfix/addExclExclCall/functionReference2.kt");
+        }
+
+        @TestMetadata("functionReference3.kt")
+        public void testFunctionReference3() throws Exception {
+            runTest("idea/testData/quickfix/addExclExclCall/functionReference3.kt");
+        }
+
         @TestMetadata("implicit.kt")
         public void testImplicit() throws Exception {
             runTest("idea/testData/quickfix/addExclExclCall/implicit.kt");