K2: implement BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION

#KT-59390 Fixed

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDataClassConverters.kt

@@ -4976,6 +4976,14 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
+    add(FirErrors.BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION) { firDiagnostic ->
+        BuilderInferenceMultiLambdaRestrictionImpl(
+            firDiagnostic.a,
+            firDiagnostic.b,
+            firDiagnostic as KtPsiDiagnostic,
+            token,
+        )
+    }
     add(FirJvmErrors.OVERRIDE_CANNOT_BE_STATIC) { firDiagnostic ->
         OverrideCannotBeStaticImpl(
             firDiagnostic as KtPsiDiagnostic,

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnostics.kt

@@ -3463,6 +3463,12 @@ sealed interface KtFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
         val containingDeclarationName: Name
     }
 
+    interface BuilderInferenceMultiLambdaRestriction : KtFirDiagnostic<PsiElement> {
+        override val diagnosticClass get() = BuilderInferenceMultiLambdaRestriction::class
+        val typeParameterName: Name
+        val containingDeclarationName: Name
+    }
+
     interface OverrideCannotBeStatic : KtFirDiagnostic<PsiElement> {
         override val diagnosticClass get() = OverrideCannotBeStatic::class
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnosticsImpl.kt

@@ -4176,6 +4176,13 @@ internal class BuilderInferenceStubReceiverImpl(
     token: KtLifetimeToken,
 ) : KtAbstractFirDiagnostic<PsiElement>(firDiagnostic, token), KtFirDiagnostic.BuilderInferenceStubReceiver
 
+internal class BuilderInferenceMultiLambdaRestrictionImpl(
+    override val typeParameterName: Name,
+    override val containingDeclarationName: Name,
+    firDiagnostic: KtPsiDiagnostic,
+    token: KtLifetimeToken,
+) : KtAbstractFirDiagnostic<PsiElement>(firDiagnostic, token), KtFirDiagnostic.BuilderInferenceMultiLambdaRestriction
+
 internal class OverrideCannotBeStaticImpl(
     firDiagnostic: KtPsiDiagnostic,
     token: KtLifetimeToken,

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java

@@ -16733,6 +16733,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/builderInference"), Pattern.compile("^(.+)\\.(kt|kts)$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true);
                 }
 
+                @Test
+                @TestMetadata("builderInferenceAnnotationInLambdaWithTVExpectedType.kt")
+                public void testBuilderInferenceAnnotationInLambdaWithTVExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/builderInferenceAnnotationInLambdaWithTVExpectedType.kt");
+                }
+
                 @Test
                 @TestMetadata("changingResolveIfDontUseBuilderInference.kt")
                 public void testChangingResolveIfDontUseBuilderInference() throws Exception {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java

@@ -16733,6 +16733,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/builderInference"), Pattern.compile("^(.+)\\.(kt|kts)$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true);
                 }
 
+                @Test
+                @TestMetadata("builderInferenceAnnotationInLambdaWithTVExpectedType.kt")
+                public void testBuilderInferenceAnnotationInLambdaWithTVExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/builderInferenceAnnotationInLambdaWithTVExpectedType.kt");
+                }
+
                 @Test
                 @TestMetadata("changingResolveIfDontUseBuilderInference.kt")
                 public void testChangingResolveIfDontUseBuilderInference() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java

@@ -16727,6 +16727,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/builderInference"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
                 }
 
+                @Test
+                @TestMetadata("builderInferenceAnnotationInLambdaWithTVExpectedType.kt")
+                public void testBuilderInferenceAnnotationInLambdaWithTVExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/builderInferenceAnnotationInLambdaWithTVExpectedType.kt");
+                }
+
                 @Test
                 @TestMetadata("changingResolveIfDontUseBuilderInference.kt")
                 public void testChangingResolveIfDontUseBuilderInference() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java

@@ -16733,6 +16733,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/builderInference"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true, "multiplatform");
                 }
 
