FIR: Fix leakage of type variables in builder inference for lambdas returning Unit

When BI/incomplete call is present in return argument, it will be
added to the main call-tree, leading to requirement violation in
`ConstraintSystemCompleter.getOrderedAllTypeVariables`
as after `FirBuilderInferenceSession.inferPostponedVariables` it will
be completed, and its type variables couldn't be found anymore

In K1, we actually do the same, but we are able to find type variables
of such calls due to architecture difference:
In K2, `getOrderedAllTypeVariables` uses FIR as the main source to
lookup
In K1, `getOrderedAllTypeVariables` uses resolution atom tree, where
candidate/CS of the call is still available after
`inferPostponedVariables`

In fact, all incomplete calls were analyzed in FULL mode according to
the contract of `FirBuilderInferenceSession.addPartiallyResolvedCall`.
Thus, it means we normally shouldn't add them to the main call-tree, but
accidentally do it as incomplete calls contain non-completed candidate

This particular commit addresses the problem partially, only
in cases when the expected return type for the lambda is Unit and when
the incomplete call is located in the last expression of the lambda

In such cases, we can skip the call from the last expression completely,
since all potential calls there were analyzed in FULL mode,
and couldn't introduce any useful info to the CS of the main call-tree

^KT-54294

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

@@ -15191,6 +15191,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/errorOnStubReceiver.kt");
                 }
 
+                @Test
+                @TestMetadata("incompleteCallInReturnArgumentsWithProperExpectType.kt")
+                public void testIncompleteCallInReturnArgumentsWithProperExpectType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/incompleteCallInReturnArgumentsWithProperExpectType.kt");
+                }
+
                 @Test
                 @TestMetadata("invalidateKeys.kt")
                 public void testInvalidateKeys() throws Exception {

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

@@ -15191,6 +15191,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/errorOnStubReceiver.kt");
                 }
 
+                @Test
+                @TestMetadata("incompleteCallInReturnArgumentsWithProperExpectType.kt")
+                public void testIncompleteCallInReturnArgumentsWithProperExpectType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/incompleteCallInReturnArgumentsWithProperExpectType.kt");
+                }
+
                 @Test
                 @TestMetadata("invalidateKeys.kt")
                 public void testInvalidateKeys() throws Exception {

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

@@ -15191,6 +15191,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/errorOnStubReceiver.kt");
                 }
 
+                @Test
+                @TestMetadata("incompleteCallInReturnArgumentsWithProperExpectType.kt")
+                public void testIncompleteCallInReturnArgumentsWithProperExpectType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/incompleteCallInReturnArgumentsWithProperExpectType.kt");
+                }
+
                 @Test
                 @TestMetadata("invalidateKeys.kt")
                 public void testInvalidateKeys() throws Exception {

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

@@ -15197,6 +15197,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/errorOnStubReceiver.kt");
                 }
 
