NI: do substitution stub types during updating trace for lambda (these stub types can appear from the builder inference as postponed variables)

^KT-40060 Fixed

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.builtins.createFunctionType
 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.KtExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
@@ -209,13 +210,24 @@ 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) {
+                typeApproximator.approximateDeclarationType(
+                    resultSubstitutor.safeSubstitute(it),
+                    local = true,
+                    languageVersionSettings = topLevelCallContext.languageVersionSettings
+                )
+            } else it
+        }
+
         val approximatedReturnType =
             typeApproximator.approximateDeclarationType(
                 returnType,
                 local = true,
                 languageVersionSettings = topLevelCallContext.languageVersionSettings
             )
-        updateTraceForLambda(lambda, topLevelTrace, approximatedReturnType)
+        updateTraceForLambda(lambda, topLevelTrace, approximatedReturnType, approximatedValueParameterTypes)
 
         for (lambdaResult in resultArgumentsInfo.nonErrorArguments) {
             val resultValueArgument = lambdaResult as? PSIKotlinCallArgument ?: continue
@@ -235,7 +247,12 @@ class ResolvedAtomCompleter(
         }
     }
 
-    private fun updateTraceForLambda(lambda: ResolvedLambdaAtom, trace: BindingTrace, returnType: UnwrappedType) {
+    private fun updateTraceForLambda(
+        lambda: ResolvedLambdaAtom,
+        trace: BindingTrace,
+        returnType: UnwrappedType,
+        valueParameters: List<UnwrappedType>
+    ) {
         val psiCallArgument = lambda.atom.psiCallArgument
 
         val ktArgumentExpression: KtExpression
@@ -254,8 +271,15 @@ class ResolvedAtomCompleter(
 
         val functionDescriptor = trace.bindingContext.get(BindingContext.FUNCTION, ktFunction) as? FunctionDescriptorImpl
             ?: throw AssertionError("No function descriptor for resolved lambda argument")
+
         functionDescriptor.setReturnType(returnType)
 
+        for ((i, valueParameter) in functionDescriptor.valueParameters.withIndex()) {
+            if (valueParameter !is ValueParameterDescriptorImpl || valueParameter.type !is StubType)
+                continue
+            valueParameter.setOutType(valueParameters[i])
+        }
+
         val existingLambdaType = trace.getType(ktArgumentExpression)
         if (existingLambdaType == null) {
             if (ktFunction is KtNamedFunction && ktFunction.nameIdentifier != null) return // it's a statement