+                @Test
+                @TestMetadata("builderInferenceAnnotationInLambdaWithTVExpectedType.kt")
+                public void testBuilderInferenceAnnotationInLambdaWithTVExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/builderInferenceAnnotationInLambdaWithTVExpectedType.kt");
+                }
+
                 @Test
                 @TestMetadata("changingResolveIfDontUseBuilderInference.kt")
                 public void testChangingResolveIfDontUseBuilderInference() throws Exception {

compiler/fir/checkers/checkers-component-generator/src/org/jetbrains/kotlin/fir/checkers/generator/diagnostics/FirDiagnosticsList.kt

@@ -1770,6 +1770,10 @@ object DIAGNOSTICS_LIST : DiagnosticList("FirErrors") {
             parameter<Name>("typeParameterName")
             parameter<Name>("containingDeclarationName")
         }
+        val BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION by error<PsiElement> {
+            parameter<Name>("typeParameterName")
+            parameter<Name>("containingDeclarationName")
+        }
     }
 }
 

compiler/fir/checkers/gen/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt

@@ -871,6 +871,7 @@ object FirErrors {
 
     // Builder inference
     val BUILDER_INFERENCE_STUB_RECEIVER by error2<PsiElement, Name, Name>()
+    val BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION by error2<PsiElement, Name, Name>()
 
     init {
         RootDiagnosticRendererFactory.registerFactory(FirErrorsDefaultMessages)

compiler/fir/checkers/gen/org/jetbrains/kotlin/fir/analysis/diagnostics/FirNonSuppressibleErrorNames.kt

@@ -572,6 +572,7 @@ val FIR_NON_SUPPRESSIBLE_ERROR_NAMES: Set<String> = setOf(
     "PRE_RELEASE_CLASS",
     "IR_WITH_UNSTABLE_ABI_COMPILED_CLASS",
     "BUILDER_INFERENCE_STUB_RECEIVER",
+    "BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION",
     "OVERRIDE_CANNOT_BE_STATIC",
     "JVM_STATIC_NOT_IN_OBJECT_OR_CLASS_COMPANION",
     "JVM_STATIC_NOT_IN_OBJECT_OR_COMPANION",

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrorsDefaultMessages.kt

@@ -112,6 +112,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.BOUNDS_NOT_ALLOWE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.BOUND_ON_TYPE_ALIAS_PARAMETER_NOT_ALLOWED
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.BREAK_OR_CONTINUE_JUMPS_ACROSS_FUNCTION_BOUNDARY
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.BREAK_OR_CONTINUE_OUTSIDE_A_LOOP
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.BUILDER_INFERENCE_STUB_RECEIVER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.CALLABLE_REFERENCE_LHS_NOT_A_CLASS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.CALLABLE_REFERENCE_TO_ANNOTATION_CONSTRUCTOR
@@ -2577,6 +2578,12 @@ object FirErrorsDefaultMessages : BaseDiagnosticRendererFactory() {
             "The type of a receiver hasn''t been inferred yet. Please specify type argument for generic parameter `{0}` of `{1}` explicitly",
             TO_STRING,
             TO_STRING
-        );
+        )
+        map.put(
+            BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION,
+            "Unstable inference behaviour with multiple lambdas. Please either specify the type argument for generic parameter `{0}` of `{1}` explicitly",
+            TO_STRING,
+            TO_STRING
+        )
     }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/coneDiagnosticToFirDiagnostic.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.isLocalMember
 import org.jetbrains.kotlin.fir.analysis.getChild
 import org.jetbrains.kotlin.fir.builder.FirSyntaxErrors
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
 import org.jetbrains.kotlin.fir.declarations.utils.*
 import org.jetbrains.kotlin.fir.diagnostics.*
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.diagnostics.*
+import org.jetbrains.kotlin.fir.resolve.inference.AnonymousFunctionBasedMultiLambdaBuilderInferenceRestriction
 import org.jetbrains.kotlin.fir.resolve.inference.ConeTypeParameterBasedTypeVariable
 import org.jetbrains.kotlin.fir.resolve.inference.ConeTypeVariableForLambdaReturnType
 import org.jetbrains.kotlin.fir.resolve.inference.model.ConeArgumentConstraintPosition
@@ -529,6 +531,18 @@ private fun ConstraintSystemError.toDiagnostic(
             )
         }
 
