Reuse information from already recorder descriptor for some callable references and don't rewrite at slice

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt

@@ -11,10 +11,7 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl
 import org.jetbrains.kotlin.descriptors.impl.ReceiverParameterDescriptorImpl
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
-import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.ValueArgument
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.MissingSupertypesResolver
@@ -24,6 +21,7 @@ import org.jetbrains.kotlin.resolve.calls.NewCommonSuperTypeCalculator
 import org.jetbrains.kotlin.resolve.calls.checkers.CallCheckerContext
 import org.jetbrains.kotlin.resolve.calls.commonSuperType
 import org.jetbrains.kotlin.resolve.calls.components.CallableReferenceAdaptation
+import org.jetbrains.kotlin.resolve.calls.components.CallableReferenceCandidate
 import org.jetbrains.kotlin.resolve.calls.components.SuspendConversionStrategy
 import org.jetbrains.kotlin.resolve.calls.components.isVararg
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
@@ -69,6 +67,14 @@ class ResolvedAtomCompleter(
     )
     private val topLevelTrace = topLevelCallCheckerContext.trace
 
+    private data class CallableReferenceResultTypeInfo(
+        val dispatchReceiver: ReceiverValue?,
+        val extensionReceiver: ReceiverValue?,
+        val explicitReceiver: ReceiverValue?,
+        val substitutor: TypeSubstitutor,
+        val resultType: KotlinType
+    )
+
     private fun complete(resolvedAtom: ResolvedAtom) {
         if (topLevelCallContext.inferenceSession.callCompleted(resolvedAtom)) {
             return
@@ -328,21 +334,16 @@ class ResolvedAtomCompleter(
         }
     }
 
-    private fun completeCallableReference(
+    private fun updateCallableReferenceResultType(
-        resolvedAtom: ResolvedCallableReferenceAtom
+        callableCandidate: CallableReferenceCandidate,
-    ) {
+        callableReferenceExpression: KtCallableReferenceExpression
-        val callableCandidate = resolvedAtom.candidate
+    ): CallableReferenceResultTypeInfo {
-        if (callableCandidate == null || resolvedAtom.completed) {
-            // todo report meanfull diagnostic here
-            return
-        }
         val resultTypeParameters =
             callableCandidate.freshSubstitutor!!.freshVariables.map { resultSubstitutor.safeSubstitute(it.defaultType) }
 
-        val typeParametersSubstitutor =
+        val typeParametersSubstitutor = NewTypeSubstitutorByConstructorMap(
-            NewTypeSubstitutorByConstructorMap(
+            (callableCandidate.candidate.typeParameters.map { it.typeConstructor } zip resultTypeParameters).toMap()
-                (callableCandidate.candidate.typeParameters.map { it.typeConstructor } zip resultTypeParameters).toMap()
+        )
-            )
 
         val resultSubstitutor = if (callableCandidate.candidate.isSupportedForCallableReference()) {
             val firstSubstitution = typeParametersSubstitutor.toOldSubstitution()
@@ -350,49 +351,16 @@ class ResolvedAtomCompleter(
             TypeSubstitutor.createChainedSubstitutor(firstSubstitution, secondSubstitution)
         } else TypeSubstitutor.EMPTY
 
-        val psiCallArgument = resolvedAtom.atom.psiCallArgument as CallableReferenceKotlinCallArgumentImpl
-        val callableReferenceExpression = psiCallArgument.ktCallableReferenceExpression
-
         // write down type for callable reference expression
         val resultType = resultSubstitutor.safeSubstitute(callableCandidate.reflectionCandidateType, Variance.INVARIANT)
+
         argumentTypeResolver.updateResultArgumentTypeIfNotDenotable(
-            topLevelTrace, expressionTypingServices.statementFilter,
+            topLevelTrace, expressionTypingServices.statementFilter, resultType, callableReferenceExpression
-            resultType,
-            callableReferenceExpression
         )
-        val reference = callableReferenceExpression.callableReference
-
-        val explicitCallableReceiver = when (callableCandidate.explicitReceiverKind) {
-            ExplicitReceiverKind.DISPATCH_RECEIVER -> callableCandidate.dispatchReceiver
-            ExplicitReceiverKind.EXTENSION_RECEIVER -> callableCandidate.extensionReceiver
-            else -> null
-        }
-
-        val explicitReceiver = explicitCallableReceiver?.receiver?.receiverValue?.updateReceiverValue(resultSubstitutor)
-        val psiCall = CallMaker.makeCall(reference, explicitReceiver, null, reference, emptyList())
-
-        val tracing = TracingStrategyImpl.create(reference, psiCall)
-        val temporaryTrace = TemporaryBindingTrace.create(topLevelTrace, "callable reference fake call")
 
         val dispatchReceiver = callableCandidate.dispatchReceiver?.receiver?.receiverValue?.updateReceiverValue(resultSubstitutor)
         val extensionReceiver = callableCandidate.extensionReceiver?.receiver?.receiverValue?.updateReceiverValue(resultSubstitutor)
 
-        val resolvedCall = ResolvedCallImpl(
-            psiCall, callableCandidate.candidate, dispatchReceiver,
-            extensionReceiver, callableCandidate.explicitReceiverKind,
-            null, temporaryTrace, tracing, MutableDataFlowInfoForArguments.WithoutArgumentsCheck(DataFlowInfo.EMPTY)
-        )
-        resolvedCall.setResultingSubstitutor(resultSubstitutor)
-
-        recordArgumentAdaptationForCallableReference(resolvedCall, callableCandidate.callableReferenceAdaptation)
-
-        tracing.bindCall(topLevelTrace, psiCall)
-        tracing.bindReference(topLevelTrace, resolvedCall)
-        tracing.bindResolvedCall(topLevelTrace, resolvedCall)
-
-        resolvedCall.setStatusToSuccess()
-        resolvedCall.markCallAsCompleted()
-
         when (callableCandidate.candidate) {
             is FunctionDescriptor -> doubleColonExpressionResolver.bindFunctionReference(
                 callableReferenceExpression,
@@ -407,17 +375,90 @@ class ResolvedAtomCompleter(
             )
         }
 
-        // TODO: probably we should also record key 'DATA_FLOW_INFO_BEFORE', see ExpressionTypingVisitorDispatcher.getTypeInfo
-        val typeInfo = createTypeInfo(resultType, resolvedAtom.atom.psiCallArgument.dataFlowInfoAfterThisArgument)
-        topLevelTrace.record(BindingContext.EXPRESSION_TYPE_INFO, callableReferenceExpression, typeInfo)
-        topLevelTrace.record(BindingContext.PROCESSED, callableReferenceExpression)
-
         doubleColonExpressionResolver.checkReferenceIsToAllowedMember(
             callableCandidate.candidate,
             topLevelCallContext.trace,
             callableReferenceExpression
         )
 
+        val explicitCallableReceiver = when (callableCandidate.explicitReceiverKind) {
+            ExplicitReceiverKind.DISPATCH_RECEIVER -> callableCandidate.dispatchReceiver
+            ExplicitReceiverKind.EXTENSION_RECEIVER -> callableCandidate.extensionReceiver
+            else -> null
+        }
+        val explicitReceiver = explicitCallableReceiver?.receiver?.receiverValue?.updateReceiverValue(resultSubstitutor)
+
+        return CallableReferenceResultTypeInfo(dispatchReceiver, extensionReceiver, explicitReceiver, resultSubstitutor, resultType)
+    }
+
+    private fun extractCallableReferenceResultTypeInfoFromDescriptor(
+        callableCandidate: CallableReferenceCandidate,
+        recorderDescriptor: CallableDescriptor
+    ): CallableReferenceResultTypeInfo {
+        val explicitCallableReceiver = when (callableCandidate.explicitReceiverKind) {
+            ExplicitReceiverKind.DISPATCH_RECEIVER -> callableCandidate.dispatchReceiver
+            ExplicitReceiverKind.EXTENSION_RECEIVER -> callableCandidate.extensionReceiver
+            else -> null
+        }
+        return CallableReferenceResultTypeInfo(
+            recorderDescriptor.dispatchReceiverParameter?.value,
+            recorderDescriptor.extensionReceiverParameter?.value,
+            explicitCallableReceiver?.receiver?.receiverValue,
+            TypeSubstitutor.EMPTY,
+            callableCandidate.reflectionCandidateType
+        )
+    }
+
+    private fun completeCallableReference(resolvedAtom: ResolvedCallableReferenceAtom) {
+        val psiCallArgument = resolvedAtom.atom.psiCallArgument as CallableReferenceKotlinCallArgumentImpl
+        val callableReferenceExpression = psiCallArgument.ktCallableReferenceExpression
+        val callableCandidate = resolvedAtom.candidate
+        if (callableCandidate == null || resolvedAtom.completed) {
+            // todo report meanfull diagnostic here
+            return
+        }
+        val recorderDescriptor = when (callableCandidate.candidate) {
+            is FunctionDescriptor -> topLevelCallContext.trace.get(BindingContext.FUNCTION, callableReferenceExpression)
+            is PropertyDescriptor -> topLevelCallContext.trace.get(BindingContext.VARIABLE, callableReferenceExpression)
+            else -> null
+        }
+
+        // For some callable references we can already have recorder descriptor (see `DoubleColonExpressionResolver.getCallableReferenceType`)
+        val resultTypeInfo = if (recorderDescriptor != null) {
+            extractCallableReferenceResultTypeInfoFromDescriptor(callableCandidate, recorderDescriptor)
+        } else {
+            updateCallableReferenceResultType(callableCandidate, psiCallArgument.ktCallableReferenceExpression)
+        }
+
+        val reference = callableReferenceExpression.callableReference
+        val psiCall = CallMaker.makeCall(reference, resultTypeInfo.explicitReceiver, null, reference, emptyList())
+
+        val tracing = TracingStrategyImpl.create(reference, psiCall)
+        val temporaryTrace = TemporaryBindingTrace.create(topLevelTrace, "callable reference fake call")
+
+        val resolvedCall = ResolvedCallImpl(
+            psiCall, callableCandidate.candidate, resultTypeInfo.dispatchReceiver,
+            resultTypeInfo.extensionReceiver, callableCandidate.explicitReceiverKind,
+            null, temporaryTrace, tracing, MutableDataFlowInfoForArguments.WithoutArgumentsCheck(DataFlowInfo.EMPTY)
+        )
+
+        resolvedCall.setResultingSubstitutor(resultTypeInfo.substitutor)
+
+        recordArgumentAdaptationForCallableReference(resolvedCall, callableCandidate.callableReferenceAdaptation)
+
+        tracing.bindCall(topLevelTrace, psiCall)
+        tracing.bindReference(topLevelTrace, resolvedCall)
+        tracing.bindResolvedCall(topLevelTrace, resolvedCall)
+
+        resolvedCall.setStatusToSuccess()
+        resolvedCall.markCallAsCompleted()
+
+        // TODO: probably we should also record key 'DATA_FLOW_INFO_BEFORE', see ExpressionTypingVisitorDispatcher.getTypeInfo
+        val typeInfo = createTypeInfo(resultTypeInfo.resultType, resolvedAtom.atom.psiCallArgument.dataFlowInfoAfterThisArgument)
+
+        topLevelTrace.record(BindingContext.EXPRESSION_TYPE_INFO, callableReferenceExpression, typeInfo)
+        topLevelTrace.record(BindingContext.PROCESSED, callableReferenceExpression)
+
         kotlinToResolvedCallTransformer.runCallCheckers(resolvedCall, topLevelCallCheckerContext)
         resolvedAtom.completed = true
     }

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

@@ -0,0 +1,77 @@
+// 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 poll71(): Flow<String> {
+    return flow {
+        val inv = ::bar2!!
+        inv()
+    }
+}
+
+fun poll73(): Flow<String> {
+    return flow {
+        val inv = ::bar4!!
+        inv
+    }
+}
+
+fun poll75(): Flow<String> {
+    return flow {
+        val inv = ::Foo6!!
+        inv
+    }
+}
+
+fun poll81(): Flow<String> {
+    return flow {
+        val inv = ::bar2 in setOf(::foo2)
+        <!UNRESOLVED_REFERENCE!>inv<!>()
+    }
+}
+
+fun poll83(): Flow<String> {
+    return flow {
+        val inv = ::bar4 in setOf(::foo4)
+        inv
+    }
+}
+
+fun poll85(): Flow<String> {
+    return flow {
+        val inv = ::Foo6 in setOf(::Foo6)
+        inv
+    }
+}

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

@@ -0,0 +1,77 @@
+// 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 poll71(): Flow<String> {
+    return flow {
+        val inv = ::bar2<!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        inv()
+    }
+}
+
+fun poll73(): Flow<String> {
+    return flow {
+        val inv = ::bar4<!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        inv
+    }
+}
+
+fun poll75(): Flow<String> {
+    return flow {
+        val inv = ::Foo6<!NOT_NULL_ASSERTION_ON_CALLABLE_REFERENCE!>!!<!>
+        inv
+    }
+}
+
+fun poll81(): Flow<String> {
+    return flow {
+        val inv = ::bar2 <!TYPE_INFERENCE_ONLY_INPUT_TYPES_WARNING!>in<!> setOf(::foo2)
+        <!UNRESOLVED_REFERENCE_WRONG_RECEIVER!>inv<!>()
+    }
+}
+
+fun poll83(): Flow<String> {
+    return flow {
+        val inv = ::bar4 <!TYPE_INFERENCE_ONLY_INPUT_TYPES_WARNING!>in<!> setOf(::foo4)
+        inv
+    }
+}
+
+fun poll85(): Flow<String> {
+    return flow {
+        val inv = ::Foo6 in setOf(::Foo6)
+        inv
+    }
+}