Implement top-down completion for nested builder inference calls

^KT-42742 Fixed

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

@@ -16987,6 +16987,18 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @Test
+            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
+            }
+
+            @Test
+            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
+            }
+
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {

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

@@ -56,10 +56,18 @@ class CoroutineInferenceSession(
 ) : ManyCandidatesResolver<CallableDescriptor>(
     psiCallResolver, postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter, callComponents, builtIns
 ) {
+    var nestedBuilderInferenceSession: CoroutineInferenceSession? = null
+
+    init {
+        if (topLevelCallContext.inferenceSession is CoroutineInferenceSession) {
+            topLevelCallContext.inferenceSession.nestedBuilderInferenceSession = this
+        }
+    }
+
     private val commonCalls = arrayListOf<PSICompletedCallInfo>()
 
     // Simple calls are calls which might not have gone through type inference, but may contain unsubstituted postponed variables inside their types.
-    private val simpleCommonCalls = arrayListOf<KtExpression>()
+    private val oldCallableReferenceCalls = arrayListOf<KtExpression>()
 
     private var hasInapplicableCall = false
 
@@ -105,8 +113,8 @@ class CoroutineInferenceSession(
         }
     }
 
-    fun addSimpleCall(callExpression: KtExpression) {
+    fun addOldCallableReferenceCalls(callExpression: KtExpression) {
-        simpleCommonCalls.add(callExpression)
+        oldCallableReferenceCalls.add(callExpression)
     }
 
     override fun addCompletedCallInfo(callInfo: CompletedCallInfo) {
@@ -139,6 +147,8 @@ class CoroutineInferenceSession(
                 descriptor.extensionReceiverParameter?.type?.contains { it is StubType } == true
     }
 
+    private fun isTopLevelBuilderInferenceCall() = topLevelCallContext.inferenceSession !is CoroutineInferenceSession
+
     fun hasInapplicableCall(): Boolean = hasInapplicableCall
 
     override fun writeOnlyStubs(callInfo: SingleCallResolutionResult): Boolean {
@@ -178,42 +188,63 @@ class CoroutineInferenceSession(
     ): Map<TypeConstructor, UnwrappedType>? {
         val (commonSystem, effectivelyEmptyConstraintSystem) = buildCommonSystem(initialStorage)
         val initialStorageSubstitutor = initialStorage.buildResultingSubstitutor(commonSystem, transformTypeVariablesToErrorTypes = false)
+
         if (effectivelyEmptyConstraintSystem) {
-            updateCalls(
+            if (isTopLevelBuilderInferenceCall()) {
-                lambda,
+                updateAllCalls(initialStorageSubstitutor, commonSystem, lambda)
-                initialStorageSubstitutor,
+            }
-                commonSystem.errors
-            )
             return null
         }
 
-        val context = commonSystem.asConstraintSystemCompleterContext()
         kotlinConstraintSystemCompleter.completeConstraintSystem(
-            context,
+            commonSystem.asConstraintSystemCompleterContext(),
             builtIns.unitType,
             partiallyResolvedCallsInfo.map { it.callResolutionResult.resultCallAtom },
             completionMode,
             diagnosticsHolder
         )
 
-        val resultingSubstitutor =
+        if (isTopLevelBuilderInferenceCall()) {
-            ComposedSubstitutor(initialStorageSubstitutor, commonSystem.buildCurrentSubstitutor() as NewTypeSubstitutor)
+            updateAllCalls(initialStorageSubstitutor, commonSystem, lambda)
-        updateCalls(lambda, resultingSubstitutor, commonSystem.errors)
+        }
 
         return commonSystem.fixedTypeVariables.cast() // TODO: SUB
     }
 
+    /*
+     * We update calls in top-down way:
+     * - updating calls within top-level builder inference call
+     * - ...
+     * - updating calls within the deepest builder inference call
+     */
+    private fun updateAllCalls(substitutor: NewTypeSubstitutor, commonSystem: NewConstraintSystemImpl, lambda: ResolvedLambdaAtom) {
+        val resultingSubstitutor = ComposedSubstitutor(substitutor, commonSystem.buildCurrentSubstitutor() as NewTypeSubstitutor)
+
+        updateCalls(lambda, resultingSubstitutor, commonSystem.errors)
+
+        nestedBuilderInferenceSession?.updateAllCalls(substitutor, commonSystem, lambda)
+    }
+
     override fun shouldCompleteResolvedSubAtomsOf(resolvedCallAtom: ResolvedCallAtom) = true
 
-    private fun createNonFixedTypeToVariableSubstitutor(): NewTypeSubstitutorByConstructorMap {
+    private fun createNonFixedTypeToVariableMap(): Map<TypeConstructor, UnwrappedType> {
         val bindings = hashMapOf<TypeConstructor, UnwrappedType>()
-        for ((variable, nonFixedType) in stubsForPostponedVariables) {
+
+        for ((variable, nonFixedType) in stubsForPostponedVariables) { // do it for nested sessions
             bindings[nonFixedType.constructor] = variable.defaultType
         }
 
-        return NewTypeSubstitutorByConstructorMap(bindings)
+        val parentBuilderInferenceCallSession = topLevelCallContext.inferenceSession as? CoroutineInferenceSession
+
+        if (parentBuilderInferenceCallSession != null) {
+            bindings.putAll(parentBuilderInferenceCallSession.createNonFixedTypeToVariableMap())
+        }
+
+        return bindings
     }
 
+    private fun createNonFixedTypeToVariableSubstitutor() = NewTypeSubstitutorByConstructorMap(createNonFixedTypeToVariableMap())
+
     private fun integrateConstraints(
         commonSystem: NewConstraintSystemImpl,
         storage: ConstraintStorage,
@@ -299,34 +330,6 @@ class CoroutineInferenceSession(
         )
     }
 
-    private fun updateCalls(lambda: ResolvedLambdaAtom, substitutor: NewTypeSubstitutor, errors: List<ConstraintSystemError>) {
-        val nonFixedToVariablesSubstitutor = createNonFixedTypeToVariableSubstitutor()
-
-        val nonFixedTypesToResult = nonFixedToVariablesSubstitutor.map.mapValues { substitutor.safeSubstitute(it.value) }
-        val nonFixedTypesToResultSubstitutor = ComposedSubstitutor(substitutor, nonFixedToVariablesSubstitutor)
-
-        val atomCompleter = createResolvedAtomCompleter(nonFixedTypesToResultSubstitutor, topLevelCallContext)
-
-        for (completedCall in commonCalls) {
-            updateCall(completedCall, nonFixedTypesToResultSubstitutor, nonFixedTypesToResult)
-            reportErrors(completedCall, completedCall.resolvedCall, errors)
-        }
-
-        for (callInfo in partiallyResolvedCallsInfo) {
-            val resolvedCall = completeCall(callInfo, atomCompleter) ?: continue
-            reportErrors(callInfo, resolvedCall, errors)
-        }
-
-        for (simpleCall in simpleCommonCalls) {
-            when (simpleCall) {
-                is KtCallableReferenceExpression -> updateCallableReferenceType(simpleCall, nonFixedTypesToResultSubstitutor)
-                else -> throw Exception("Unsupported call expression type")
-            }
-        }
-
-        atomCompleter.completeAll(lambda)
-    }
-
     private fun updateCall(
         completedCall: PSICompletedCallInfo,
         nonFixedTypesToResultSubstitutor: NewTypeSubstitutor,
@@ -346,26 +349,12 @@ class CoroutineInferenceSession(
         completeCall(completedCall, atomCompleter)
     }
 
-    private fun updateCallableReferenceType(expression: KtCallableReferenceExpression, substitutor: NewTypeSubstitutor) {
+    private fun completeCallableReference(expression: KtCallableReferenceExpression, substitutor: NewTypeSubstitutor) {
-        val functionDescriptor = trace.get(BindingContext.DECLARATION_TO_DESCRIPTOR, expression) as? SimpleFunctionDescriptorImpl ?: return
+        createResolvedAtomCompleter(substitutor, topLevelCallContext).substituteFunctionLiteralDescriptor(
-        val returnType = functionDescriptor.returnType
+            resolvedAtom = null,
-
+            descriptor = trace.get(BindingContext.DECLARATION_TO_DESCRIPTOR, expression) as? SimpleFunctionDescriptorImpl ?: return,
-        fun KotlinType.substituteAndApproximate() = typeApproximator.approximateDeclarationType(
+            substitutor
-            substitutor.safeSubstitute(this.unwrap()),
-            local = true,
-            languageVersionSettings = topLevelCallContext.languageVersionSettings
         )
-
-        if (returnType != null && returnType.contains { it is StubType }) {
-            functionDescriptor.setReturnType(returnType.substituteAndApproximate())
-        }
-
-        for (valueParameter in functionDescriptor.valueParameters) {
-            if (valueParameter !is ValueParameterDescriptorImpl || valueParameter.type !is StubType)
-                continue
-
-            valueParameter.setOutType(valueParameter.type.substituteAndApproximate())
-        }
     }
 
     private fun completeCall(
@@ -408,6 +397,40 @@ class CoroutineInferenceSession(
             containingElement.anyDescendantOfType<KtElement> { it == atom.psiCall.callElement }
         }
     }
+
+    companion object {
+        private fun CoroutineInferenceSession.updateCalls(
+            lambda: ResolvedLambdaAtom,
+            substitutor: NewTypeSubstitutor,
+            errors: List<ConstraintSystemError>
+        ) {
+            val nonFixedToVariablesSubstitutor = createNonFixedTypeToVariableSubstitutor()
+
+            val nonFixedTypesToResult = nonFixedToVariablesSubstitutor.map.mapValues { substitutor.safeSubstitute(it.value) }
+            val nonFixedTypesToResultSubstitutor = ComposedSubstitutor(substitutor, nonFixedToVariablesSubstitutor)
+
+            val atomCompleter = createResolvedAtomCompleter(nonFixedTypesToResultSubstitutor, topLevelCallContext)
+
+            for (completedCall in commonCalls) {
+                updateCall(completedCall, nonFixedTypesToResultSubstitutor, nonFixedTypesToResult)
+                reportErrors(completedCall, completedCall.resolvedCall, errors)
+            }
+
+            for (callInfo in partiallyResolvedCallsInfo) {
+                val resolvedCall = completeCall(callInfo, atomCompleter) ?: continue
+                reportErrors(callInfo, resolvedCall, errors)
+            }
+
+            for (simpleCall in oldCallableReferenceCalls) {
+                when (simpleCall) {
+                    is KtCallableReferenceExpression -> completeCallableReference(simpleCall, nonFixedTypesToResultSubstitutor)
+                    else -> throw Exception("Unsupported call expression type")
+                }
+            }
+
+            atomCompleter.completeAll(lambda)
+        }
+    }
 }
 
 class ComposedSubstitutor(val left: NewTypeSubstitutor, val right: NewTypeSubstitutor) : NewTypeSubstitutor {

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

@@ -8,8 +8,8 @@ package org.jetbrains.kotlin.resolve.calls.tower
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 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.SimpleFunctionDescriptorImpl
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.diagnostics.reportDiagnosticOnce
@@ -48,7 +48,7 @@ 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.shouldBeSubstituted
+import org.jetbrains.kotlin.types.typeUtil.shouldBeUpdated
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class ResolvedAtomCompleter(
@@ -212,52 +212,93 @@ class ResolvedAtomCompleter(
             return commonReturnType.isUnit()
         }
 
-    private fun completeLambda(lambda: ResolvedLambdaAtom) {
+    private fun KotlinType.substituteAndApproximate(substitutor: NewTypeSubstitutor): FunctionLiteralTypes.ProcessedType {
-        @Suppress("NAME_SHADOWING")
+        val substitutedType = substitutor.safeSubstitute(this.unwrap())
-        val lambda = lambda.unwrap()
+
+        return FunctionLiteralTypes.ProcessedType(
+            substitutedType,
+            approximatedType = typeApproximator.approximateDeclarationType(
+                substitutedType, local = true, languageVersionSettings = topLevelCallContext.languageVersionSettings
+            )
+        )
+    }
+
+    fun substituteFunctionLiteralDescriptor(
+        resolvedAtom: ResolvedLambdaAtom?, // null is for callable references resolved though the old type inference
+        descriptor: SimpleFunctionDescriptorImpl,
+        substitutor: NewTypeSubstitutor
+    ): FunctionLiteralTypes {
+        val returnType =
+            (if (resolvedAtom?.isCoercedToUnit == true) builtIns.unitType else resolvedAtom?.returnType) ?: descriptor.returnType
+        val receiverType = resolvedAtom?.receiver ?: descriptor.extensionReceiverParameter?.type
+        val valueParameterTypes = resolvedAtom?.parameters ?: descriptor.valueParameters.map { it.type }
+
+        require(returnType != null)
+
+        val substitutedReturnType = returnType.substituteAndApproximate(substitutor).also {
+            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 substitutedValueParameterTypes = descriptor.valueParameters.mapIndexed { i, valueParameter ->
+            valueParameterTypes[i].substituteAndApproximate(substitutor).also {
+                if (valueParameter is ValueParameterDescriptorImpl && valueParameter.type.shouldBeUpdated()) {
+                    valueParameter.setOutType(it.approximatedType)
+                }
+            }
+        }
+
+        return FunctionLiteralTypes(substitutedReturnType, substitutedValueParameterTypes, substitutedReceiverType)
+    }
+
+    private fun completeLambda(resolvedAtom: ResolvedLambdaAtom) {
+        val lambda = resolvedAtom.unwrap()
         val resultArgumentsInfo = lambda.resultArgumentsInfo!!
-        val returnType = if (lambda.isCoercedToUnit) {
-            builtIns.unitType
-        } else {
-            resultSubstitutor.safeSubstitute(lambda.returnType)
-        }
-        val receiverType = lambda.receiver
 
-        val approximatedValueParameterTypes = lambda.parameters.map { parameterType ->
+        val psiCallArgument = lambda.atom.psiCallArgument
-            if (parameterType.shouldBeSubstituted()) {
+        val (ktArgumentExpression, ktFunction) = when (psiCallArgument) {
-                typeApproximator.approximateDeclarationType(
+            is LambdaKotlinCallArgumentImpl -> psiCallArgument.ktLambdaExpression to psiCallArgument.ktLambdaExpression.functionLiteral
-                    resultSubstitutor.safeSubstitute(parameterType),
+            is FunctionExpressionImpl -> psiCallArgument.ktFunction to psiCallArgument.ktFunction
-                    local = true,
+            else -> throw AssertionError("Unexpected psiCallArgument for resolved lambda argument: $psiCallArgument")
-                    languageVersionSettings = topLevelCallContext.languageVersionSettings
-                )
-            } else parameterType
         }
 
-        val approximatedReturnType =
+        val descriptor = topLevelTrace.bindingContext.get(BindingContext.FUNCTION, ktFunction) as? SimpleFunctionDescriptorImpl
-            typeApproximator.approximateDeclarationType(
+            ?: throw AssertionError("No function descriptor for resolved lambda argument")
-                returnType,
-                local = true,
-                languageVersionSettings = topLevelCallContext.languageVersionSettings
-            )
 
-        val approximatedReceiverType = if (receiverType != null) {
+        val substitutedLambdaTypes = substituteFunctionLiteralDescriptor(lambda, descriptor, resultSubstitutor)
-            typeApproximator.approximateDeclarationType(
-                resultSubstitutor.safeSubstitute(receiverType),
-                local = true,
-                languageVersionSettings = topLevelCallContext.languageVersionSettings
-            )
-        } else null
 
-        updateTraceForLambda(lambda, topLevelTrace, approximatedReturnType, approximatedValueParameterTypes, approximatedReceiverType)
+        val existingLambdaType = topLevelTrace.getType(ktArgumentExpression)
+
+        if (existingLambdaType == null) {
+            if (ktFunction is KtNamedFunction && ktFunction.nameIdentifier != null) return // it's a statement
+            throw AssertionError("No type for resolved lambda argument: ${ktArgumentExpression.text}")
+        }
+
+        val substitutedFunctionalType = createFunctionType(
+            builtIns,
+            existingLambdaType.annotations,
+            substitutedLambdaTypes.receiverType?.substitutedType,
+            substitutedLambdaTypes.parameterTypes.map { it.substitutedType },
+            null, // parameter names transforms to special annotations, so they are already taken from parameter types
+            substitutedLambdaTypes.returnType.substitutedType,
+            lambda.isSuspend
+        )
+
+        topLevelTrace.recordType(ktArgumentExpression, substitutedFunctionalType)
 
         for (lambdaResult in resultArgumentsInfo.nonErrorArguments) {
             val resultValueArgument = lambdaResult as? PSIKotlinCallArgument ?: continue
-            val newContext =
+            val newContext = topLevelCallContext.replaceDataFlowInfo(resultValueArgument.dataFlowInfoAfterThisArgument)
-                topLevelCallContext.replaceDataFlowInfo(resultValueArgument.dataFlowInfoAfterThisArgument)
+                .replaceExpectedType(substitutedLambdaTypes.returnType.approximatedType)
-                    .replaceExpectedType(approximatedReturnType)
+                .replaceBindingTrace(topLevelTrace)
-                    .replaceBindingTrace(topLevelTrace)
-
             val argumentExpression = resultValueArgument.valueArgument.getArgumentExpression() ?: continue
+
             kotlinToResolvedCallTransformer.updateRecordedType(
                 argumentExpression,
                 parameter = null,
@@ -268,64 +309,6 @@ class ResolvedAtomCompleter(
         }
     }
 
-    private fun updateTraceForLambda(
-        lambda: ResolvedLambdaAtom,
-        trace: BindingTrace,
-        returnType: UnwrappedType,
-        valueParameters: List<UnwrappedType>,
-        receiverType: UnwrappedType?
-    ) {
-        val psiCallArgument = lambda.atom.psiCallArgument
-
-        val ktArgumentExpression: KtExpression
-        val ktFunction: KtElement
-        when (psiCallArgument) {
-            is LambdaKotlinCallArgumentImpl -> {
-                ktArgumentExpression = psiCallArgument.ktLambdaExpression
-                ktFunction = ktArgumentExpression.functionLiteral
-            }
-            is FunctionExpressionImpl -> {
-                ktArgumentExpression = psiCallArgument.ktFunction
-                ktFunction = ktArgumentExpression
-            }
-            else -> throw AssertionError("Unexpected psiCallArgument for resolved lambda argument: $psiCallArgument")
-        }
-
-        val functionDescriptor = trace.bindingContext.get(BindingContext.FUNCTION, ktFunction) as? FunctionDescriptorImpl
-            ?: throw AssertionError("No function descriptor for resolved lambda argument")
-
-        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])
-        }
-
-        val existingLambdaType = trace.getType(ktArgumentExpression)
-        if (existingLambdaType == null) {
-            if (ktFunction is KtNamedFunction && ktFunction.nameIdentifier != null) return // it's a statement
-
-            throw AssertionError("No type for resolved lambda argument: ${ktArgumentExpression.text}")
-        }
-        val substitutedFunctionalType = createFunctionType(
-            builtIns,
-            existingLambdaType.annotations,
-            lambda.receiver?.let { resultSubstitutor.safeSubstitute(it) },
-            lambda.parameters.map { resultSubstitutor.safeSubstitute(it) },
-            null, // parameter names transforms to special annotations, so they are already taken from parameter types
-            returnType,
-            lambda.isSuspend
-        )
-
-        trace.recordType(ktArgumentExpression, substitutedFunctionalType)
-    }
-
     private fun NewTypeSubstitutor.toOldSubstitution(): TypeSubstitution = object : TypeSubstitution() {
         override fun get(key: KotlinType): TypeProjection? {
             return safeSubstitute(key.unwrap()).takeIf { it !== key }?.asTypeProjection()
@@ -567,3 +550,11 @@ class ResolvedAtomCompleter(
         expressionTypingServices.getTypeInfo(psiCallArgument.collectionLiteralExpression, actualContext)
     }
 }
+
+class FunctionLiteralTypes(
+    val returnType: ProcessedType,
+    val parameterTypes: List<ProcessedType>,
+    val receiverType: ProcessedType?
+) {
+    class ProcessedType(val substitutedType: KotlinType, val approximatedType: KotlinType)
+}

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

@@ -554,7 +554,7 @@ class DoubleColonExpressionResolver(
         val dataFlowInfo = (lhs as? DoubleColonLHS.Expression)?.dataFlowInfo ?: c.dataFlowInfo
 
         if (c.inferenceSession is CoroutineInferenceSession && result?.contains { it is StubType } == true) {
-            c.inferenceSession.addSimpleCall(expression)
+            c.inferenceSession.addOldCallableReferenceCalls(expression)
         }
 
         return dataFlowAnalyzer.checkType(createTypeInfo(result, dataFlowInfo), expression, c)

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

@@ -0,0 +1,38 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER -EXPERIMENTAL_IS_NOT_ENABLED
+// WITH_RUNTIME
+
+import kotlin.experimental.ExperimentalTypeInference
+
+@OptIn(ExperimentalTypeInference::class)
+class A1<T> {
+    fun <BT1> builder1(@BuilderInference configure: A2<BT1>.() -> Unit) {}
+}
+
+@OptIn(ExperimentalTypeInference::class)
+class A2<A2_BT1> {
+    fun <BT2> builder2(@BuilderInference configure: A3<A2_BT1, BT2>.() -> Unit) {}
+}
+
+@OptIn(ExperimentalTypeInference::class)
+class A3<A3_BT1, A3_BT2> {
+    fun <BT3> builder3(@BuilderInference configure: A4<A3_BT1, A3_BT2, BT3>.() -> Unit) {}
+}
+
+class A4<A3_BT1, A3_BT2, A3_BT3> {
+    fun resolver(x: A3_BT3) {}
+}
+
+fun foo(x: A1<String>) {
+    x.builder1<String> {
+        builder2<String> {
+            builder3 {
+                resolver("")
+            }
+        }
+    }
+}
+
+fun box(): String {
+    foo(A1())
+    return "OK"
+}

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

@@ -0,0 +1,29 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER -EXPERIMENTAL_IS_NOT_ENABLED
+// WITH_RUNTIME
+
+import kotlin.experimental.ExperimentalTypeInference
+
+@OptIn(ExperimentalTypeInference::class)
+class A1<T> {
+    fun <BT1> builder1(@BuilderInference configure: A2<BT1>.() -> Unit) {}
+}
+
+@OptIn(ExperimentalTypeInference::class)
+class A2<A2_BT1> {
+    fun <BT2> builder2(@BuilderInference configure: A3<A2_BT1, BT2>.() -> Unit) {}
+}
+
+class A3<A3_BT1, A3_BT2> {
+    fun resolver(x: A3_BT2) {}
+}
+
+fun foo(x: A1<String>) {
+    x.builder1<String> {
+        builder2 { resolver("") }
+    }
+}
+
+fun box(): String {
+    foo(A1<String>())
+    return "OK"
+}

compiler/testData/diagnostics/tests/AutoCreatedIt.kt

@@ -3,8 +3,8 @@ fun text() {
     "direct:a" to "mock:a"
     "direct:a" on {it.body == "<hello/>"} to "mock:a"
     "direct:a" on {it -> it.body == "<hello/>"} to "mock:a"
-    bar <!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>{<!>1}
+    bar <!TYPE_MISMATCH!><!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>{<!>1}<!>
-    bar <!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>{<!><!UNRESOLVED_REFERENCE!>it<!> <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE{OI}, DEBUG_INFO_MISSING_UNRESOLVED{NI}!>+<!> 1}
+    bar <!TYPE_MISMATCH!><!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>{<!><!UNRESOLVED_REFERENCE!>it<!> <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE{OI}, DEBUG_INFO_MISSING_UNRESOLVED{NI}!>+<!> 1}<!>
     bar {it, <!UNUSED_ANONYMOUS_PARAMETER!>it1<!> -> it}
 
     bar1 {1}
@@ -13,7 +13,7 @@ fun text() {
     bar2 <!TYPE_MISMATCH{NI}!>{<!TYPE_MISMATCH!><!>}<!>
     bar2 {1}
     bar2 <!TYPE_MISMATCH{NI}!>{<!UNRESOLVED_REFERENCE!>it<!>}<!>
-    bar2 <!TYPE_MISMATCH{NI}!>{<!CANNOT_INFER_PARAMETER_TYPE, EXPECTED_PARAMETERS_NUMBER_MISMATCH!>it<!> -> <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE!>it<!>}<!>
+    bar2 <!TYPE_MISMATCH{NI}!>{<!CANNOT_INFER_PARAMETER_TYPE, EXPECTED_PARAMETERS_NUMBER_MISMATCH!>it<!> -> <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE{OI}, TYPE_MISMATCH!>it<!>}<!>
 }
 
 fun bar(<!UNUSED_PARAMETER!>f<!> :  (Int, Int) -> Int) {}

compiler/testData/diagnostics/tests/functionLiterals/ExpectedParametersTypesMismatch.kt

@@ -30,9 +30,9 @@ fun test1() {
     }<!>
 
 
-    foo2 <!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>{<!>
+    foo2 <!TYPE_MISMATCH!><!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>{<!>
         ""
-    }
+    }<!>
     foo2 <!TYPE_MISMATCH{NI}!>{
         <!EXPECTED_PARAMETERS_NUMBER_MISMATCH, EXPECTED_PARAMETER_TYPE_MISMATCH!><!UNUSED_ANONYMOUS_PARAMETER!>s<!>: String<!> -> ""
     }<!>

compiler/testData/diagnostics/tests/inference/completionOfMultipleLambdas.kt

@@ -11,7 +11,7 @@ fun test() {
     <!DEBUG_INFO_EXPRESSION_TYPE("(A, A) -> kotlin.Unit")!>select3(
         { a: A, b: A -> Unit },
         { b, a -> Unit },
-        <!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>{<!> <!UNRESOLVED_REFERENCE!>it<!>; Unit }
+        <!TYPE_MISMATCH!><!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>{<!> <!UNRESOLVED_REFERENCE!>it<!>; Unit }<!>
     )<!>
 }
 

compiler/testData/diagnostics/tests/regressions/kt30245.kt

@@ -46,9 +46,9 @@ fun test1() { // to extension lambda 0
 
     val w11 = W1 <!TYPE_MISMATCH!>{ <!EXPECTED_PARAMETERS_NUMBER_MISMATCH!>i: Int<!> -> i }<!> // oi- ni-
     val i11: E0 = id { i: Int -> i } // o1+ ni+
-    val w12 = W1 <!TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE, EXPECTED_PARAMETERS_NUMBER_MISMATCH!>i<!> -> <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE!>i<!> }<!> // oi- ni-
+    val w12 = W1 <!TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE, EXPECTED_PARAMETERS_NUMBER_MISMATCH!>i<!> -> <!TYPE_MISMATCH!>i<!> }<!> // oi- ni-
-    val i12: E0 = id <!TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE, EXPECTED_PARAMETERS_NUMBER_MISMATCH!>i<!> -> <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE!>i<!> }<!> // oi- ni-
+    val i12: E0 = id <!TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE, EXPECTED_PARAMETERS_NUMBER_MISMATCH!>i<!> -> <!TYPE_MISMATCH!>i<!> }<!> // oi- ni-
-    val j12 = id<E0> <!TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE, EXPECTED_PARAMETERS_NUMBER_MISMATCH!>i<!> -> <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE!>i<!> }<!> // oi- ni-
+    val j12 = id<E0> <!TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE, EXPECTED_PARAMETERS_NUMBER_MISMATCH!>i<!> -> <!TYPE_MISMATCH!>i<!> }<!> // oi- ni-
 
     // yet unsupported cases - considering lambdas as extension ones unconditionally
 //    val w13 = W1 { it } // this or it: oi- ni-

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

@@ -16987,6 +16987,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @Test
+            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
+            }
+
+            @Test
+            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
+            }
+
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {

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

@@ -16987,6 +16987,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @Test
+            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
+            }
+
+            @Test
+            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
+            }
+
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {

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

@@ -14067,6 +14067,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
+            }
+
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");

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

@@ -57,7 +57,7 @@ public class ReceiverParameterDescriptorImpl extends AbstractReceiverParameterDe
     }
 
     public void setOutType(@NotNull KotlinType outType) {
-        assert TypeUtilsKt.shouldBeSubstituted(this.value.getType());
+        assert TypeUtilsKt.shouldBeUpdated(this.value.getType());
         this.value = value.replaceType(outType);
     }
 }

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

@@ -50,7 +50,7 @@ public abstract class VariableDescriptorImpl extends DeclarationDescriptorNonRoo
     }
 
     public void setOutType(KotlinType outType) {
-        assert this.outType == null || TypeUtilsKt.shouldBeSubstituted(this.outType);
+        assert this.outType == null || TypeUtilsKt.shouldBeUpdated(this.outType);
         this.outType = outType;
     }
 

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

@@ -308,4 +308,4 @@ private fun NewCapturedType.unCaptureTopLevelType(): UnwrappedType {
     return constructor.projection.type.unwrap()
 }
 
-fun KotlinType.shouldBeSubstituted() = contains { it is StubType || it.constructor is TypeVariableTypeConstructorMarker }
+fun KotlinType.shouldBeUpdated() = contains { it is StubType || it.constructor is TypeVariableTypeConstructorMarker || it.isError }

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

@@ -12356,6 +12356,16 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
+            }
+
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");

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

@@ -11777,6 +11777,16 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
+            }
+
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");

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

@@ -11842,6 +11842,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
+            }
+
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");

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

@@ -6288,6 +6288,16 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
                 runTest("compiler/testData/codegen/box/inference/builderInference/lackOfNullCheckOnNullableInsideBuild.kt");
             }
 
+            @TestMetadata("multiStepCompletionWithinThreeBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinThreeBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinThreeBuilderInferenceCalls.kt");
+            }
+
+            @TestMetadata("multiStepCompletionWithinTwoBuilderInferenceCalls.kt")
+            public void testMultiStepCompletionWithinTwoBuilderInferenceCalls() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/multiStepCompletionWithinTwoBuilderInferenceCalls.kt");
+            }
+
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");