FIR: Fix inference of builder-inference function from expect type

Previously, such calls were being completed with FULL mode and incorrect
INFERENCE_NO_INFORMATION_FOR_PARAMETER has been reported

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/KtFirCompositeScope.fir.txt

@@ -20,7 +20,7 @@ FILE: KtFirCompositeScope.kt
         public final fun getAllNames(): R|kotlin/collections/Set<kotlin/String>| {
             ^getAllNames R|/withValidityAssertion|<R|kotlin/collections/Set<kotlin/String>|>(<L> = withValidityAssertion@fun <anonymous>(): R|kotlin/collections/Set<kotlin/String>| <inline=Inline, kind=UNKNOWN>  {
                 ^ R|/buildSet|<R|kotlin/String|>(<L> = buildSet@fun R|kotlin/collections/MutableSet<kotlin/String>|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=UNKNOWN>  {
-                    this@R|/KtFirCompositeScope|.R|/KtFirCompositeScope.subScopes|.R|kotlin/collections/flatMapTo|<R|KtScope|, R|kotlin/String|, R|kotlin/collections/MutableSet<kotlin/String>|>(this@R|special/anonymous|, <L> = flatMapTo@fun <anonymous>(it: R|KtScope|): R|kotlin/collections/Iterable<kotlin/String>| <inline=Inline, kind=UNKNOWN>  {
+                    this@R|/KtFirCompositeScope|.R|/KtFirCompositeScope.subScopes|.R|kotlin/collections/flatMapTo|<R|KtScope|, R|kotlin/String|, R|kotlin/collections/MutableSet<Stub: TypeVariable(E)>|>(this@R|special/anonymous|, <L> = flatMapTo@fun <anonymous>(it: R|KtScope|): R|kotlin/collections/Iterable<kotlin/String>| <inline=Inline, kind=UNKNOWN>  {
                         ^ R|<local>/it|.R|/KtScope.getAllNames|()
                     }
                     )

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

@@ -17677,6 +17677,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/inference/builderInference/constraintsBetweenTwoStubVariables.kt");
             }
 
+            @Test
+            @TestMetadata("inferFromExpectedType.kt")
+            public void testInferFromExpectedType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/inferFromExpectedType.kt");
+            }
+
             @Test
             @TestMetadata("intersect.kt")
             public void testIntersect() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirOverloadByLambdaReturnTypeResolver.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir
 
-import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
@@ -138,7 +137,8 @@ class FirOverloadByLambdaReturnTypeResolver(
             val results = postponedArgumentsAnalyzer.analyzeLambda(
                 firstCandidate.system.asPostponedArgumentsAnalyzerContext(),
                 firstAtom,
-                firstCandidate
+                firstCandidate,
+                ConstraintSystemCompletionMode.FULL,
             )
             while (iterator.hasNext()) {
                 val (candidate, atom) = iterator.next()
@@ -147,7 +147,8 @@ class FirOverloadByLambdaReturnTypeResolver(
                     candidate.system.asPostponedArgumentsAnalyzerContext(),
                     atom,
                     candidate,
-                    results
+                    results,
+                    ConstraintSystemCompletionMode.FULL,
                 )
             }
 

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

@@ -108,7 +108,8 @@ class FirBuilderInferenceSession(
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage
+        initialStorage: ConstraintStorage,
+        completionMode: ConstraintSystemCompletionMode
     ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? {
         val (commonSystem, effectivelyEmptyConstraintSystem) = buildCommonSystem(initialStorage)
         if (effectivelyEmptyConstraintSystem) {
@@ -120,7 +121,7 @@ class FirBuilderInferenceSession(
         @Suppress("UNCHECKED_CAST")
         components.callCompleter.completer.complete(
             context,
-            ConstraintSystemCompletionMode.FULL,
+            completionMode,
             partiallyResolvedCalls.map { it.first as FirStatement },
             components.session.builtinTypes.unitType.type, resolutionContext,
             collectVariablesFromContext = true

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

@@ -167,7 +167,7 @@ class FirCallCompleter(
             initialType,
             transformer.resolutionContext
         ) {
-            analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it, candidate)
+            analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it, candidate, completionMode)
         }
     }
 

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

@@ -53,7 +53,8 @@ class FirDelegatedPropertyInferenceSession(
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage
+        initialStorage: ConstraintStorage,
+        completionMode: ConstraintSystemCompletionMode
     ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? = null
 
     override fun <T> shouldCompleteResolvedSubAtomsOf(call: T): Boolean where T : FirResolvable, T : FirStatement = true
@@ -76,7 +77,8 @@ class FirDelegatedPropertyInferenceSession(
                 postponedArgumentsAnalyzer.analyze(
                     commonSystem.asPostponedArgumentsAnalyzerContext(),
                     it,
-                    resolvedCalls.first().candidate
+                    resolvedCalls.first().candidate,
+                    ConstraintSystemCompletionMode.FULL,
                 )
             }
         }

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeTypeVariableTypeConstructor
+import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 
 abstract class FirInferenceSession {
@@ -27,6 +28,7 @@ abstract class FirInferenceSession {
     abstract fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
         initialStorage: ConstraintStorage,
+        completionMode: ConstraintSystemCompletionMode,
         // TODO: diagnostic holder
     ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>?
 
@@ -47,7 +49,8 @@ abstract class FirStubInferenceSession : FirInferenceSession() {
 
     override fun inferPostponedVariables(
         lambda: ResolvedLambdaAtom,
-        initialStorage: ConstraintStorage
+        initialStorage: ConstraintStorage,
+        completionMode: ConstraintSystemCompletionMode
     ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType>? = null
 
     override fun <T> writeOnlyStubs(call: T): Boolean where T : FirResolvable, T : FirStatement = false

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

@@ -16,11 +16,15 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedReferenceError
 import org.jetbrains.kotlin.fir.resolve.inference.model.ConeLambdaArgumentConstraintPosition
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.StoreNameReference
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeVariable
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.isMarkedNullable
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzerContext
 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.CoroutinePosition
 import org.jetbrains.kotlin.types.model.StubTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
@@ -52,14 +56,15 @@ class PostponedArgumentsAnalyzer(
     fun analyze(
         c: PostponedArgumentsAnalyzerContext,
         argument: PostponedResolvedAtom,
-        candidate: Candidate
+        candidate: Candidate,
+        completionMode: ConstraintSystemCompletionMode
     ) {
         when (argument) {
             is ResolvedLambdaAtom ->
-                analyzeLambda(c, argument, candidate)
+                analyzeLambda(c, argument, candidate, completionMode)
 
             is LambdaWithTypeVariableAsExpectedTypeAtom ->
-                analyzeLambda(c, argument.transformToResolvedLambda(c.getBuilder(), resolutionContext), candidate)
+                analyzeLambda(c, argument.transformToResolvedLambda(c.getBuilder(), resolutionContext), candidate, completionMode)
 
             is ResolvedCallableReferenceAtom -> processCallableReference(argument, candidate)
 
@@ -106,7 +111,8 @@ class PostponedArgumentsAnalyzer(
     fun analyzeLambda(
         c: PostponedArgumentsAnalyzerContext,
         lambda: ResolvedLambdaAtom,
-        candidate: Candidate
+        candidate: Candidate,
+        completionMode: ConstraintSystemCompletionMode,
         //diagnosticHolder: KotlinDiagnosticsHolder
     ): ReturnArgumentsAnalysisResult {
         val unitType = components.session.builtinTypes.unitType.type
@@ -135,7 +141,15 @@ class PostponedArgumentsAnalyzer(
             expectedTypeForReturnArguments,
             stubsForPostponedVariables
         )
-        applyResultsOfAnalyzedLambdaToCandidateSystem(c, lambda, candidate, results, expectedTypeForReturnArguments, ::substitute)
+        applyResultsOfAnalyzedLambdaToCandidateSystem(
+            c,
+            lambda,
+            candidate,
+            results,
+            completionMode,
+            expectedTypeForReturnArguments,
+            ::substitute
+        )
         return results
     }
 
@@ -144,6 +158,7 @@ class PostponedArgumentsAnalyzer(
         lambda: ResolvedLambdaAtom,
         candidate: Candidate,
         results: ReturnArgumentsAnalysisResult,
+        completionMode: ConstraintSystemCompletionMode,
         expectedReturnType: ConeKotlinType? = null,
         substitute: (ConeKotlinType) -> ConeKotlinType = c.createSubstituteFunctorForLambdaAnalysis()
     ) {
@@ -151,7 +166,7 @@ class PostponedArgumentsAnalyzer(
         if (inferenceSession != null) {
             val storageSnapshot = c.getBuilder().currentStorage()
 
-            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot)
+            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot, completionMode)
 
             if (postponedVariables == null) {
                 c.getBuilder().removePostponedVariables()

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

@@ -0,0 +1,17 @@
+// WITH_RUNTIME
+// TARGET_BACKEND: JVM_IR
+
+@OptIn(ExperimentalStdlibApi::class)
+fun getAllPossibleNames(subScopes: List<List<String>>): Set<String> = withValidityAssertion {
+    buildSet {
+        subScopes.flatMapTo(this) { it }
+    }
+}
+
+inline fun <R> withValidityAssertion(action: () -> R): R {
+    return action()
+}
+
+fun box(): String {
+    return getAllPossibleNames(listOf(listOf("O"), listOf("K"))).joinToString("")
+}

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

@@ -1,6 +1,5 @@
 // ISSUE: KT-41164
 // WITH_RUNTIME
-// IGNORE_BACKEND_FIR: JVM_IR
 // DONT_TARGET_EXACT_BACKEND: WASM
 
 import kotlin.experimental.ExperimentalTypeInference

compiler/testData/diagnostics/testsWithStdLib/builderInference/resolveUsualCallWithBuilderInference.fir.kt

@@ -1,37 +0,0 @@
-// !LANGUAGE: +UnrestrictedBuilderInference
-// !DIAGNOSTICS: -UNUSED_PARAMETER
-// FILE: annotation.kt
-
-package kotlin
-
-annotation class BuilderInference
-
-// FILE: test.kt
-
-class Builder<T> {
-    fun add(t: T) {}
-}
-
-fun <S> build(@BuilderInference g: Builder<S>.() -> Unit): List<S> = TODO()
-fun <S> wrongBuild(g: Builder<S>.() -> Unit): List<S> = TODO()
-
-fun <S> Builder<S>.extensionAdd(s: S) {}
-
-@BuilderInference
-fun <S> Builder<S>.safeExtensionAdd(s: S) {}
-
-val member = build {
-    add(42)
-}
-
-val memberWithoutAnn = wrongBuild {
-    add(<!ARGUMENT_TYPE_MISMATCH!>42<!>)
-}
-
-val extension = build {
-    extensionAdd("foo")
-}
-
-val safeExtension = build {
-    safeExtensionAdd("foo")
-}

compiler/testData/diagnostics/testsWithStdLib/builderInference/resolveUsualCallWithBuilderInference.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +UnrestrictedBuilderInference
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // FILE: annotation.kt

compiler/testData/ir/irText/expressions/kt47082.fir.txt

@@ -77,11 +77,11 @@ FILE fqName:<root> fileName:/kt47082.kt
               $receiver: VALUE_PARAMETER name:<this> type:<root>.Derived<R of <root>.foo>
               BLOCK_BODY
                 TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                  CALL 'public final fun toChannel <E, C> (destination: C of <root>.toChannel): C of <root>.toChannel declared in <root>' type=<root>.Derived<R of <root>.foo> origin=null
+                  CALL 'public final fun toChannel <E, C> (destination: C of <root>.toChannel): C of <root>.toChannel declared in <root>' type=<root>.Derived<IrErrorType(null)> origin=null
                     <E>: R of <root>.foo
-                    <C>: <root>.Derived<R of <root>.foo>
+                    <C>: <root>.Derived<IrErrorType(null)>
                     $receiver: GET_VAR 'r: <root>.Receiver<R of <root>.foo> declared in <root>.foo' type=<root>.Receiver<R of <root>.foo> origin=null
-                    destination: GET_VAR '<this>: <root>.Derived<R of <root>.foo> declared in <root>.foo.<anonymous>' type=<root>.Derived<R of <root>.foo> origin=null
+                    destination: GET_VAR '<this>: <root>.Derived<R of <root>.foo> declared in <root>.foo.<anonymous>' type=<root>.Derived<IrErrorType(null)> origin=null
   FUN name:box visibility:public modality:FINAL <> () returnType:kotlin.String
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun box (): kotlin.String declared in <root>'

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

@@ -17677,6 +17677,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/inference/builderInference/constraintsBetweenTwoStubVariables.kt");
             }
 
+            @Test
+            @TestMetadata("inferFromExpectedType.kt")
+            public void testInferFromExpectedType() throws Exception {
+                runTest("compiler/testData/codegen/box/inference/builderInference/inferFromExpectedType.kt");
+            }
+
             @Test
             @TestMetadata("intersect.kt")
             public void testIntersect() throws Exception {