Inject stub type variables of a different builder inference call properly

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

@@ -54,7 +54,8 @@ class BuilderInferenceSession(
     private val deprecationResolver: DeprecationResolver,
     private val moduleDescriptor: ModuleDescriptor,
     private val typeApproximator: TypeApproximator,
-    private val missingSupertypesResolver: MissingSupertypesResolver
+    private val missingSupertypesResolver: MissingSupertypesResolver,
+    private val lambdaArgument: LambdaKotlinCallArgument
 ) : ManyCandidatesResolver<CallableDescriptor>(
     psiCallResolver, postponedArgumentsAnalyzer, kotlinConstraintSystemCompleter, callComponents, builtIns
 ) {
@@ -204,9 +205,10 @@ class BuilderInferenceSession(
         return lhs.isAncestor(callElement)
     }
 
-    override fun currentConstraintSystem(): ConstraintStorage {
-        return ConstraintStorage.Empty
-    }
+    override fun currentConstraintSystem() = ConstraintStorage.Empty
+
+    fun getNotFixedToInferredTypesSubstitutor(): NewTypeSubstitutor =
+        ComposedSubstitutor(commonSystem.buildCurrentSubstitutor() as NewTypeSubstitutor, createNonFixedTypeToVariableSubstitutor())
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
@@ -214,14 +216,18 @@ class BuilderInferenceSession(
         completionMode: ConstraintSystemCompletionMode,
         diagnosticsHolder: KotlinDiagnosticsHolder,
     ): Map<TypeConstructor, UnwrappedType>? {
-        val effectivelyEmptyConstraintSystem = initializeCommonSystem(initialStorage)
-        val initialStorageSubstitutor = initialStorage.buildResultingSubstitutor(commonSystem, transformTypeVariablesToErrorTypes = false)
-
         this.lambda = lambda
 
+        fun getResultingSubstitutor(): NewTypeSubstitutor {
+            val storageSubstitutor = initialStorage.buildResultingSubstitutor(commonSystem, transformTypeVariablesToErrorTypes = false)
+            return ComposedSubstitutor(storageSubstitutor, commonSystem.buildCurrentSubstitutor() as NewTypeSubstitutor)
+        }
+
+        val effectivelyEmptyConstraintSystem = initializeCommonSystem(initialStorage)
+
         if (effectivelyEmptyConstraintSystem) {
             if (isTopLevelBuilderInferenceCall()) {
-                updateAllCalls(initialStorageSubstitutor)
+                updateAllCalls(getResultingSubstitutor())
             }
             return null
         }
@@ -235,7 +241,7 @@ class BuilderInferenceSession(
         )
 
         if (isTopLevelBuilderInferenceCall()) {
-            updateAllCalls(initialStorageSubstitutor)
+            updateAllCalls(getResultingSubstitutor())
         }
 
         return commonSystem.fixedTypeVariables.cast() // TODO: SUB
@@ -250,12 +256,15 @@ class BuilderInferenceSession(
     private fun updateAllCalls(substitutor: NewTypeSubstitutor) {
         updateCalls(
             lambda,
-            substitutor = ComposedSubstitutor(substitutor, commonSystem.buildCurrentSubstitutor() as NewTypeSubstitutor),
+            substitutor = substitutor,
             commonSystem.errors
         )
 
         for (nestedSession in nestedBuilderInferenceSessions) {
-            nestedSession.updateAllCalls(substitutor)
+            // TODO: exclude injected variables
+            nestedSession.updateAllCalls(
+                ComposedSubstitutor(nestedSession.commonSystem.buildCurrentSubstitutor() as NewTypeSubstitutor, substitutor)
+            )
         }
     }
 
@@ -291,7 +300,14 @@ class BuilderInferenceSession(
         }
 
         for (parentSession in findAllParentBuildInferenceSessions()) {
-            parentSession.stubsForPostponedVariables.keys.forEach { commonSystem.registerVariable(it) }
+            for ((variable, stubType) in parentSession.stubsForPostponedVariables) {
+                commonSystem.registerVariable(variable)
+                commonSystem.addSubtypeConstraint(
+                    variable.defaultType,
+                    stubType,
+                    InjectedAnotherStubTypeConstraintPositionImpl(lambdaArgument)
+                )
+            }
         }
 
         /*
@@ -435,7 +451,7 @@ class BuilderInferenceSession(
         return currentSession.topLevelCallContext.trace
     }
 
-    private fun completeDoubleColonExpression(expression: KtDoubleColonExpression, substitutor: NewTypeSubstitutor) {
+    fun completeDoubleColonExpression(expression: KtDoubleColonExpression, substitutor: NewTypeSubstitutor) {
         val atomCompleter = createResolvedAtomCompleter(substitutor, topLevelCallContext)
         val declarationDescriptor = trace.get(BindingContext.DECLARATION_TO_DESCRIPTOR, expression)
 

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

@@ -167,7 +167,7 @@ class KotlinResolutionCallbacksImpl(
                     callComponents, builtIns, topLevelCallContext, stubsForPostponedVariables, trace,
                     kotlinToResolvedCallTransformer, expressionTypingServices, argumentTypeResolver,
                     doubleColonExpressionResolver, deprecationResolver, moduleDescriptor, typeApproximator,
-                    missingSupertypesResolver
+                    missingSupertypesResolver, lambdaArgument
                 )
             } else {
                 null

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

@@ -18,6 +18,11 @@ abstract class ExplicitTypeParameterConstraintPosition<T>(val typeArgument: T) :
     override fun toString(): String = "TypeParameter $typeArgument"
 }
 
+abstract class InjectedAnotherStubTypeConstraintPosition<T>(private val builderInferenceLambdaOfInjectedStubType: T) : ConstraintPosition(),
+    OnlyInputTypeConstraintPosition {
+    override fun toString(): String = "Injected from $builderInferenceLambdaOfInjectedStubType builder inference call"
+}
+
 abstract class ExpectedTypeConstraintPosition<T>(val topLevelCall: T) : ConstraintPosition(), OnlyInputTypeConstraintPosition {
     override fun toString(): String = "ExpectedType for call $topLevelCall"
 }

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

@@ -15,6 +15,9 @@ class ExplicitTypeParameterConstraintPositionImpl(
     typeArgument: SimpleTypeArgument
 ) : ExplicitTypeParameterConstraintPosition<SimpleTypeArgument>(typeArgument)
 
+class InjectedAnotherStubTypeConstraintPositionImpl(builderInferenceLambdaOfInjectedStubType: LambdaKotlinCallArgument) :
+    InjectedAnotherStubTypeConstraintPosition<LambdaKotlinCallArgument>(builderInferenceLambdaOfInjectedStubType)
+
 class ExpectedTypeConstraintPositionImpl(topLevelCall: KotlinCall) : ExpectedTypeConstraintPosition<KotlinCall>(topLevelCall)
 
 class DeclaredUpperBoundConstraintPositionImpl(