Complete old inference's callable references properly, by updating descriptor, resolved call and recorded expression's type

^KT-45721 Fixed
^KT-44994 Fixed

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -17071,6 +17071,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
             }
 
+            @Test
+            @TestMetadata("callableReferencesProperCompletion.kt")
+            public void testCallableReferencesProperCompletion() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/callableReferencesProperCompletion.kt");
+            }
+
             @Test
             @TestMetadata("kt41164.kt")
             public void testKt41164() throws Exception {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/inference/BuilderInferenceSession.kt

@@ -15,9 +15,9 @@ import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
+import org.jetbrains.kotlin.resolve.calls.callUtil.toOldSubstitution
 import org.jetbrains.kotlin.resolve.calls.components.*
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
-import org.jetbrains.kotlin.resolve.calls.inference.BuilderInferenceSession.Companion.updateCalls
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
 import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
@@ -68,7 +68,7 @@ class BuilderInferenceSession(
     private val commonCalls = arrayListOf<PSICompletedCallInfo>()
 
     // These calls come from the old type inference
-    private val oldCallableReferenceCalls = arrayListOf<KtExpression>()
+    private val oldDoubleColonExpressionCalls = arrayListOf<KtExpression>()
 
     private var hasInapplicableCall = false
 
@@ -117,7 +117,7 @@ class BuilderInferenceSession(
     }
 
     fun addOldCallableReferenceCalls(callExpression: KtExpression) {
-        oldCallableReferenceCalls.add(callExpression)
+        oldDoubleColonExpressionCalls.add(callExpression)
     }
 
     override fun addCompletedCallInfo(callInfo: CompletedCallInfo) {
@@ -378,12 +378,32 @@ class BuilderInferenceSession(
         return currentSession.topLevelCallContext.trace
     }
 
-    private fun completeCallableReference(expression: KtCallableReferenceExpression, substitutor: NewTypeSubstitutor) {
-        createResolvedAtomCompleter(substitutor, topLevelCallContext).substituteFunctionLiteralDescriptor(
-            resolvedAtom = null,
-            descriptor = trace.get(BindingContext.DECLARATION_TO_DESCRIPTOR, expression) as? SimpleFunctionDescriptorImpl ?: return,
-            substitutor
-        )
+    private fun completeDoubleColonExpression(expression: KtDoubleColonExpression, substitutor: NewTypeSubstitutor) {
+        val atomCompleter = createResolvedAtomCompleter(substitutor, topLevelCallContext)
+        val declarationDescriptor = trace.get(BindingContext.DECLARATION_TO_DESCRIPTOR, expression)
+
+        if (declarationDescriptor is SimpleFunctionDescriptorImpl) {
+            atomCompleter.substituteFunctionLiteralDescriptor(resolvedAtom = null, descriptor = declarationDescriptor, substitutor)
+        }
+
+        val recordedType = trace.getType(expression)
+
+        if (recordedType != null) {
+            trace.recordType(expression, substitutor.safeSubstitute(recordedType.unwrap()))
+        }
+
+        val targetExpression = when (expression) {
+            is KtCallableReferenceExpression -> expression.callableReference
+            is KtClassLiteralExpression -> expression.receiverExpression
+            else -> throw IllegalStateException("Unsupported double colon expression")
+        }
+
+        val call = trace.get(BindingContext.CALL, targetExpression) ?: return
+        val resolvedCall = trace.get(BindingContext.RESOLVED_CALL, call)
+
+        if (resolvedCall is ResolvedCallImpl<*>) {
+            resolvedCall.setResultingSubstitutor(substitutor.toOldSubstitution().buildSubstitutor())
+        }
     }
 
     private fun completeCall(
@@ -453,9 +473,9 @@ class BuilderInferenceSession(
                 reportErrors(callInfo, resolvedCall, errors)
             }
 
-            for (simpleCall in oldCallableReferenceCalls) {
-                when (simpleCall) {
-                    is KtCallableReferenceExpression -> completeCallableReference(simpleCall, nonFixedTypesToResultSubstitutor)
+            for (call in oldDoubleColonExpressionCalls) {
+                when (call) {
+                    is KtDoubleColonExpression -> completeDoubleColonExpression(call, nonFixedTypesToResultSubstitutor)
                     else -> throw Exception("Unsupported call expression type")
                 }
             }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/model/ResolvedCallImpl.java

@@ -35,12 +35,11 @@ import org.jetbrains.kotlin.resolve.calls.results.ResolutionStatus;
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind;
 import org.jetbrains.kotlin.resolve.calls.tasks.ResolutionCandidate;
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy;
-import org.jetbrains.kotlin.resolve.scopes.receivers.CastImplicitClassReceiver;
-import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitClassReceiver;
-import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue;
+import org.jetbrains.kotlin.resolve.scopes.receivers.*;
 import org.jetbrains.kotlin.types.KotlinType;
 import org.jetbrains.kotlin.types.TypeProjection;
 import org.jetbrains.kotlin.types.TypeSubstitutor;
+import org.jetbrains.kotlin.types.Variance;
 
 import java.util.*;
 
@@ -63,7 +62,7 @@ public class ResolvedCallImpl<D extends CallableDescriptor> implements MutableRe
     private final Call call;
     private final D candidateDescriptor;
     private D resultingDescriptor; // Probably substituted
-    private final ReceiverValue dispatchReceiver; // receiver object of a method
+    private ReceiverValue dispatchReceiver; // receiver object of a method
     private ReceiverValue extensionReceiver; // receiver of an extension function
     private final ExplicitReceiverKind explicitReceiverKind;
     private final TypeSubstitutor knownTypeParametersSubstitutor;
@@ -240,6 +239,14 @@ public class ResolvedCallImpl<D extends CallableDescriptor> implements MutableRe
             assert substitutedVersion != null : valueParameterDescriptor;
             argumentToParameterMap.put(entry.getKey(), argumentMatch.replaceValueParameter(substitutedVersion));
         }
+
+        if (dispatchReceiver instanceof ExpressionReceiver) {
+            dispatchReceiver = dispatchReceiver.replaceType(substitutor.safeSubstitute(dispatchReceiver.getType(), Variance.IN_VARIANCE));
+        }
+        if (extensionReceiver instanceof ExtensionReceiver) {
+            extensionReceiver =
+                    extensionReceiver.replaceType(substitutor.safeSubstitute(extensionReceiver.getType(), Variance.IN_VARIANCE));
+        }
     }
 
     @Override

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.resolve.MissingSupertypesResolver
 import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.NewCommonSuperTypeCalculator
+import org.jetbrains.kotlin.resolve.calls.callUtil.toOldSubstitution
 import org.jetbrains.kotlin.resolve.calls.checkers.CallCheckerContext
 import org.jetbrains.kotlin.resolve.calls.commonSuperType
 import org.jetbrains.kotlin.resolve.calls.components.CallableReferenceAdaptation
@@ -45,7 +46,6 @@ import org.jetbrains.kotlin.types.expressions.CoercionStrategy
 import org.jetbrains.kotlin.types.expressions.DoubleColonExpressionResolver
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.createTypeInfo
-import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.contains
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.types.typeUtil.shouldBeUpdated
@@ -230,7 +230,8 @@ class ResolvedAtomCompleter(
     ): FunctionLiteralTypes {
         val returnType =
             (if (resolvedAtom?.isCoercedToUnit == true) builtIns.unitType else resolvedAtom?.returnType) ?: descriptor.returnType
-        val receiverType = resolvedAtom?.receiver ?: descriptor.extensionReceiverParameter?.type
+        val extensionReceiverType = resolvedAtom?.receiver ?: descriptor.extensionReceiverParameter?.type
+        val dispatchReceiverType = descriptor.dispatchReceiverParameter?.type
         val valueParameterTypes = resolvedAtom?.parameters ?: descriptor.valueParameters.map { it.type }
 
         require(returnType != null)
@@ -239,10 +240,17 @@ class ResolvedAtomCompleter(
             descriptor.setReturnType(it.approximatedType)
         }
 
-        val receiverFromDescriptor = descriptor.extensionReceiverParameter
-        val substitutedReceiverType = receiverType?.substituteAndApproximate(substitutor)?.also {
-            if (receiverFromDescriptor is ReceiverParameterDescriptorImpl && receiverFromDescriptor.type.shouldBeUpdated()) {
-                receiverFromDescriptor.setOutType(it.approximatedType)
+        val extensionReceiverFromDescriptor = descriptor.extensionReceiverParameter
+        val substitutedReceiverType = extensionReceiverType?.substituteAndApproximate(substitutor)?.also {
+            if (extensionReceiverFromDescriptor is ReceiverParameterDescriptorImpl && extensionReceiverFromDescriptor.type.shouldBeUpdated()) {
+                extensionReceiverFromDescriptor.setOutType(it.approximatedType)
+            }
+        }
+
+        val dispatchReceiverFromDescriptor = descriptor.dispatchReceiverParameter
+        dispatchReceiverType?.substituteAndApproximate(substitutor)?.also {
+            if (dispatchReceiverFromDescriptor is ReceiverParameterDescriptorImpl && dispatchReceiverFromDescriptor.type.shouldBeUpdated()) {
+                dispatchReceiverFromDescriptor.setOutType(it.approximatedType)
             }
         }
 
@@ -309,16 +317,6 @@ class ResolvedAtomCompleter(
         }
     }
 
-    private fun NewTypeSubstitutor.toOldSubstitution(): TypeSubstitution = object : TypeSubstitution() {
-        override fun get(key: KotlinType): TypeProjection? {
-            return safeSubstitute(key.unwrap()).takeIf { it !== key }?.asTypeProjection()
-        }
-
-        override fun isEmpty(): Boolean {
-            return isEmpty
-        }
-    }
-
     private fun updateCallableReferenceResultType(
         callableCandidate: CallableReferenceCandidate,
         callableReferenceExpression: KtCallableReferenceExpression

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/util/callUtil.kt

@@ -33,11 +33,15 @@ import org.jetbrains.kotlin.resolve.StatementFilter
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.CallTransformer
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
+import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
 import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.resolve.calls.tower.NewResolvedCallImpl
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.TypeProjection
+import org.jetbrains.kotlin.types.TypeSubstitution
 import org.jetbrains.kotlin.types.isError
+import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.utils.sure
 
@@ -321,3 +325,13 @@ inline fun BindingTrace.reportTrailingLambdaErrorOr(
         }
     }
 }
+
+fun NewTypeSubstitutor.toOldSubstitution(): TypeSubstitution = object : TypeSubstitution() {
+    override fun get(key: KotlinType): TypeProjection? {
+        return safeSubstitute(key.unwrap()).takeIf { it !== key }?.asTypeProjection()
+    }
+
+    override fun isEmpty(): Boolean {
+        return isEmpty
+    }
+}

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/DoubleColonExpressionResolver.kt

@@ -113,6 +113,11 @@ class DoubleColonExpressionResolver(
             c.trace.report(UNSUPPORTED.on(expression, "Class literals with empty left hand side are not yet supported"))
         } else {
             val result = resolveDoubleColonLHS(expression, c)
+
+            if (c.inferenceSession is BuilderInferenceSession && result?.type?.contains { it is StubType } == true) {
+                c.inferenceSession.addOldCallableReferenceCalls(expression)
+            }
+
             if (result != null && !result.type.isError) {
                 val inherentType = result.type
                 val dataFlowInfo = (result as? DoubleColonLHS.Expression)?.dataFlowInfo ?: c.dataFlowInfo

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

@@ -0,0 +1,72 @@
+// DONT_TARGET_EXACT_BACKEND: WASM
+// WITH_RUNTIME
+
+fun nonGenericId(x: Any?) = x
+
+fun <T> id(x: T) = x
+
+fun test1(): Sequence<String> = sequence {
+    yield("")
+    this::class
+}
+
+fun <T> sequence2(block: suspend SequenceScope<T>.() -> Unit): Sequence<T> = Sequence { iterator(block) }
+
+fun foo() {
+    sequence2<String> {
+        id(this::class)
+    }
+}
+
+fun test2(): Sequence<String> = sequence {
+    yield("")
+//    id(this::class) // TODO
+}
+
+fun test3(): Sequence<String> = sequence {
+    yield("")
+    nonGenericId(this::class)
+}
+
+fun test4(): Sequence<String> = sequence {
+    yield("")
+    this::`yield`
+}
+
+fun test5(): Sequence<String> = sequence {
+    yield("")
+    id(this::`yield`)
+}
+
+fun test6(): Sequence<String> = sequence {
+    yield("")
+    nonGenericId(this::`yield`)
+}
+
+fun test7(): Sequence<String> = sequence {
+    yield("")
+    ::`yield`
+}
+
+fun test8(): Sequence<String> = sequence {
+    yield("")
+    id(::`yield`)
+}
+
+fun test9(): Sequence<String> = sequence {
+    yield("")
+    nonGenericId(::`yield`)
+}
+
+fun box(): String {
+    test1()
+    test2()
+    test3()
+    test4()
+    test5()
+    test6()
+    test7()
+    test8()
+    test9()
+    return "OK"
+}

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

@@ -1,8 +1,5 @@
 // DONT_TARGET_EXACT_BACKEND: WASM
 // WASM_MUTE_REASON: COROUTINES
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND: JS_IR
-// IGNORE_BACKEND: NATIVE
 // KJS_WITH_FULL_RUNTIME
 // WITH_RUNTIME
 // WITH_COROUTINES

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -17053,6 +17053,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
             }
 
+            @Test
+            @TestMetadata("callableReferencesProperCompletion.kt")
+            public void testCallableReferencesProperCompletion() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/callableReferencesProperCompletion.kt");
+            }
+
             @Test
             @TestMetadata("kt41164.kt")
             public void testKt41164() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -17071,6 +17071,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
             }
 
+            @Test
+            @TestMetadata("callableReferencesProperCompletion.kt")
+            public void testCallableReferencesProperCompletion() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/callableReferencesProperCompletion.kt");
+            }
+
             @Test
             @TestMetadata("kt41164.kt")
             public void testKt41164() throws Exception {

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

@@ -14122,6 +14122,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
             }
 
+            @TestMetadata("callableReferencesProperCompletion.kt")
+            public void testCallableReferencesProperCompletion() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/callableReferencesProperCompletion.kt");
+            }
+
             @TestMetadata("kt41164.kt")
             public void testKt41164() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt41164.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/es6/semantics/IrJsCodegenBoxES6TestGenerated.java

@@ -12406,6 +12406,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
             }
 
+            @TestMetadata("callableReferencesProperCompletion.kt")
+            public void testCallableReferencesProperCompletion() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/callableReferencesProperCompletion.kt");
+            }
+
             @TestMetadata("kt41164.kt")
             public void testKt41164() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt41164.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -11827,6 +11827,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
             }
 
+            @TestMetadata("callableReferencesProperCompletion.kt")
+            public void testCallableReferencesProperCompletion() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/callableReferencesProperCompletion.kt");
+            }
+
             @TestMetadata("kt41164.kt")
             public void testKt41164() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt41164.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -11892,6 +11892,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/callableReferenceAndCoercionToUnit.kt");
             }
 
+            @TestMetadata("callableReferencesProperCompletion.kt")
+            public void testCallableReferencesProperCompletion() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/callableReferencesProperCompletion.kt");
+            }
+
             @TestMetadata("kt41164.kt")
             public void testKt41164() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt41164.kt");