KT-35042 Fix InsertHandlerProvider.kt to work with suspend functional parameters

- ^KT-35042 Fixed
- Also, fix review suggestions for KOTLIN-CR-3495

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionUtils.kt

@@ -250,13 +250,13 @@ fun returnExpressionItems(bindingContext: BindingContext, position: KtElement):
 }
 
 private fun KtElement.isDirectlyInLoopBody(): Boolean {
-    val possibleLoop = when (parent) {
-        is KtBlockExpression -> parent.parent?.parent
-        is KtContainerNodeForControlStructureBody -> parent.parent
+    val loopContainer = when (val blockOrContainer = parent) {
+        is KtBlockExpression -> blockOrContainer.parent as? KtContainerNodeForControlStructureBody
+        is KtContainerNodeForControlStructureBody -> blockOrContainer
         else -> null
     }
 
-    return possibleLoop is KtLoopExpression
+    return loopContainer?.parent is KtLoopExpression
 }
 
 fun KtElement.isRightOperandInElvis(): Boolean {
@@ -292,9 +292,9 @@ private fun KtElement.inReturnExpression(): Boolean = findReturnExpression(this)
 private tailrec fun findReturnExpression(expression: PsiElement?): KtReturnExpression? =
     when (val parent = expression?.parent) {
         is KtReturnExpression -> parent
-        is KtBinaryExpression -> findReturnExpression(parent.takeIf { it.operationToken == KtTokens.ELVIS })
+        is KtBinaryExpression -> if (parent.operationToken == KtTokens.ELVIS) findReturnExpression(parent) else null
         is KtContainerNodeForControlStructureBody, is KtIfExpression -> findReturnExpression(parent)
-        is KtBlockExpression -> findReturnExpression(parent.takeIf { expression.isLastOrSingleStatement() })
+        is KtBlockExpression -> if (expression.isLastOrSingleStatement()) findReturnExpression(parent) else null
         is KtWhenEntry -> findReturnExpression(parent.parent)
         else -> null
     }

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/InsertHandlerProvider.kt

@@ -20,7 +20,6 @@ import com.intellij.codeInsight.completion.InsertHandler
 import com.intellij.codeInsight.lookup.LookupElement
 import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
 import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
-import org.jetbrains.kotlin.builtins.isFunctionType
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.completion.handlers.*
 import org.jetbrains.kotlin.idea.core.ExpectedInfo
@@ -102,7 +101,7 @@ class InsertHandlerProvider(
                 descriptor.upperBounds.filter { it.arguments.isNotEmpty() }.forEach(::addPotentiallyInferred)
             }
 
-            if (type.isFunctionType && getValueParametersCountFromFunctionType(type) <= 1) {
+            if (type.isBuiltinFunctionalType && getValueParametersCountFromFunctionType(type) <= 1) {
                 // do not rely on inference from input of function type with one or no arguments - use only return type of functional type
                 addPotentiallyInferred(type.getReturnTypeFromFunctionType())
                 return

idea/idea-completion/testData/handlers/basic/typeArgsForCall/FunctionTypeParameter3.kt

@@ -0,0 +1,7 @@
+fun <T1, T2> T1.foo(handler: suspend (T2) -> Boolean) {}
+
+fun f() {
+    "".<caret>
+}
+
+// ELEMENT: foo

idea/idea-completion/testData/handlers/basic/typeArgsForCall/FunctionTypeParameter3.kt.after

@@ -0,0 +1,7 @@
+fun <T1, T2> T1.foo(handler: suspend (T2) -> Boolean) {}
+
+fun f() {
+    "".foo<<caret>> {  }
+}
+
+// ELEMENT: foo

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/handlers/BasicCompletionHandlerTestGenerated.java

@@ -973,6 +973,11 @@ public class BasicCompletionHandlerTestGenerated extends AbstractBasicCompletion
             runTest("idea/idea-completion/testData/handlers/basic/typeArgsForCall/FunctionTypeParameter2.kt");
         }
 
+        @TestMetadata("FunctionTypeParameter3.kt")
+        public void testFunctionTypeParameter3() throws Exception {
+            runTest("idea/idea-completion/testData/handlers/basic/typeArgsForCall/FunctionTypeParameter3.kt");
+        }
+
         @TestMetadata("HasExpectedType.kt")
         public void testHasExpectedType() throws Exception {
             runTest("idea/idea-completion/testData/handlers/basic/typeArgsForCall/HasExpectedType.kt");

idea/performanceTests/org/jetbrains/kotlin/idea/perf/PerformanceBasicCompletionHandlerTestGenerated.java

@@ -973,6 +973,11 @@ public class PerformanceBasicCompletionHandlerTestGenerated extends AbstractPerf
             runTest("idea/idea-completion/testData/handlers/basic/typeArgsForCall/FunctionTypeParameter2.kt");
         }
 
+        @TestMetadata("FunctionTypeParameter3.kt")
+        public void testFunctionTypeParameter3() throws Exception {
+            runTest("idea/idea-completion/testData/handlers/basic/typeArgsForCall/FunctionTypeParameter3.kt");
+        }
+
         @TestMetadata("HasExpectedType.kt")
         public void testHasExpectedType() throws Exception {
             runTest("idea/idea-completion/testData/handlers/basic/typeArgsForCall/HasExpectedType.kt");