+        is AnonymousFunctionBasedMultiLambdaBuilderInferenceRestriction -> {
+            val typeParameterSymbol = (typeParameter as ConeTypeParameterLookupTag).typeParameterSymbol
+            FirErrors.BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION.createOn(
+                anonymous.source ?: source,
+                typeParameterSymbol.name,
+                @OptIn(SymbolInternals::class)
+                (typeParameterSymbol.containingDeclarationSymbol.fir as FirMemberDeclaration).nameOrSpecialName
+            )
+        }
+
+        is MultiLambdaBuilderInferenceRestriction<*> -> shouldNotBeCalled()
+
         else -> null
     }
 }

compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/inference/AnonymousFunctionBasedMultiLambdaBuilderInferenceRestriction.kt

@@ -0,0 +1,15 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.resolve.inference
+
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.resolve.calls.inference.model.MultiLambdaBuilderInferenceRestriction
+import org.jetbrains.kotlin.types.model.TypeParameterMarker
+
+class AnonymousFunctionBasedMultiLambdaBuilderInferenceRestriction(
+    anonymous: FirAnonymousFunction,
+    typeParameter: TypeParameterMarker
+) : MultiLambdaBuilderInferenceRestriction<FirAnonymousFunction>(anonymous, typeParameter)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolutionStages.kt

@@ -20,10 +20,7 @@ import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.FirThisReference
 import org.jetbrains.kotlin.fir.resolve.directExpansionType
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
-import org.jetbrains.kotlin.fir.resolve.inference.ConeTypeParameterBasedTypeVariable
-import org.jetbrains.kotlin.fir.resolve.inference.ResolvedCallableReferenceAtom
-import org.jetbrains.kotlin.fir.resolve.inference.csBuilder
-import org.jetbrains.kotlin.fir.resolve.inference.hasBuilderInferenceAnnotation
+import org.jetbrains.kotlin.fir.resolve.inference.*
 import org.jetbrains.kotlin.fir.resolve.inference.model.ConeExplicitTypeParameterConstraintPosition
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.scopes.FirTypeScope
@@ -686,8 +683,17 @@ internal object PostponedVariablesInitializerResolutionStage : ResolutionStage()
     override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val argumentMapping = candidate.argumentMapping ?: return
         if (candidate.typeArgumentMapping is TypeArgumentMapping.Mapped) return
-        for (parameter in argumentMapping.values) {
+
+        val atomsToMark = mutableSetOf<FirAnonymousFunction>()
+
+        for ((argument, parameter) in argumentMapping) {
             if (!parameter.hasBuilderInferenceAnnotation(context.session)) continue
+            val unwrapped = argument.unwrapArgument()
+            if (unwrapped is FirAnonymousFunctionExpression) {
+                atomsToMark.add(unwrapped.anonymousFunction)
+            }
+
+            // TODO: This is effectively dead-code, since UseBuilderInferenceOnlyIfNeeded is always enabled
             val type = parameter.returnTypeRef.coneType
             val receiverType = type.receiverType(callInfo.session) ?: continue
             val dontUseBuilderInferenceIfPossible =
@@ -706,6 +712,23 @@ internal object PostponedVariablesInitializerResolutionStage : ResolutionStage()
                 }
             }
         }