+                @Test
+                @TestMetadata("incompleteCallInReturnArgumentsWithProperExpectType.kt")
+                public void testIncompleteCallInReturnArgumentsWithProperExpectType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/incompleteCallInReturnArgumentsWithProperExpectType.kt");
+                }
+
                 @Test
                 @TestMetadata("invalidateKeys.kt")
                 public void testInvalidateKeys() throws Exception {

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

@@ -20704,6 +20704,24 @@ public class FirLightTreeBlackBoxCodegenTestGenerated extends AbstractFirLightTr
                 runTest("compiler/testData/codegen/box/inference/builderInference/nullability.kt");
             }
 
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgument.kt")
+            public void testPartiallyResolvedCallInReturnArgument() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgument.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonLast.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonLast() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonLast.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonUnit.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonUnit() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonUnit.kt");
+            }
+
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {

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

@@ -20704,6 +20704,24 @@ public class FirPsiBlackBoxCodegenTestGenerated extends AbstractFirPsiBlackBoxCo
                 runTest("compiler/testData/codegen/box/inference/builderInference/nullability.kt");
             }
 
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgument.kt")
+            public void testPartiallyResolvedCallInReturnArgument() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgument.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonLast.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonLast() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonLast.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonUnit.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonUnit() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonUnit.kt");
+            }
+
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {

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

@@ -174,9 +174,16 @@ class PostponedArgumentsAnalyzer(
         val returnTypeRef = lambda.atom.returnTypeRef.let {
             it as? FirResolvedTypeRef ?: it.resolvedTypeFromPrototype(substitute(lambda.returnType))
         }
+        val lambdaExpectedTypeIsUnit = returnTypeRef.type.isUnitOrFlexibleUnit
         returnArguments.forEach {
-            val haveSubsystem = c.addSubsystemFromExpression(it)
             // If the lambda returns Unit, the last expression is not returned and should not be constrained.
+            val isLastExpression = it == lastExpression
+
+            // Don't add last call for builder-inference
+            // Note, that we don't use the same condition "(returnTypeRef.type.isUnitOrFlexibleUnit || lambda.atom.shouldReturnUnit(returnArguments))"
+            // as in the if below, because "lambda.atom.shouldReturnUnit(returnArguments)" might mean that the last statement is not completed
+            if (isLastExpression && lambdaExpectedTypeIsUnit && inferenceSession is FirBuilderInferenceSession) return@forEach
+
             // TODO (KT-55837) questionable moment inherited from FE1.0 (the `haveSubsystem` case):
             //    fun <T> foo(): T
             //    run {
@@ -184,8 +191,10 @@ class PostponedArgumentsAnalyzer(
             //      foo() // T = Unit, even though there is no implicit return
             //    }
             //  Things get even weirder if T has an upper bound incompatible with Unit.
-            if (it == lastExpression && !haveSubsystem &&
-                (returnTypeRef.type.isUnitOrFlexibleUnit || lambda.atom.shouldReturnUnit(returnArguments))
+            // Not calling `addSubsystemFromExpression` for builder-inference is crucial
+            val haveSubsystem = c.addSubsystemFromExpression(it)
+            if (isLastExpression && !haveSubsystem &&
+                (lambdaExpectedTypeIsUnit || lambda.atom.shouldReturnUnit(returnArguments))
             ) return@forEach
 
             hasExpressionInReturnArguments = true
@@ -202,7 +211,7 @@ class PostponedArgumentsAnalyzer(
             }
         }
 
-        if (!hasExpressionInReturnArguments && !returnTypeRef.type.isUnitOrFlexibleUnit) {
+        if (!hasExpressionInReturnArguments && !lambdaExpectedTypeIsUnit) {
             builder.addSubtypeConstraint(
                 components.session.builtinTypes.unitType.type,
                 returnTypeRef.type,

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

@@ -0,0 +1,34 @@
+// WITH_STDLIB
+// TARGET_BACKEND: JVM
+class Out<out V>(val v: V)
+class Box<R> {
+    var boxed: R? = null
+
+    fun set(newValue: R) {
+        boxed = newValue
+    }
+}
+
+fun <R> buildBox(fill: Box<R>.() -> Unit): Box<R> {
+    return Box<R>().also {
+        it.fill()
+    }
+}
+
+fun foo() =
+    buildBox {
+        set(select(Out("OK"), makeOut())) // problem is here
+        // we keep set in PARTIALLY analyzed state, while we're also trying to incorporate it both into return-type constraining
+        // and builder inference, which introduces non-trivial loop of codependency between those two steps
+        // Due to that, we leak type variables from set call-graph to the system of buildBox, while those variables are already
+        // fixed
+    }
+
+fun <S> select(a: S, b: S): S = a
+
+fun <M> makeOut(): Out<M>? = null
+
+fun box(): String {
+    if (foo().boxed!!.v != "OK") return "FAIL"
+    return "OK"
+}

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

@@ -0,0 +1,34 @@
+// WITH_STDLIB
+// TARGET_BACKEND: JVM
+// IGNORE_BACKEND_K2: JVM_IR
+// FIR status: KT-58742
+
+class Out<out V>(val v: V)
+class Box<R> {
+    var boxed: R? = null
+
+    fun set(newValue: R) {
+        boxed = newValue
+    }
+}
+
+fun <R> buildBox(fill: Box<R>.() -> Unit): Box<R> {
+    return Box<R>().also {
+        it.fill()
+    }
+}
+
+fun foo() =
+    buildBox {
+        if (true) return@buildBox set(select(Out("OK"), makeOut())) // problem is here
+        Unit
+    }
+
+fun <S> select(a: S, b: S): S = a
+
+fun <M> makeOut(): Out<M>? = null
+
+fun box(): String {
+    if (foo().boxed!!.v != "OK") return "FAIL"
+    return "OK"
+}

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

@@ -0,0 +1,36 @@
+// WITH_STDLIB
+// TARGET_BACKEND: JVM
+// IGNORE_BACKEND_K2: JVM_IR
+// FIR status: KT-58741
+
+fun foo() =
+    buildBox {
+        set(select(Out("OK"), makeOut()))
+    }
+
+
+class Out<out V>(val v: V)
+class Box<R> {
+    var boxed: R? = null
+
+    fun set(newValue: R): String {
+        boxed = newValue
+        return "set"
+    }
+}
+
+fun <R> buildBox(fill: Box<R>.() -> String): Box<R> {
+    return Box<R>().also {
+        require(it.fill() == "set")
+    }
+}
+
+fun <S> select(a: S, b: S): S = a
+
+fun <M> makeOut(): Out<M>? = null
+
+
+fun box(): String {
+    if (foo().boxed!!.v != "OK") return "FAIL"
+    return "OK"
+}

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

@@ -0,0 +1,43 @@
+// FIR_IDENTICAL
+// WITH_STDLIB
+// SKIP_TXT
+
+fun test1() =
+    buildBoxUnit {
+        set(select(Out(""), makeOut()))
+        mat() // Should be able to infer Unit here
+    }
+
+fun test2() =
+    buildBoxProperType {
+        set(select(Out(""), makeOut()))
+        mat() // Should be able to infer String here
+    }
+
+
+class Out<out V>(val v: V)
+class Box<R> {
+    var boxed: R? = null
+
+    fun set(newValue: R) {
+        boxed = newValue
+    }
+}
+
+fun <R> buildBoxUnit(fill: Box<R>.() -> Unit): Box<R> {
+    return Box<R>().also {
+        it.fill()
+    }
+}
+
+fun <R> buildBoxProperType(fill: Box<R>.() -> String): Box<R> {
+    return Box<R>().also {
+        it.fill()
+    }
+}
+
+fun <S> select(a: S, b: S): S = a
+
+fun <M> makeOut(): Out<M>? = null
+
+fun <M> mat(): M = null!!

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

@@ -15197,6 +15197,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/inference/builderInference/errorOnStubReceiver.kt");
                 }
 
+                @Test
+                @TestMetadata("incompleteCallInReturnArgumentsWithProperExpectType.kt")
+                public void testIncompleteCallInReturnArgumentsWithProperExpectType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/builderInference/incompleteCallInReturnArgumentsWithProperExpectType.kt");
+                }
+
                 @Test
                 @TestMetadata("invalidateKeys.kt")
                 public void testInvalidateKeys() throws Exception {

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

@@ -19780,6 +19780,24 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/inference/builderInference/nullability.kt");
             }
 
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgument.kt")
+            public void testPartiallyResolvedCallInReturnArgument() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgument.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonLast.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonLast() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonLast.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonUnit.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonUnit() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonUnit.kt");
+            }
+
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {

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

@@ -20704,6 +20704,24 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/nullability.kt");
             }
 
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgument.kt")
+            public void testPartiallyResolvedCallInReturnArgument() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgument.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonLast.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonLast() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonLast.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonUnit.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonUnit() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonUnit.kt");
+            }
+
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {

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

@@ -20704,6 +20704,24 @@ public class IrBlackBoxCodegenWithIrInlinerTestGenerated extends AbstractIrBlack
                 runTest("compiler/testData/codegen/box/inference/builderInference/nullability.kt");
             }
 
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgument.kt")
+            public void testPartiallyResolvedCallInReturnArgument() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgument.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonLast.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonLast() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonLast.kt");
+            }
+
+            @Test
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonUnit.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonUnit() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonUnit.kt");
+            }
+
             @Test
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {

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

@@ -16484,6 +16484,21 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/nullability.kt");
             }
 
+            @TestMetadata("partiallyResolvedCallInReturnArgument.kt")
+            public void testPartiallyResolvedCallInReturnArgument() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgument.kt");
+            }
+
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonLast.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonLast() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonLast.kt");
+            }
+
+            @TestMetadata("partiallyResolvedCallInReturnArgumentNonUnit.kt")
+            public void testPartiallyResolvedCallInReturnArgumentNonUnit() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/partiallyResolvedCallInReturnArgumentNonUnit.kt");
+            }
+
             @TestMetadata("propagateInferenceSessionIntoDeclarationAnalyzers.kt")
             public void testPropagateInferenceSessionIntoDeclarationAnalyzers() throws Exception {
                 runTest("compiler/testData/codegen/box/inference/builderInference/propagateInferenceSessionIntoDeclarationAnalyzers.kt");

docs/fir/builder_inference.md

@@ -99,6 +99,10 @@ See `org.jetbrains.kotlin.fir.resolve.inference.PostponedArgumentsAnalyzer.apply
 
 Once the lambda body was analyzed return arguments are added into the call-tree
 
+> ##### Note: Incomplete calls in return arguments
+> We don't add last expression of lambda to the call-tree as a return argument if its functional-type has Unit return-type, to
+> avoid situations when last expression contains incomplete call, see [Incomplete call in return arguments](#incomplete-call-in-return-arguments)
+
 Then, we perform inference of postponed type variables
 
 See `org.jetbrains.kotlin.fir.resolve.inference.FirBuilderInferenceSession.inferPostponedVariables`
@@ -207,4 +211,14 @@ type variable from the call-tree and solution from integration CS, loosing infor
 It causes unsound solutions in the call-tree CS
 #### Resulting substitution is unclear
 Its unclear why we mix stub type substitutor with type-variable-type substitutor, and why we need to manually handle substitution to error
-types
+types
+#### Incomplete call in return arguments
+If incomplete call is present among return arguments during [Lambda analysis finalization](#lambda-analysis-finalization) we add it to the 
+main call-tree, but then complete it as part of [result writing](#result-write).
+
+It leads to violation of contract in `org.jetbrains.kotlin.fir.resolve.inference.ConstraintSystemCompleter.getOrderedAllTypeVariables` as type variables for completed call couldn't be found anymore
+
+To avoid such problem, we have [workaround for lambdas with Unit return-type](#Note-Incomplete-calls-in-return-arguments), but problem still 
+occurs:
+- When we have [lambdas with non-Unit return-type](https://youtrack.jetbrains.com/issue/KT-58741)
+- When the [incomplete call is present in a return argument that isn't last expression](https://youtrack.jetbrains.com/issue/KT-58742)