Compute type for callable references and lambdas inside not null assertion if couldn't get already recorded type

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/BasicExpressionTypingVisitor.java

@@ -823,10 +823,7 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
             DiagnosticFactory0<PsiElement> diagnosticFactory =
                     isFunctionLiteral ? NOT_NULL_ASSERTION_ON_LAMBDA_EXPRESSION : NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE;
             context.trace.report(diagnosticFactory.on(operationSign));
-            if (baseTypeInfo == null) {
+            return baseTypeInfo != null ? baseTypeInfo : components.expressionTypingServices.getTypeInfo(baseExpression, context);
-                return TypeInfoFactoryKt.createTypeInfo(ErrorUtils.createErrorType("Unresolved lambda expression"), context);
-            }
-            return baseTypeInfo;
         }
         assert baseTypeInfo != null : "Base expression was not processed: " + expression;
         KotlinType baseType = baseTypeInfo.getType();

compiler/testData/codegen/box/inference/builderInference/specialCallsWithCallableReferencesErrorType.kt

@@ -0,0 +1,113 @@
+// WITH_RUNTIME
+
+import kotlin.experimental.ExperimentalTypeInference
+
+fun <K> FlowCollector<K>.bar(): K = null as K
+fun <K> FlowCollector<K>.foo(): K = null as K
+
+fun bar2(): Int = 1
+fun foo2(): Float = 1f
+
+fun <T> materialize() = null as T
+
+interface FlowCollector<in T> {}
+
+@Suppress("EXPERIMENTAL_API_USAGE_ERROR")
+fun <L> flow(@BuilderInference block: suspend FlowCollector<L>.() -> Unit) = Flow(block)
+
+class Flow<out R>(private val block: suspend FlowCollector<R>.() -> Unit)
+
+fun <R> select(vararg x: R) = x[0]
+
+fun poll01(): Flow<String> {
+    return flow {
+        val inv = select(::bar2, ::foo2)
+        inv()
+    }
+}
+
+fun poll1(flag: Boolean): Flow<String> {
+    return flow {
+        val inv = if (flag) { ::bar2 } else { ::foo2 }
+        inv()
+    }
+}
+
+fun poll11(flag: Boolean): Flow<String> {
+    return flow {
+        val inv = if (flag) { ::bar2 } else { ::foo2 }
+        inv()
+    }
+}
+
+
+fun poll21(flag: Boolean): Flow<String> {
+    return flow {
+        val inv = when (flag) { true -> ::bar2 else -> ::foo2 }
+        inv()
+    }
+}
+
+fun poll31(flag: Boolean): Flow<String> {
+    return flow {
+        val inv = when (flag) { true -> ::bar2 false -> ::foo2 }
+        inv()
+    }
+}
+
+fun poll41(): Flow<String> {
+    return flow {
+        val inv = try { ::bar2 } finally { ::foo2 }
+        inv()
+    }
+}
+
+fun poll51(): Flow<String> {
+    return flow {
+        val inv = try { ::bar2 } catch (e: Exception) { ::foo2 } finally { ::foo2 }
+        inv()
+    }
+}
+
+fun poll61(): Flow<String> {
+    return flow {
+        val inv = ::bar2
+        inv
+    }
+}
+
+fun poll71(): Flow<String> {
+    return flow {
+        val inv = ::bar2!!
+        inv()
+    }
+}
+
+fun poll81(): Flow<String> {
+    return flow {
+        val inv = ::bar2 in setOf(::foo2)
+        inv
+    }
+}
+
+fun poll91(): Flow<String> {
+    return flow {
+        val inv = ::foo2 in setOf(::foo2)
+        inv
+    }
+}
+
+fun box(): String {
+    poll01()
+    poll1(true)
+    poll11(true)
+    poll21(true)
+    poll31(true)
+    poll41()
+    poll51()
+    poll61()
+    poll71()
+    poll81()
+    poll91()
+    return "OK"
+}

compiler/testData/diagnostics/tests/inference/builderInference/specialCallsWithCallableReferencesErrorType.fir.kt

@@ -0,0 +1,86 @@
+// WITH_RUNTIME
+// SKIP_TXT
+// !DIAGNOSTICS: -CAST_NEVER_SUCCEEDS -UNCHECKED_CAST -UNUSED_PARAMETER -UNUSED_VARIABLE -EXPERIMENTAL_API_USAGE_ERROR -UNUSED_EXPRESSION
+
+import kotlin.experimental.ExperimentalTypeInference
+
+fun <K> FlowCollector<K>.bar(): K = null as K
+fun <K> FlowCollector<K>.foo(): K = null as K
+
+fun <K> K.bar3(): K = null as K
+fun <K> K.foo3(): K = null as K
+
+fun bar2(): Int = 1
+fun foo2(): Float = 1f
+
+val bar4: Int
+    get() = 1
+
+var foo4: Float
+    get() = 1f
+    set(value) {}
+
+val <K> FlowCollector<K>.bar5: K get() = null as K
+val <K> FlowCollector<K>.foo5: K get() = null as K
+
+class Foo6
+
+class Foo7<T>
+fun foo7() = null as Foo7<Int>
+
+interface FlowCollector<in T> {}
+
+fun <L> flow(@BuilderInference block: suspend FlowCollector<L>.() -> Unit) = Flow(block)
+
+class Flow<out R>(private val block: suspend FlowCollector<R>.() -> Unit)
+
+
+fun poll7(): Flow<String> {
+    return flow {
+        val inv = ::bar!!
+        <!INAPPLICABLE_CANDIDATE!>inv<!>()
+    }
+}
+
+fun poll71(): Flow<String> {
+    return flow {
+        val inv = ::bar2!!
+        inv()
+    }
+}
+
+fun poll72(): Flow<String> {
+    return flow {
+        val inv = ::bar3!!
+        inv()
+    }
+}
+
+fun poll73(): Flow<String> {
+    return flow {
+        val inv = ::bar4!!
+        inv
+    }
+}
+
+fun poll74(): Flow<String> {
+    return flow {
+        val inv = ::bar5!!
+        inv
+    }
+}
+
+fun poll75(): Flow<String> {
+    return flow {
+        val inv = ::Foo6!!
+        inv
+    }
+}
+
+fun poll76(): Flow<String> {
+    return flow {
+        val inv = ::Foo7!!
+        inv
+    }
+}
+

