[NI] Fix construction of common system for builder inference

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

@@ -111,9 +111,11 @@ class KotlinToResolvedCallTransformer(
                 context.trace.record(BindingContext.ONLY_RESOLVED_CALL, psiCall, PartialCallContainer(baseResolvedCall))
                 context.trace.record(BindingContext.PARTIAL_CALL_RESOLUTION_CONTEXT, psiCall, context)
 
-                context.inferenceSession.addPartialCallInfo(
-                    PSIPartialCallInfo(baseResolvedCall, context, tracingStrategy)
-                )
+                if (baseResolvedCall.forwardToInferenceSession) {
+                    context.inferenceSession.addPartialCallInfo(
+                        PSIPartialCallInfo(baseResolvedCall, context, tracingStrategy),
+                    )
+                }
 
                 createStubResolvedCallAndWriteItToTrace(candidate, context.trace, baseResolvedCall.diagnostics, substitutor = null)
             }

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

@@ -51,14 +51,22 @@ class KotlinCallCompleter(
         candidate.addExpectedTypeFromCastConstraint(returnType, resolutionCallbacks)
         candidate.checkSamWithVararg(diagnosticHolder)
 
-        return if (resolutionCallbacks.inferenceSession.shouldRunCompletion(candidate))
-            candidate.runCompletion(
-                CompletionModeCalculator.computeCompletionMode(candidate, expectedType, returnType, trivialConstraintTypeInferenceOracle),
-                diagnosticHolder,
-                resolutionCallbacks
-            )
-        else
-            candidate.asCallResolutionResult(ConstraintSystemCompletionMode.PARTIAL, diagnosticHolder)
+        val completionMode =
+            CompletionModeCalculator.computeCompletionMode(candidate, expectedType, returnType, trivialConstraintTypeInferenceOracle)
+
+        return when (completionMode) {
+            ConstraintSystemCompletionMode.FULL -> {
+                if (resolutionCallbacks.inferenceSession.shouldRunCompletion(candidate)) {
+                    candidate.runCompletion(completionMode, diagnosticHolder, resolutionCallbacks)
+                } else {
+                    candidate.asCallResolutionResult(
+                        ConstraintSystemCompletionMode.PARTIAL, diagnosticHolder, forwardToInferenceSession = true
+                    )
+                }
+            }
+            ConstraintSystemCompletionMode.PARTIAL ->
+                candidate.runCompletion(completionMode, diagnosticHolder, resolutionCallbacks)
+        }
     }
 
     private fun KotlinResolutionCandidate.checkSamWithVararg(diagnosticHolder: KotlinDiagnosticsHolder.SimpleHolder) {
@@ -210,7 +218,8 @@ class KotlinCallCompleter(
 
     fun KotlinResolutionCandidate.asCallResolutionResult(
         type: ConstraintSystemCompletionMode,
-        diagnosticsHolder: KotlinDiagnosticsHolder.SimpleHolder
+        diagnosticsHolder: KotlinDiagnosticsHolder.SimpleHolder,
+        forwardToInferenceSession: Boolean = false
     ): CallResolutionResult {
         val systemStorage = getSystem().asReadOnlyStorage()
         val allDiagnostics = diagnosticsHolder.getDiagnostics() + this.diagnosticsFromResolutionParts
@@ -222,7 +231,7 @@ class KotlinCallCompleter(
         return if (type == ConstraintSystemCompletionMode.FULL) {
             CompletedCallResolutionResult(resolvedCall, allDiagnostics, systemStorage)
         } else {
-            PartialCallResolutionResult(resolvedCall, allDiagnostics, systemStorage)
+            PartialCallResolutionResult(resolvedCall, allDiagnostics, systemStorage, forwardToInferenceSession)
         }
     }
 

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

@@ -225,7 +225,8 @@ open class SingleCallResolutionResult(
 class PartialCallResolutionResult(
     resultCallAtom: ResolvedCallAtom,
     diagnostics: List<KotlinCallDiagnostic>,
-    constraintSystem: ConstraintStorage
+    constraintSystem: ConstraintStorage,
+    val forwardToInferenceSession: Boolean = false
 ) : SingleCallResolutionResult(resultCallAtom, diagnostics, constraintSystem)
 
 class CompletedCallResolutionResult(

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/builderInferenceForMaterializeWithExpectedType.kt

@@ -0,0 +1,11 @@
+// !LANGUAGE: +NewInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE
+
+fun foo(x: List<String>) {}
+fun <K> materialize(): K = TODO()
+
+fun test() {
+    val x: Sequence<String> = sequence {
+        foo(materialize())
+    }
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/builderInferenceForMaterializeWithExpectedType.txt

@@ -0,0 +1,5 @@
+package
+
+public fun foo(/*0*/ x: kotlin.collections.List<kotlin.String>): kotlin.Unit
+public fun </*0*/ K> materialize(): K
+public fun test(): kotlin.Unit

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/kt35684.kt

@@ -14,7 +14,7 @@ fun test_1() {
 
 fun test_2() {
     <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>sequence<!> {
-        yield(<!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>materialize<!>())
+        <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>yield<!>(<!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>materialize<!>())
     }
 }
 

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestWithStdLibGenerated.java

@@ -1882,6 +1882,11 @@ public class DiagnosticsTestWithStdLibGenerated extends AbstractDiagnosticsTestW
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/applyInsideCoroutine.kt");
             }
 
+            @TestMetadata("builderInferenceForMaterializeWithExpectedType.kt")
+            public void testBuilderInferenceForMaterializeWithExpectedType() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/builderInferenceForMaterializeWithExpectedType.kt");
+            }
+
             @TestMetadata("callableReferenceToASuspendFunction.kt")
             public void testCallableReferenceToASuspendFunction() throws Exception {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/callableReferenceToASuspendFunction.kt");

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsTestWithStdLibUsingJavacGenerated.java

@@ -1882,6 +1882,11 @@ public class DiagnosticsTestWithStdLibUsingJavacGenerated extends AbstractDiagno
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/applyInsideCoroutine.kt");
             }
 
+            @TestMetadata("builderInferenceForMaterializeWithExpectedType.kt")
+            public void testBuilderInferenceForMaterializeWithExpectedType() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/builderInferenceForMaterializeWithExpectedType.kt");
+            }
+
             @TestMetadata("callableReferenceToASuspendFunction.kt")
             public void testCallableReferenceToASuspendFunction() throws Exception {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/callableReferenceToASuspendFunction.kt");