NI: do substitution type variables during updating trace for lambda (these type variables can appear after the builder inference)

^KT-42450 Fixed

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

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

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

@@ -47,6 +47,7 @@ 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.isUnit
+import org.jetbrains.kotlin.types.typeUtil.shouldBeSubstituted
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class ResolvedAtomCompleter(
@@ -212,15 +213,14 @@ class ResolvedAtomCompleter(
             resultSubstitutor.safeSubstitute(lambda.returnType)
         }
 
-        val approximatedValueParameterTypes = lambda.parameters.map {
-            // Do substitution and approximation only for stub types, which can appear from builder inference (as postponed variables)
-            if (it is StubType) {
+        val approximatedValueParameterTypes = lambda.parameters.map { parameterType ->
+            if (parameterType.shouldBeSubstituted()) {
                 typeApproximator.approximateDeclarationType(
-                    resultSubstitutor.safeSubstitute(it),
+                    resultSubstitutor.safeSubstitute(parameterType),
                     local = true,
                     languageVersionSettings = topLevelCallContext.languageVersionSettings
                 )
-            } else it
+            } else parameterType
         }
 
         val approximatedReturnType =
@@ -277,8 +277,7 @@ class ResolvedAtomCompleter(
         functionDescriptor.setReturnType(returnType)
 
         for ((i, valueParameter) in functionDescriptor.valueParameters.withIndex()) {
-            if (valueParameter !is ValueParameterDescriptorImpl || valueParameter.type !is StubType)
-                continue
+            if (valueParameter !is ValueParameterDescriptorImpl || !valueParameter.type.shouldBeSubstituted()) continue
             valueParameter.setOutType(valueParameters[i])
         }
 

compiler/testData/codegen/box/builderInference/substituteTypeVariableIntolambdaParameterDescriptor.kt

@@ -0,0 +1,25 @@
+// TARGET_BACKEND: JVM
+// IGNORE_BACKEND_FIR: JVM_IR
+// WITH_RUNTIME
+
+class Foo<A>
+
+fun <K> bar(x: Foo<K>): Unit {}
+
+fun <E> foo(block: (Foo<E>) -> Unit): E = null as E
+
+interface FlowCollector<T> {
+    fun emit(value: T)
+}
+
+@Suppress("EXPERIMENTAL_API_USAGE_ERROR")
+fun <I> flow(@BuilderInference block: FlowCollector<I>.() -> Unit): I = null as I
+
+fun adapt(): Unit = flow {
+    emit(foo { coroutine -> bar(coroutine) })
+}
+
+fun box(): String {
+    adapt()
+    return "OK"
+}

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

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

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

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

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

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

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

@@ -22,7 +22,7 @@ import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.types.KotlinType;
-import org.jetbrains.kotlin.types.StubType;
+import org.jetbrains.kotlin.types.typeUtil.TypeUtilsKt;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -50,7 +50,7 @@ public abstract class VariableDescriptorImpl extends DeclarationDescriptorNonRoo
     }
 
     public void setOutType(KotlinType outType) {
-        assert this.outType == null || this.outType instanceof StubType;
+        assert this.outType == null || TypeUtilsKt.shouldBeSubstituted(this.outType);
         this.outType = outType;
     }
 

core/descriptors/src/org/jetbrains/kotlin/types/TypeUtils.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.isCaptured
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.checker.*
+import org.jetbrains.kotlin.types.model.TypeVariableTypeConstructorMarker
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import java.util.*
 
@@ -304,3 +305,5 @@ private fun NewCapturedType.unCaptureTopLevelType(): UnwrappedType {
 
     return constructor.projection.type.unwrap()
 }
+
+fun KotlinType.shouldBeSubstituted() = contains { it is StubType || it.constructor is TypeVariableTypeConstructorMarker }