[FE] Substitute fixed type variables with inferred stub types

Actually, a type variable might be fixed into a stub type. Such stub type should be substituted before sub calls completion

^KT-51988 Fixed

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

@@ -19480,6 +19480,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt51988.kt")
+            public void testKt51988() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");
+            }
+
             @Test
             @TestMetadata("labaledCall.kt")
             public void testLabaledCall() throws Exception {

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
@@ -108,12 +109,14 @@ class FirBuilderInferenceSession(
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage,
+        constraintSystemBuilder: ConstraintSystemBuilder,
         completionMode: ConstraintSystemCompletionMode
     ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? {
-        val (commonSystem, effectivelyEmptyConstraintSystem) = buildCommonSystem(initialStorage)
+        val (commonSystem, effectivelyEmptyConstraintSystem) = buildCommonSystem(constraintSystemBuilder.currentStorage())
+        val resultingSubstitutor by lazy { getResultingSubstitutor(commonSystem) }
+
         if (effectivelyEmptyConstraintSystem) {
-            updateCalls(commonSystem)
+            updateCalls(resultingSubstitutor)
             return null
         }
 
@@ -129,7 +132,11 @@ class FirBuilderInferenceSession(
             error("Shouldn't be called in complete constraint system mode")
         }
 
-        updateCalls(commonSystem)
+        if (completionMode == ConstraintSystemCompletionMode.FULL) {
+            constraintSystemBuilder.substituteFixedVariables(resultingSubstitutor)
+        }
+
+        updateCalls(resultingSubstitutor)
 
         @Suppress("UNCHECKED_CAST")
         return commonSystem.fixedTypeVariables as Map<ConeTypeVariableTypeConstructor, ConeKotlinType>
@@ -217,16 +224,18 @@ class FirBuilderInferenceSession(
         return introducedConstraint
     }
 
-    private fun updateCalls(commonSystem: NewConstraintSystemImpl) {
+    private fun getResultingSubstitutor(commonSystem: NewConstraintSystemImpl): ChainedSubstitutor {
         val nonFixedToVariablesSubstitutor = createNonFixedTypeToVariableSubstitutor()
         val commonSystemSubstitutor = commonSystem.buildCurrentSubstitutor() as ConeSubstitutor
-        val nonFixedTypesToResultSubstitutor = ChainedSubstitutor(nonFixedToVariablesSubstitutor, commonSystemSubstitutor)
+        return ChainedSubstitutor(nonFixedToVariablesSubstitutor, commonSystemSubstitutor)
+    }
 
-        val stubTypeSubstitutor = FirStubTypeTransformer(nonFixedTypesToResultSubstitutor)
+    private fun updateCalls(substitutor: ConeSubstitutor) {
+        val stubTypeSubstitutor = FirStubTypeTransformer(substitutor)
         lambda.transformSingle(stubTypeSubstitutor, null)
         // TODO: support diagnostics, see [CoroutineInferenceSession#updateCalls]
 
-        val completionResultsWriter = components.callCompleter.createCompletionResultsWriter(nonFixedTypesToResultSubstitutor)
+        val completionResultsWriter = components.callCompleter.createCompletionResultsWriter(substitutor)
         for ((call, _) in partiallyResolvedCalls) {
             call.transformSingle(completionResultsWriter, null)
             // TODO: support diagnostics, see [CoroutineInferenceSession#updateCalls]

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.resolve.substitution.ChainedSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.NotFixedTypeToVariableSubstitutorForDelegateInference
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionContext
@@ -108,7 +109,7 @@ class FirDelegatedPropertyInferenceSession(
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage,
+        constraintSystemBuilder: ConstraintSystemBuilder,
         completionMode: ConstraintSystemCompletionMode
     ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? = null
 

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeStubType
 import org.jetbrains.kotlin.fir.types.ConeTypeVariableTypeConstructor
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionContext
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
@@ -41,7 +42,7 @@ abstract class FirInferenceSession {
 
     abstract fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage,
+        constraintSystemBuilder: ConstraintSystemBuilder,
         completionMode: ConstraintSystemCompletionMode,
         // TODO: diagnostic holder
     ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>?
@@ -61,7 +62,7 @@ abstract class FirStubInferenceSession : FirInferenceSession() {
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage,
+        constraintSystemBuilder: ConstraintSystemBuilder,
         completionMode: ConstraintSystemCompletionMode
     ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? = null
 

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

@@ -209,9 +209,8 @@ class PostponedArgumentsAnalyzer(
         lambda.returnStatements = returnArguments
 
         if (inferenceSession != null) {
-            val storageSnapshot = c.getBuilder().currentStorage()
+            val constraintSystemBuilder = c.getBuilder()
-
+            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, constraintSystemBuilder, completionMode)
-            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot, completionMode)
 
             if (postponedVariables == null) {
                 c.getBuilder().removePostponedVariables()
@@ -219,7 +218,7 @@ class PostponedArgumentsAnalyzer(
             }
 
             for ((constructor, resultType) in postponedVariables) {
-                val variableWithConstraints = storageSnapshot.notFixedTypeVariables[constructor] ?: continue
+                val variableWithConstraints = constraintSystemBuilder.currentStorage().notFixedTypeVariables[constructor] ?: continue
                 val variable = variableWithConstraints.typeVariable as ConeTypeVariable
 
                 c.getBuilder().unmarkPostponedVariable(variable)

compiler/frontend/src/org/jetbrains/kotlin/resolve/DelegateInferenceSession.kt

@@ -97,7 +97,7 @@ class DelegateInferenceSession(
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage,
+        constraintSystemBuilder: ConstraintSystemBuilder,
         completionMode: ConstraintSystemCompletionMode,
         diagnosticsHolder: KotlinDiagnosticsHolder
     ): Map<TypeConstructor, UnwrappedType> = emptyMap()
@@ -124,7 +124,7 @@ class InferenceSessionForExistingCandidates(
     override fun currentConstraintSystem(): ConstraintStorage = ConstraintStorage.Empty
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage,
+        constraintSystemBuilder: ConstraintSystemBuilder,
         completionMode: ConstraintSystemCompletionMode,
         diagnosticsHolder: KotlinDiagnosticsHolder
     ): Map<TypeConstructor, UnwrappedType> = emptyMap()

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

@@ -233,22 +233,23 @@ class BuilderInferenceSession(
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage,
+        constraintSystemBuilder: ConstraintSystemBuilder,
         completionMode: ConstraintSystemCompletionMode,
         diagnosticsHolder: KotlinDiagnosticsHolder,
     ): Map<TypeConstructor, UnwrappedType>? {
         initializeLambda(lambda)
 
-        fun getResultingSubstitutor(): NewTypeSubstitutor {
+        val initialStorage = constraintSystemBuilder.currentStorage()
+        val resultingSubstitutor by lazy {
             val storageSubstitutor = initialStorage.buildResultingSubstitutor(commonSystem, transformTypeVariablesToErrorTypes = false)
-            return ComposedSubstitutor(storageSubstitutor, commonSystem.buildCurrentSubstitutor() as NewTypeSubstitutor)
+            ComposedSubstitutor(storageSubstitutor, commonSystem.buildCurrentSubstitutor() as NewTypeSubstitutor)
         }
 
         val effectivelyEmptyConstraintSystem = initializeCommonSystem(initialStorage)
 
         if (effectivelyEmptyConstraintSystem) {
             if (isTopLevelBuilderInferenceCall()) {
-                updateAllCalls(getResultingSubstitutor())
+                updateAllCalls(resultingSubstitutor)
             }
             return null
         }
@@ -261,8 +262,14 @@ class BuilderInferenceSession(
             diagnosticsHolder
         )
 
+        if (completionMode == ConstraintSystemCompletionMode.FULL) {
+            constraintSystemBuilder.substituteFixedVariables(
+                ComposedSubstitutor(resultingSubstitutor, createNonFixedTypeToVariableSubstitutor())
+            )
+        }
+
         if (isTopLevelBuilderInferenceCall()) {
-            updateAllCalls(getResultingSubstitutor())
+            updateAllCalls(resultingSubstitutor)
         }
 
         return commonSystem.fixedTypeVariables.cast() // TODO: SUB

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

@@ -18,6 +18,7 @@ interface ConstraintSystemOperation {
     fun markCouldBeResolvedWithUnrestrictedBuilderInference()
     fun unmarkPostponedVariable(variable: TypeVariableMarker)
     fun removePostponedVariables()
+    fun substituteFixedVariables(substitutor: TypeSubstitutorMarker)
 
     fun getBuiltFunctionalExpectedTypeForPostponedArgument(
         topLevelVariable: TypeConstructorMarker,

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

@@ -144,6 +144,10 @@ class NewConstraintSystemImpl(
         storage.postponedTypeVariables.clear()
     }
 
+    override fun substituteFixedVariables(substitutor: TypeSubstitutorMarker) {
+        storage.fixedTypeVariables.replaceAll { _, type -> substitutor.safeSubstitute(type) }
+    }
+
     override fun putBuiltFunctionalExpectedTypeForPostponedArgument(
         topLevelVariable: TypeConstructorMarker,
         pathToExpectedType: List<Pair<TypeConstructorMarker, Int>>,

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.resolve.calls.components
 
 import org.jetbrains.kotlin.resolve.calls.components.candidate.ResolutionCandidate
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.model.*
@@ -26,7 +27,7 @@ interface InferenceSession {
             override fun currentConstraintSystem(): ConstraintStorage = ConstraintStorage.Empty
             override fun inferPostponedVariables(
                 lambda: ResolvedLambdaAtom,
-                initialStorage: ConstraintStorage,
+                constraintSystemBuilder: ConstraintSystemBuilder,
                 completionMode: ConstraintSystemCompletionMode,
                 diagnosticsHolder: KotlinDiagnosticsHolder
             ): Map<TypeConstructor, UnwrappedType> = emptyMap()
@@ -51,7 +52,7 @@ interface InferenceSession {
     fun currentConstraintSystem(): ConstraintStorage
     fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage,
+        constraintSystemBuilder: ConstraintSystemBuilder,
         completionMode: ConstraintSystemCompletionMode,
         diagnosticsHolder: KotlinDiagnosticsHolder
     ): Map<TypeConstructor, UnwrappedType>?

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

@@ -191,11 +191,11 @@ class PostponedArgumentsAnalyzer(
                 || languageVersionSettings.supportsFeature(LanguageFeature.UseBuilderInferenceWithoutAnnotation)
 
         if (inferenceSession != null && shouldUseBuilderInference) {
-            val storageSnapshot = c.getBuilder().currentStorage()
+            val constraintSystemBuilder = c.getBuilder()
 
             val postponedVariables = inferenceSession.inferPostponedVariables(
                 lambda,
-                storageSnapshot,
+                constraintSystemBuilder,
                 completionMode,
                 diagnosticHolder
             )
@@ -205,7 +205,7 @@ class PostponedArgumentsAnalyzer(
             }
 
             for ((constructor, resultType) in postponedVariables) {
-                val variableWithConstraints = storageSnapshot.notFixedTypeVariables[constructor] ?: continue
+                val variableWithConstraints = constraintSystemBuilder.currentStorage().notFixedTypeVariables[constructor] ?: continue
                 val variable = variableWithConstraints.typeVariable
 
                 c.getBuilder().unmarkPostponedVariable(variable)

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

@@ -0,0 +1,37 @@
+// WITH_STDLIB
+// IGNORE_BACKEND_FIR: JVM_IR
+
+import kotlin.experimental.ExperimentalTypeInference
+
+@OptIn(ExperimentalTypeInference::class)
+fun <AB1, BB1> build(@BuilderInference block: BuilderScope<AB1>.() -> BB1): ResultProvider<AB1, BB1> = object : ResultProvider<AB1, BB1> {
+    override fun provideResult(): AB1 = "OK" as AB1
+}
+
+@OptIn(ExperimentalTypeInference::class)
+fun <AB2, BB2> build2(@BuilderInference block: BuilderScope<AB2>.() -> BB2): ResultProvider<AB2, BB2> = object : ResultProvider<AB2, BB2> {
+    override fun provideResult(): AB2 = "OK" as AB2
+}
+
+interface BuilderScope<BS> {
+    fun <B1> getResult(result: ResultProvider<BS, B1>): B1
+    fun <B2> getResult2(result: ResultProvider<BS, B2>): B2
+}
+
+interface ResultProvider<AR, BR> {
+    fun provideResult(): AR
+}
+
+val resultProvider: ResultProvider<Any, Any> = object : ResultProvider<Any, Any> {
+    override fun provideResult(): Any = "NOK"
+}
+
+val result = build {
+    getResult(build2 {
+        getResult2(resultProvider)
+    })
+}
+
+fun box(): String {
+    return result.provideResult().toString()
+}

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

@@ -19048,6 +19048,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt51988.kt")
+            public void testKt51988() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");
+            }
+
             @Test
             @TestMetadata("labaledCall.kt")
             public void testLabaledCall() throws Exception {

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

@@ -19480,6 +19480,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt51988.kt")
+            public void testKt51988() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");
+            }
+
             @Test
             @TestMetadata("labaledCall.kt")
             public void testLabaledCall() throws Exception {

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

@@ -15843,6 +15843,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @TestMetadata("kt51988.kt")
+            public void testKt51988() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");
+            }
+
             @TestMetadata("labaledCall.kt")
             public void testLabaledCall() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/labaledCall.kt");

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

@@ -14826,6 +14826,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt51988.kt")
+            public void testKt51988() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");
+            }
+
             @Test
             @TestMetadata("labaledCall.kt")
             public void testLabaledCall() throws Exception {

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

@@ -14790,6 +14790,12 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @Test
+            @TestMetadata("kt51988.kt")
+            public void testKt51988() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");
+            }
+
             @Test
             @TestMetadata("labaledCall.kt")
             public void testLabaledCall() throws Exception {

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

@@ -13158,6 +13158,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
                 runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
             }
 
+            @TestMetadata("kt51988.kt")
+            public void testKt51988() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");
+            }
+
             @TestMetadata("labaledCall.kt")
             public void testLabaledCall() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/labaledCall.kt");

native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeCodegenBoxTestGenerated.java

@@ -15900,6 +15900,12 @@ public class NativeCodegenBoxTestGenerated extends AbstractNativeCodegenBoxTest
                     runTest("compiler/testData/codegen/box/inference/builderInference/kt49887.kt");
                 }
 
+                @Test
+                @TestMetadata("kt51988.kt")
+                public void testKt51988() throws Exception {
+                    runTest("compiler/testData/codegen/box/inference/builderInference/kt51988.kt");
+                }
+
                 @Test
                 @TestMetadata("labaledCall.kt")
                 public void testLabaledCall() throws Exception {