compiler/testData/diagnostics/tests/inference/builderInference/specialCallsWithCallableReferencesErrorType.kt

@@ -0,0 +1,86 @@
+// WITH_RUNTIME
+// SKIP_TXT
+// !DIAGNOSTICS: -CAST_NEVER_SUCCEEDS -UNCHECKED_CAST -UNUSED_PARAMETER -UNUSED_VARIABLE -EXPERIMENTAL_API_USAGE_ERROR -UNUSED_EXPRESSION
+
+import kotlin.experimental.ExperimentalTypeInference
+
+fun <K> FlowCollector<K>.bar(): K = null as K
+fun <K> FlowCollector<K>.foo(): K = null as K
+
+fun <K> K.bar3(): K = null as K
+fun <K> K.foo3(): K = null as K
+
+fun bar2(): Int = 1
+fun foo2(): Float = 1f
+
+val bar4: Int
+    get() = 1
+
+var foo4: Float
+    get() = 1f
+    set(value) {}
+
+val <K> FlowCollector<K>.bar5: K get() = null as K
+val <K> FlowCollector<K>.foo5: K get() = null as K
+
+class Foo6
+
+class Foo7<T>
+fun foo7() = null as Foo7<Int>
+
+interface FlowCollector<in T> {}
+
+fun <L> flow(@BuilderInference block: suspend FlowCollector<L>.() -> Unit) = Flow(block)
+
+class Flow<out R>(private val block: suspend FlowCollector<R>.() -> Unit)
+
+
+fun poll7(): Flow<String> {
+    return flow {
+        val inv = <!TYPE_INFERENCE_POSTPONED_VARIABLE_IN_RECEIVER_TYPE!>::bar<!><!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        <!DEBUG_INFO_MISSING_UNRESOLVED!>inv()<!>
+    }
+}
+
+fun poll71(): Flow<String> {
+    return flow {
+        val inv = ::bar2<!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        inv()
+    }
+}
+
+fun poll72(): Flow<String> {
+    return flow {
+        val inv = <!TYPE_INFERENCE_POSTPONED_VARIABLE_IN_RECEIVER_TYPE!>::bar3<!><!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        <!DEBUG_INFO_MISSING_UNRESOLVED!>inv()<!>
+    }
+}
+
+fun poll73(): Flow<String> {
+    return flow {
+        val inv = ::bar4<!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        inv
+    }
+}
+
+fun poll74(): Flow<String> {
+    return flow {
+        val inv = <!TYPE_INFERENCE_POSTPONED_VARIABLE_IN_RECEIVER_TYPE!>::bar5<!><!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE!>inv<!>
+    }
+}
+
+fun poll75(): Flow<String> {
+    return flow {
+        val inv = ::Foo6<!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        inv
+    }
+}
+
+fun poll76(): Flow<String> {
+    return flow {
+        val inv = <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>::<!TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>Foo7<!><!><!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE!>inv<!>
+    }
+}
+

compiler/testData/ir/irText/expressions/kt36963.txt

@@ -1,9 +1,9 @@
 FILE fqName:<root> fileName:/kt36963.kt
   FUN name:foo visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
-  FUN name:test visibility:public modality:FINAL <> () returnType:IrErrorType
+  FUN name:test visibility:public modality:FINAL <> () returnType:kotlin.reflect.KFunction0<kotlin.Unit>
     BLOCK_BODY
-      RETURN type=kotlin.Nothing from='public final fun test (): IrErrorType declared in <root>'
+      RETURN type=kotlin.Nothing from='public final fun test (): kotlin.reflect.KFunction0<kotlin.Unit> declared in <root>'
         CALL 'public final fun CHECK_NOT_NULL <T0> (arg0: T0 of kotlin.internal.ir.CHECK_NOT_NULL?): T0 of kotlin.internal.ir.CHECK_NOT_NULL declared in kotlin.internal.ir' type=kotlin.reflect.KFunction0<kotlin.Unit> origin=EXCLEXCL
           <T0>: kotlin.reflect.KFunction0<kotlin.Unit>
           arg0: FUNCTION_REFERENCE 'public final fun foo (): kotlin.Unit declared in <root>' type=kotlin.reflect.KFunction0<kotlin.Unit> origin=null reflectionTarget=<same>

compiler/tests-gen/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -13501,6 +13501,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @TestMetadata("specialCallsWithCallableReferencesErrorType.kt")
+            public void testSpecialCallsWithCallableReferencesErrorType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/specialCallsWithCallableReferencesErrorType.kt");
+            }
+
             @TestMetadata("substituteStubTypeIntoCR.kt")
             public void testSubstituteStubTypeIntoCR() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteStubTypeIntoCR.kt");