[NI] Some improvements in callable reference resolution

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/CallableReferenceResolution.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
 import org.jetbrains.kotlin.builtins.isFunctionType
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.resolve.calls.components.CreateDescriptorWithFreshTypeVariables.createToFreshVariableSubstitutorAndAddInitialConstraints
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemOperation
 import org.jetbrains.kotlin.resolve.calls.inference.components.FreshVariableNewTypeSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.model.ArgumentConstraintPosition
@@ -42,6 +43,7 @@ import org.jetbrains.kotlin.types.typeUtil.immediateSupertypes
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.types.upperIfFlexible
 import org.jetbrains.kotlin.utils.SmartList
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 sealed class CallableReceiver(val receiver: ReceiverValueWithSmartCastInfo) {
     class UnboundReference(val qualifier: QualifierReceiver, receiver: ReceiverValueWithSmartCastInfo) : CallableReceiver(receiver)
@@ -108,6 +110,45 @@ fun createCallableReferenceProcessor(factory: CallableReferencesCandidateFactory
     }
 }
 
+fun ConstraintSystemOperation.checkCallableReference(
+        argument: CallableReferenceKotlinCallArgument,
+        dispatchReceiver: CallableReceiver?,
+        extensionReceiver: CallableReceiver?,
+        candidateDescriptor: CallableDescriptor,
+        reflectionCandidateType: UnwrappedType,
+        expectedType: UnwrappedType?,
+        ownerDescriptor: DeclarationDescriptor
+): Pair<FreshVariableNewTypeSubstitutor, KotlinCallDiagnostic?> {
+    val invisibleMember = Visibilities.findInvisibleMember(dispatchReceiver?.asReceiverValueForVisibilityChecks,
+                                                           candidateDescriptor, ownerDescriptor)
+
+    val position = ArgumentConstraintPosition(argument)
+
+    val toFreshSubstitutor = createToFreshVariableSubstitutorAndAddInitialConstraints(candidateDescriptor, this, kotlinCall = null)
+    val reflectionType = toFreshSubstitutor.safeSubstitute(reflectionCandidateType)
+
+    if (expectedType != null) {
+        addSubtypeConstraint(reflectionType, expectedType, position)
+    }
+
+    addReceiverConstraint(toFreshSubstitutor, dispatchReceiver, candidateDescriptor.dispatchReceiverParameter, position)
+    addReceiverConstraint(toFreshSubstitutor, extensionReceiver, candidateDescriptor.extensionReceiverParameter, position)
+
+    return toFreshSubstitutor to invisibleMember?.let(::VisibilityError)
+}
+
+
+private fun ConstraintSystemOperation.addReceiverConstraint(
+        toFreshSubstitutor: FreshVariableNewTypeSubstitutor,
+        receiver: CallableReceiver?,
+        candidateReceiver: ReceiverParameterDescriptor?,
+        position: ArgumentConstraintPosition
+) {
+    val expectedType = toFreshSubstitutor.safeSubstitute(candidateReceiver?.value?.type?.unwrap() ?: return)
+    val receiverType = receiver?.receiver?.stableType ?: return
+    addSubtypeConstraint(receiverType, expectedType, position)
+}
+
 class CallableReferencesCandidateFactory(
         val argument: CallableReferenceKotlinCallArgument,
         val outerCallContext: KotlinCallContext,
@@ -142,40 +183,23 @@ class CallableReferencesCandidateFactory(
 
         val diagnostics = SmartList<KotlinCallDiagnostic>()
         diagnostics.addAll(towerCandidate.diagnostics)
-        val invisibleMember = Visibilities.findInvisibleMember(dispatchCallableReceiver?.asReceiverValueForVisibilityChecks,
-                                                               candidateDescriptor, outerCallContext.scopeTower.lexicalScope.ownerDescriptor)
-        if (invisibleMember != null) {
-            diagnostics.add(VisibilityError(invisibleMember))
-        }
         // todo smartcast on receiver diagnostic and CheckInstantiationOfAbstractClass
 
         compatibilityChecker {
-            if (it.hasContradiction || expectedType == null) return@compatibilityChecker
+            if (it.hasContradiction) return@compatibilityChecker
 
-            val toFreshSubstitutor = CreateDescriptorWithFreshTypeVariables.createToFreshVariableSubstitutorAndAddInitialConstraints(candidateDescriptor, it, kotlinCall = null)
+            val (_, visibilityError) = it.checkCallableReference(argument, dispatchCallableReceiver, extensionCallableReceiver, candidateDescriptor,
-            val reflectionType = toFreshSubstitutor.safeSubstitute(reflectionCandidateType)
+                                      reflectionCandidateType, expectedType, outerCallContext.scopeTower.lexicalScope.ownerDescriptor)
-            it.addSubtypeConstraint(reflectionType, expectedType, position)
 
-            it.addUnboundConstraint(toFreshSubstitutor, dispatchCallableReceiver, candidateDescriptor.dispatchReceiverParameter)
+            diagnostics.addIfNotNull(visibilityError)
-            it.addUnboundConstraint(toFreshSubstitutor, extensionCallableReceiver, candidateDescriptor.extensionReceiverParameter)
 
-            if (it.hasContradiction) diagnostics.add(CallableReferenceNotCompatible(argument, candidateDescriptor, expectedType, reflectionType))
+            if (it.hasContradiction) diagnostics.add(CallableReferenceNotCompatible(argument, candidateDescriptor, expectedType, reflectionCandidateType))
         }
 
         return CallableReferenceCandidate(candidateDescriptor, dispatchCallableReceiver, extensionCallableReceiver,
                                           explicitReceiverKind, reflectionCandidateType, defaults, ResolutionCandidateStatus(diagnostics))
     }
 
-    private fun ConstraintSystemOperation.addUnboundConstraint(
-            toFreshSubstitutor: FreshVariableNewTypeSubstitutor,
-            receiver: CallableReceiver?,
-            candidateReceiver: ReceiverParameterDescriptor?
-    ) {
-        val expectedType = toFreshSubstitutor.safeSubstitute(candidateReceiver?.value?.type?.unwrap() ?: return)
-        val receiverType = receiver?.receiver?.stableType ?: return
-        addSubtypeConstraint(receiverType, expectedType, position)
-    }
-
     private fun getArgumentAndReturnTypeUseMappingByExpectedType(
             descriptor: FunctionDescriptor,
             expectedType: UnwrappedType?,
@@ -235,12 +259,10 @@ class CallableReferencesCandidateFactory(
         val argumentsAndReceivers = ArrayList<KotlinType>(descriptor.valueParameters.size + 2)
 
         if (dispatchReceiver is CallableReceiver.UnboundReference) {
-            argumentsAndReceivers.add(descriptor.dispatchReceiverParameter?.type
+            argumentsAndReceivers.add(dispatchReceiver.receiver.stableType)
-                                      ?: error("Dispatch receiver should be not null for descriptor: $descriptor"))
         }
         if (extensionReceiver is CallableReceiver.UnboundReference) {
-            argumentsAndReceivers.add(descriptor.extensionReceiverParameter?.type
+            argumentsAndReceivers.add(extensionReceiver.receiver.stableType)
-                                      ?: error("Extension receiver should be not null for descriptor: $descriptor"))
         }
 
         val descriptorReturnType = descriptor.returnType
@@ -305,7 +327,7 @@ class CallableReferencesCandidateFactory(
 class CallableReferenceNotCompatible(
         val argument: CallableReferenceKotlinCallArgument,
         val candidate: CallableMemberDescriptor,
-        val expectedType: UnwrappedType,
+        val expectedType: UnwrappedType?,
         val callableReverenceType: UnwrappedType
 ) : KotlinCallDiagnostic(ResolutionCandidateApplicability.INAPPLICABLE) {
     override fun report(reporter: DiagnosticReporter) = reporter.onCallArgument(argument, this)

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/CheckArgumentsResolutionPart.kt

@@ -20,7 +20,6 @@ import org.jetbrains.kotlin.builtins.*
 import org.jetbrains.kotlin.descriptors.ClassifierDescriptorWithTypeParameters
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
-import org.jetbrains.kotlin.resolve.calls.components.CreateDescriptorWithFreshTypeVariables.createToFreshVariableSubstitutorAndAddInitialConstraints
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.addSubtypeConstraintIfCompatible
 import org.jetbrains.kotlin.resolve.calls.inference.model.ArgumentConstraintPosition
@@ -173,16 +172,18 @@ internal object CheckArguments : ResolutionPart {
             1 -> candidates.single()
             else -> return CallableReferenceCandidatesAmbiguity(argument, candidates)
         }
+        val (toFreshSubstitutor, diagnostic) = with(chosenCandidate) {
+            csBuilder.checkCallableReference(argument, dispatchReceiver, extensionReceiver, candidate,
+                                             reflectionCandidateType, expectedType, callContext.scopeTower.lexicalScope.ownerDescriptor)
+        }
 
+        val resolvedCallableReference = ResolvedCallableReferenceArgument(
+                kotlinCall, argument, toFreshSubstitutor.freshVariables,
+                chosenCandidate, toFreshSubstitutor.safeSubstitute(chosenCandidate.reflectionCandidateType))
 
-        val toFreshSubstitutor = createToFreshVariableSubstitutorAndAddInitialConstraints(chosenCandidate.candidate, csBuilder, kotlinCall = null)
-        val reflectionType = toFreshSubstitutor.safeSubstitute(chosenCandidate.reflectionCandidateType)
-        csBuilder.addSubtypeConstraint(reflectionType, expectedType, ArgumentConstraintPosition(argument))
-
-        val resolvedCallableReference = ResolvedCallableReferenceArgument(kotlinCall, argument, toFreshSubstitutor.freshVariables, chosenCandidate)
         csBuilder.addCallableReferenceArgument(resolvedCallableReference)
 
-        return null
+        return diagnostic
     }
 
     fun processCollectionLiteralArgument(

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/FixationOrderCalculator.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.resolve.calls.inference.components
 import org.jetbrains.kotlin.resolve.calls.inference.model.Constraint
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind
 import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
+import org.jetbrains.kotlin.resolve.calls.model.ArgumentWithPostponeResolution
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedLambdaArgument
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.checker.NewKotlinTypeChecker
@@ -42,7 +43,7 @@ class FixationOrderCalculator {
 
     interface Context {
         val notFixedTypeVariables: Map<TypeConstructor, VariableWithConstraints>
-        val lambdaArguments: List<ResolvedLambdaArgument>
+        val lambdaAndCallableReferenceArguments: List<ArgumentWithPostponeResolution>
     }
 
     fun computeCompletionOrder(
@@ -87,11 +88,11 @@ class FixationOrderCalculator {
         }
 
         private fun buildLambdaEdges() {
-            for (lambdaArgument in c.lambdaArguments) {
+            for (lambdaArgument in c.lambdaAndCallableReferenceArguments) {
-                if (lambdaArgument.analyzed) continue
+                if (lambdaArgument is ResolvedLambdaArgument && lambdaArgument.analyzed) continue // optimization
 
                 val typeVariablesInReturnType = SmartList<Variable>()
-                lambdaArgument.returnType.findTypeVariables(typeVariablesInReturnType)
+                lambdaArgument.outputType?.findTypeVariables(typeVariablesInReturnType)
 
                 if (typeVariablesInReturnType.isEmpty()) continue // optimization
 
@@ -131,7 +132,7 @@ class FixationOrderCalculator {
             topReturnType.visitType(ResolveDirection.TO_SUBTYPE) { variableWithConstraints, direction ->
                 enterToNode(variableWithConstraints, direction)
             }
-            for (resolvedLambdaArgument in c.lambdaArguments) {
+            for (resolvedLambdaArgument in c.lambdaAndCallableReferenceArguments) {
                 inner@ for (inputType in resolvedLambdaArgument.inputTypes) {
                     inputType.visitType(ResolveDirection.TO_SUBTYPE) { variableWithConstraints, direction ->
                         enterToNode(variableWithConstraints, direction)

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/model/NewConstraintSystemImpl.kt

@@ -249,12 +249,18 @@ class NewConstraintSystemImpl(val constraintInjector: ConstraintInjector, val re
         storage.errors.add(error)
     }
 
-    // FixationOrderCalculator.Context, KotlinCallCompleter.Context
+    //  KotlinCallCompleter.Context
     override val lambdaArguments: List<ResolvedLambdaArgument> get() {
         checkState(State.COMPLETION)
         return storage.lambdaArguments
     }
 
+    // FixationOrderCalculator.Context
+    override val lambdaAndCallableReferenceArguments: List<ArgumentWithPostponeResolution> get() {
+        checkState(State.COMPLETION)
+        return storage.lambdaArguments + storage.callableReferenceArguments
+    }
+
     override val callableReferenceArguments: List<ResolvedCallableReferenceArgument>
         get() {
             checkState(State.COMPLETION)

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolvedCallElements.kt

@@ -30,8 +30,6 @@ sealed class ArgumentWithPostponeResolution {
     abstract val argument: KotlinCallArgument
     abstract val inputTypes: Collection<UnwrappedType> // parameters and implicit receiver
     abstract val outputType: UnwrappedType?
-
-    var analyzed: Boolean = false
 }
 
 class ResolvedLambdaArgument(
@@ -42,6 +40,8 @@ class ResolvedLambdaArgument(
         val parameters: List<UnwrappedType>,
         val returnType: UnwrappedType
 ) : ArgumentWithPostponeResolution() {
+    var analyzed: Boolean = false
+
     val type: SimpleType = createFunctionType(returnType.builtIns, Annotations.EMPTY, receiver, parameters, null, returnType, isSuspend) // todo support annotations
 
     override val inputTypes: Collection<UnwrappedType> get() = receiver?.let { parameters + it } ?: parameters
@@ -55,10 +55,11 @@ class ResolvedCallableReferenceArgument(
         override val outerCall: KotlinCall,
         override val argument: CallableReferenceKotlinCallArgument,
         val myTypeVariables: List<NewTypeVariable>,
-        val callableResolutionCandidate: CallableReferenceCandidate
+        val callableResolutionCandidate: CallableReferenceCandidate,
+        val reflectionType: UnwrappedType // via myTypeVariables variables
 ) : ArgumentWithPostponeResolution() {
-    override val inputTypes: Collection<UnwrappedType> get() = emptyList()
+    override val inputTypes: Collection<UnwrappedType> get() = reflectionType.arguments.dropLast(1).map { it.type.unwrap() }
-    override val outputType: UnwrappedType? = null
+    override val outputType: UnwrappedType? = reflectionType.arguments.last().type.unwrap()
 }
 
 class ResolvedCollectionLiteralArgument(