+
+        candidate.postponedAtoms.forEach { postponedAtomContainer ->
+            when (postponedAtomContainer) {
+                is ResolvedLambdaAtom -> {
+                    if (postponedAtomContainer.atom in atomsToMark) {
+                        postponedAtomContainer.isCorrespondingParameterAnnotatedWithBuilderInference = true
+                    }
+                }
+                is LambdaWithTypeVariableAsExpectedTypeAtom -> {
+                    // Although, one may annotate fun foo(@BuilderInference t: T)
+                    // it makes little sense
+                }
+                is ResolvedCallableReferenceAtom -> {
+                    // Builder inference doesn't apply to such atoms.
+                }
+            }
+        }
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/ConstraintSystemCompleter.kt

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraint
 import org.jetbrains.kotlin.resolve.calls.model.PostponedAtomWithRevisableExpectedType
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
+import org.jetbrains.kotlin.types.model.TypeVariableTypeConstructorMarker
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.filterIsInstanceWithChecker
 
@@ -194,15 +195,44 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
 
         val lambdaArguments = postponedArguments.filterIsInstance<ResolvedLambdaAtom>().takeIf { it.isNotEmpty() } ?: return false
 
+        fun ResolvedLambdaAtom.notFixedInputTypeVariables(): List<TypeVariableTypeConstructorMarker> =
+            inputTypes.flatMap { it.extractTypeVariables() }.filter { it !in fixedTypeVariables }
+
+        val checkForDangerousBuilderInference =
+            !languageVersionSettings.supportsFeature(LanguageFeature.NoBuilderInferenceWithoutAnnotationRestriction)
+
+        // Let's call builder lambda (BL) a lambda that has non-zero not fixed input type variables in
+        // type arguments of it's input types
+        // ex: MutableList<T>.() -> Unit
+        // During type inference of call-site such lambda will be considered BL, if
+        // T not fixed yet
+        // Given we have two or more builder lambdas among postponed arguments, it could result in incorrect type inference due to
+        // incorrect constraint propagation into common system
+        // See KT-53740
+        // Constraint propagation into common system happens at
+        // org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzer.applyResultsOfAnalyzedLambdaToCandidateSystem
+        val dangerousBuilderInferenceWithoutAnnotation =
+            lambdaArguments.size >= 2 && lambdaArguments.count { it.notFixedInputTypeVariables().isNotEmpty() } >= 2
+
         // We assume useBuilderInferenceWithoutAnnotation = true for FIR
 
         val builder = getBuilder()
         for (argument in lambdaArguments) {
-            val notFixedInputTypeVariables = argument.inputTypes
-                .flatMap { it.extractTypeVariables() }.filter { it !in fixedTypeVariables }
+            val reallyHasBuilderInferenceAnnotation = argument.isCorrespondingParameterAnnotatedWithBuilderInference
+
+            val notFixedInputTypeVariables = argument.notFixedInputTypeVariables()
 
             if (notFixedInputTypeVariables.isEmpty()) continue
 
+            // we have dangerous inference situation
+            // if lambda annotated with BuilderInference it's probably safe, due to type shape
+            // otherwise report multi-lambda builder inference restriction diagnostic
+            if (checkForDangerousBuilderInference && dangerousBuilderInferenceWithoutAnnotation && !reallyHasBuilderInferenceAnnotation) {
+                for (variable in notFixedInputTypeVariables) {
+                    addError(AnonymousFunctionBasedMultiLambdaBuilderInferenceRestriction(argument.atom, variable.typeParameter!!))
+                }
+            }
+
             for (variable in notFixedInputTypeVariables) {
                 builder.markPostponedVariable(notFixedTypeVariables.getValue(variable).typeVariable)
             }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedAtoms.kt

@@ -77,6 +77,15 @@ class ResolvedLambdaAtom(
     fun replaceTypeVariableForLambdaReturnType(typeVariableForLambdaReturnType: ConeTypeVariableForLambdaReturnType) {
         this.typeVariableForLambdaReturnType = typeVariableForLambdaReturnType
     }
+
+    /**
+     * Set to true by resolution stage, if the corresponding parameter has BuilderInference annotation
+     * ```kotlin
+     *     fun foo(@BuilderInference b: Buildee<T>.() -> Unit) {}
+     *     fun test() = foo({ ... }) // Will be true for the lambda argument passed to b
+     * ```
+     */
+    var isCorrespondingParameterAnnotatedWithBuilderInference: Boolean = false
 }
 
 class LambdaWithTypeVariableAsExpectedTypeAtom(

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/DiagnosticReporterByTrackingStrategy.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.inference.BuilderInferenceExpectedTypeConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.*
+import org.jetbrains.kotlin.resolve.calls.model.MultiLambdaBuilderInferenceRestriction
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.calls.smartcasts.SingleSmartCast
 import org.jetbrains.kotlin.resolve.calls.smartcasts.SmartCastManager
@@ -55,6 +56,7 @@ import org.jetbrains.kotlin.types.model.freshTypeConstructor
 import org.jetbrains.kotlin.types.typeUtil.contains
 import org.jetbrains.kotlin.types.typeUtil.isNullableNothing
 import org.jetbrains.kotlin.types.typeUtil.makeNullable
+import org.jetbrains.kotlin.utils.addToStdlib.shouldNotBeCalled
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
 
@@ -812,7 +814,9 @@ class DiagnosticReporterByTrackingStrategy(
             // LowerPriorityToPreserveCompatibility is not expected to report something
             is LowerPriorityToPreserveCompatibility -> {}
             // NoSuccessfulFork does not exist in K1
-            is NoSuccessfulFork -> {}
+            is NoSuccessfulFork -> shouldNotBeCalled()
+            // MultiLambdaBuilderInferenceRestriction does not exist in K1
+            is org.jetbrains.kotlin.resolve.calls.inference.model.MultiLambdaBuilderInferenceRestriction<*> -> shouldNotBeCalled()
             // NotEnoughInformationForTypeParameterImpl is already considered above
             is NotEnoughInformationForTypeParameter<*> -> {
                 throw AssertionError("constraintError should not be called with ${error::class.java}")

compiler/resolution.common/src/org/jetbrains/kotlin/resolve/calls/inference/model/ConstraintPositionAndErrors.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability.*
 import org.jetbrains.kotlin.types.EmptyIntersectionTypeKind
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.model.TypeParameterMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 interface OnlyInputTypeConstraintPosition
@@ -173,6 +174,11 @@ class OnlyInputTypesDiagnostic(val typeVariable: TypeVariableMarker) : Constrain
 class LowerPriorityToPreserveCompatibility(val needToReportWarning: Boolean) :
     ConstraintSystemError(RESOLVED_NEED_PRESERVE_COMPATIBILITY)
 
+open class MultiLambdaBuilderInferenceRestriction<T>(
+    val anonymous: T,
+    val typeParameter: TypeParameterMarker
+) : ConstraintSystemError(RESOLVED_WITH_ERROR)
+
 fun Constraint.isExpectedTypePosition() =
     position.from is ExpectedTypeConstraintPosition<*> || position.from is DelegatedPropertyConstraintPosition<*>
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/KotlinConstraintSystemCompleter.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.resolve.calls.components.candidate.SimpleResolutionC
 import org.jetbrains.kotlin.resolve.calls.components.transformToResolvedLambda
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.*
+import org.jetbrains.kotlin.resolve.calls.model.MultiLambdaBuilderInferenceRestriction
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.error.ErrorTypeKind
 import org.jetbrains.kotlin.types.error.ErrorUtils

compiler/testData/diagnostics/tests/inference/builderInference/builderInferenceAnnotationInLambdaWithTVExpectedType.fir.kt

@@ -0,0 +1,15 @@
+// WITH_STDLIB
+package some
+
+import kotlin.experimental.ExperimentalTypeInference
+
+@OptIn(ExperimentalTypeInference::class)
+fun <T> applyBI(@BuilderInference t: T): T = t
+
+fun <V> myBuildList(a: MutableList<out V>.() -> Unit) {}
+
+fun main() {
+    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>myBuildList<!>(<!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>applyBI<!> {
+        <!CANNOT_INFER_PARAMETER_TYPE!>this<!>.<!UNRESOLVED_REFERENCE!>add<!>("1")
+    })
+}

compiler/testData/diagnostics/tests/inference/builderInference/builderInferenceAnnotationInLambdaWithTVExpectedType.kt

@@ -0,0 +1,15 @@
+// WITH_STDLIB
+package some
+
+import kotlin.experimental.ExperimentalTypeInference
+
+@OptIn(ExperimentalTypeInference::class)
+fun <T> applyBI(@BuilderInference t: T): T = t
+
+fun <V> myBuildList(a: MutableList<out V>.() -> Unit) {}
+
+fun main() {
+    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>myBuildList<!>(<!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>applyBI<!> {
+        this.<!DEBUG_INFO_MISSING_UNRESOLVED!>add<!>("1")
+    })
+}

compiler/testData/diagnostics/tests/inference/builderInference/kt53422.fir.kt

@@ -3,7 +3,7 @@
 fun test() {
     foo(
         flow { emit(0) }
-    ) { <!BUILDER_INFERENCE_STUB_RECEIVER!>it.collect <!TOO_MANY_ARGUMENTS!>{}<!><!> }
+    ) <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{ <!BUILDER_INFERENCE_STUB_RECEIVER!>it.collect <!TOO_MANY_ARGUMENTS!>{}<!><!> }<!>
 
     // 0. Initial
     // W <: Any / declared upper bound
@@ -53,4 +53,4 @@ class FlowCollectorImpl<C> : FlowCollector<C> {
     override fun emit(value: C) {}
 }
 
-fun Flow<*>.collect() {}
+fun Flow<*>.collect() {}

compiler/testData/diagnostics/tests/inference/builderInference/kt53639.fir.kt

@@ -12,7 +12,7 @@ data class Output(val source: InputWrapper<List<String>>)
 fun main2(input: InputWrapper<Unit>): Output {
     val output = input.doMapping(
         foo = { buildList { add("this is List<String>") } },
-        bar = { it.isNotEmpty() },
+        bar = <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{ it.isNotEmpty() }<!>,
     )
 
     return Output(source = output)

compiler/testData/diagnostics/tests/inference/builderInference/multiLambdaRestriction.fir.kt

@@ -9,9 +9,9 @@ fun <R> a(first: R, second: (List<R>) -> Unit) {}
 fun test1() {
     a(
         buildList { add("") },
-        second = {
+        second = <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{
             it.myGenericExt()
-        }
+        }<!>
     )
 }
 
@@ -23,9 +23,9 @@ fun test2() {
         first = {
             buildList { add("") }
         },
-        second = {
+        second = <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{
             it.myExt() // Note: must be extension to add constraints
-        }
+        }<!>
     )
 }
 
@@ -45,13 +45,13 @@ fun test3() {
     )
 
     select (
-        myBuildList { add("") },
-        myBuildList { add(1) },
+        myBuildList <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{ add("") }<!>,
+        myBuildList <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{ add(1) }<!>,
     )
 
     select (
-        run { myBuildList { add("") } },
-        myBuildList { add(1) },
+        run { myBuildList <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{ add("") }<!> },
+        myBuildList <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{ add(1) }<!>,
     )
 }
 
@@ -64,7 +64,7 @@ fun <D> buildPartList(left: MutableList<D>.() -> Unit, right: MutableList<D>.()
 
 fun test4() {
     buildPartList(
-        left = { add(1) },
-        right = { add("") }
+        left = <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{ add(1) }<!>,
+        right = <!BUILDER_INFERENCE_MULTI_LAMBDA_RESTRICTION!>{ add("") }<!>
     )
-}
+}

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

@@ -16733,6 +16733,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/builderInference"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.(reversed|fir|ll)\\.kts?$"), true);
                 }
 
+                @Test
+                @TestMetadata("builderInferenceAnnotationInLambdaWithTVExpectedType.kt")
+                public void testBuilderInferenceAnnotationInLambdaWithTVExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/builderInferenceAnnotationInLambdaWithTVExpectedType.kt");
+                }
+
                 @Test
                 @TestMetadata("changingResolveIfDontUseBuilderInference.kt")
                 public void testChangingResolveIfDontUseBuilderInference() throws Exception {