Substitute lambda's receiver type during completion including the builder inference stub variables substitution

^KT-42175 Fixed

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

@@ -13507,6 +13507,11 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             public void testSubstituteTypeVariableIntolambdaParameterDescriptor() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteTypeVariableIntolambdaParameterDescriptor.kt");
             }
+
+            @TestMetadata("substitutelambdaExtensionReceiverType.kt")
+            public void testSubstitutelambdaExtensionReceiverType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
+            }
         }
     }
 

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

@@ -221,6 +221,7 @@ class ResolvedAtomCompleter(
         } else {
             resultSubstitutor.safeSubstitute(lambda.returnType)
         }
+        val receiverType = lambda.receiver
 
         val approximatedValueParameterTypes = lambda.parameters.map { parameterType ->
             if (parameterType.shouldBeSubstituted()) {
@@ -238,7 +239,16 @@ class ResolvedAtomCompleter(
                 local = true,
                 languageVersionSettings = topLevelCallContext.languageVersionSettings
             )
-        updateTraceForLambda(lambda, topLevelTrace, approximatedReturnType, approximatedValueParameterTypes)
+
+        val approximatedReceiverType = if (receiverType != null) {
+            typeApproximator.approximateDeclarationType(
+                resultSubstitutor.safeSubstitute(receiverType),
+                local = true,
+                languageVersionSettings = topLevelCallContext.languageVersionSettings
+            )
+        } else null
+
+        updateTraceForLambda(lambda, topLevelTrace, approximatedReturnType, approximatedValueParameterTypes, approximatedReceiverType)
 
         for (lambdaResult in resultArgumentsInfo.nonErrorArguments) {
             val resultValueArgument = lambdaResult as? PSIKotlinCallArgument ?: continue
@@ -262,7 +272,8 @@ class ResolvedAtomCompleter(
         lambda: ResolvedLambdaAtom,
         trace: BindingTrace,
         returnType: UnwrappedType,
-        valueParameters: List<UnwrappedType>
+        valueParameters: List<UnwrappedType>,
+        receiverType: UnwrappedType?
     ) {
         val psiCallArgument = lambda.atom.psiCallArgument
 
@@ -285,6 +296,12 @@ class ResolvedAtomCompleter(
 
         functionDescriptor.setReturnType(returnType)
 
+        val extensionReceiverParameter = functionDescriptor.extensionReceiverParameter
+
+        if (receiverType != null && extensionReceiverParameter is ReceiverParameterDescriptorImpl && extensionReceiverParameter.type.shouldBeSubstituted()) {
+            extensionReceiverParameter.setOutType(receiverType)
+        }
+
         for ((i, valueParameter) in functionDescriptor.valueParameters.withIndex()) {
             if (valueParameter !is ValueParameterDescriptorImpl || !valueParameter.type.shouldBeSubstituted()) continue
             valueParameter.setOutType(valueParameters[i])

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

@@ -0,0 +1,22 @@
+// WITH_RUNTIME
+// DONT_TARGET_EXACT_BACKEND: WASM
+
+import kotlin.experimental.ExperimentalTypeInference
+
+operator fun <T> SequenceScope<String>.plusAssign(x: SequenceScope<T>) {}
+
+@OptIn(ExperimentalTypeInference::class)
+
+fun <T> mySequence(@BuilderInference block: suspend SequenceScope<T>.() -> Unit): Sequence<T> = Sequence { iterator(block) }
+
+fun main() {
+    val y: Sequence<String> = mySequence {
+        yield("result")
+        this += this
+    }
+}
+
+fun box(): String {
+    main()
+    return "OK"
+}

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

@@ -13507,6 +13507,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             public void testSubstituteTypeVariableIntolambdaParameterDescriptor() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteTypeVariableIntolambdaParameterDescriptor.kt");
             }
+
+            @TestMetadata("substitutelambdaExtensionReceiverType.kt")
+            public void testSubstitutelambdaExtensionReceiverType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
+            }
         }
     }
 

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

@@ -13507,6 +13507,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             public void testSubstituteTypeVariableIntolambdaParameterDescriptor() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteTypeVariableIntolambdaParameterDescriptor.kt");
             }
+
+            @TestMetadata("substitutelambdaExtensionReceiverType.kt")
+            public void testSubstitutelambdaExtensionReceiverType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
+            }
         }
     }
 

compiler/tests-gen/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -13507,6 +13507,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             public void testSubstituteTypeVariableIntolambdaParameterDescriptor() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteTypeVariableIntolambdaParameterDescriptor.kt");
             }
+
+            @TestMetadata("substitutelambdaExtensionReceiverType.kt")
+            public void testSubstitutelambdaExtensionReceiverType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
+            }
         }
     }
 

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/ReceiverParameterDescriptorImpl.java

@@ -21,10 +21,12 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor;
 import org.jetbrains.kotlin.descriptors.ReceiverParameterDescriptor;
 import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue;
+import org.jetbrains.kotlin.types.KotlinType;
+import org.jetbrains.kotlin.types.typeUtil.TypeUtilsKt;
 
 public class ReceiverParameterDescriptorImpl extends AbstractReceiverParameterDescriptor {
     private final DeclarationDescriptor containingDeclaration;
-    private final ReceiverValue value;
+    private ReceiverValue value;
 
     public ReceiverParameterDescriptorImpl(
             @NotNull DeclarationDescriptor containingDeclaration,
@@ -53,4 +55,9 @@ public class ReceiverParameterDescriptorImpl extends AbstractReceiverParameterDe
     public ReceiverParameterDescriptor copy(@NotNull DeclarationDescriptor newOwner) {
         return new ReceiverParameterDescriptorImpl(newOwner, value, getAnnotations());
     }
+
+    public void setOutType(@NotNull KotlinType outType) {
+        assert TypeUtilsKt.shouldBeSubstituted(this.value.getType());
+        this.value = value.replaceType(outType);
+    }
 }

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

@@ -11557,6 +11557,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
             public void testSubstituteStubTypeIntolambdaParameterDescriptor() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteStubTypeIntolambdaParameterDescriptor.kt");
             }
+
+            @TestMetadata("substitutelambdaExtensionReceiverType.kt")
+            public void testSubstitutelambdaExtensionReceiverType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
+            }
         }
     }
 

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

@@ -11557,6 +11557,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             public void testSubstituteStubTypeIntolambdaParameterDescriptor() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteStubTypeIntolambdaParameterDescriptor.kt");
             }
+
+            @TestMetadata("substitutelambdaExtensionReceiverType.kt")
+            public void testSubstitutelambdaExtensionReceiverType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
+            }
         }
     }
 

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

@@ -11622,6 +11622,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             public void testSubstituteStubTypeIntolambdaParameterDescriptor() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/substituteStubTypeIntolambdaParameterDescriptor.kt");
             }
+
+            @TestMetadata("substitutelambdaExtensionReceiverType.kt")
+            public void testSubstitutelambdaExtensionReceiverType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/substitutelambdaExtensionReceiverType.kt");
+            }
         }
     }