FIR: Fix builder inference case with independent calls inside lambda

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

@@ -31624,6 +31624,22 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/builderInference")
+        @TestDataPath("$PROJECT_ROOT")
+        public class BuilderInference extends AbstractFirDiagnosticTest {
+            @Test
+            public void testAllFilesPresentInBuilderInference() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/builderInference"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+            }
+
+            @Test
+            @TestMetadata("completeIrrelevantCalls.kt")
+            public void testCompleteIrrelevantCalls() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/completeIrrelevantCalls.kt");
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/builtins")
         @TestDataPath("$PROJECT_ROOT")

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolutionStages.kt

@@ -239,14 +239,13 @@ internal object CheckLowPriorityInOverloadResolution : CheckerStage() {
 }
 
 internal object PostponedVariablesInitializerResolutionStage : ResolutionStage() {
-    private val BUILDER_INFERENCE_CLASS_ID: ClassId = ClassId.fromString("kotlin/BuilderInference")
 
     override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val argumentMapping = candidate.argumentMapping ?: return
         // TODO: convert type argument mapping to map [FirTypeParameterSymbol, FirTypedProjection?]
         if (candidate.typeArgumentMapping is TypeArgumentMapping.Mapped) return
         for (parameter in argumentMapping.values) {
-            if (!parameter.hasBuilderInferenceMarker()) continue
+            if (!parameter.hasBuilderInferenceAnnotation()) continue
             val type = parameter.returnTypeRef.coneType
             val receiverType = type.receiverType(callInfo.session) ?: continue
 
@@ -263,8 +262,4 @@ internal object PostponedVariablesInitializerResolutionStage : ResolutionStage()
             }
         }
     }
-
-    private fun FirValueParameter.hasBuilderInferenceMarker(): Boolean {
-        return this.hasAnnotation(BUILDER_INFERENCE_CLASS_ID)
-    }
 }

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

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.fir.resolve.inference
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
+import org.jetbrains.kotlin.fir.declarations.hasAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
@@ -13,13 +15,18 @@ import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.visitors.*
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.name.ClassId
 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.ConstraintKind
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.inference.model.CoroutinePosition
 import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintSystemImpl
+import org.jetbrains.kotlin.resolve.descriptorUtil.BUILDER_INFERENCE_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 
 class FirBuilderInferenceSession(
@@ -33,6 +40,8 @@ class FirBuilderInferenceSession(
         val system = candidate.system
 
         if (system.hasContradiction) return true
+        if (!candidate.isSuitableForBuilderInference()) return true
+
 
         val storage = system.getBuilder().currentStorage()
 
@@ -45,6 +54,23 @@ class FirBuilderInferenceSession(
         }
     }
 
+    private fun Candidate.isSuitableForBuilderInference(): Boolean {
+        val extensionReceiver = extensionReceiverValue
+        val dispatchReceiver = dispatchReceiverValue
+        return when {
+            extensionReceiver == null && dispatchReceiver == null -> false
+            dispatchReceiver?.type?.containsStubType() == true -> true
+            extensionReceiver?.type?.containsStubType() == true -> symbol.fir.hasBuilderInferenceAnnotation()
+            else -> false
+        }
+    }
+
+    private fun ConeKotlinType.containsStubType(): Boolean {
+        return this.contains {
+            it is ConeStubType
+        }
+    }
+
     private fun FirStatement.hasPostponed(): Boolean {
         var result = false
         processAllContainingCallCandidates(processBlocks = false) {
@@ -243,3 +269,8 @@ class FirStubTypeTransformer(
     override fun transformArgumentList(argumentList: FirArgumentList, data: Nothing?): CompositeTransformResult<FirArgumentList> =
         argumentList.transformArguments(this, data).compose()
 }
+
+private val BUILDER_INFERENCE_ANNOTATION_CLASS_ID = ClassId.topLevel(BUILDER_INFERENCE_ANNOTATION_FQ_NAME)
+
+fun FirElement.hasBuilderInferenceAnnotation(): Boolean =
+    (this as? FirAnnotatedDeclaration)?.hasAnnotation(BUILDER_INFERENCE_ANNOTATION_CLASS_ID) == true

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

@@ -0,0 +1,16 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+
+class A {
+    lateinit var m: Map<String, Int>
+
+    @ExperimentalStdlibApi
+    fun foo(xs: Collection<List<String>>) {
+        m = buildMap {
+            // flatMap calls might be completed on early phase
+            for (x in xs.flatMap { it.toList() }) {
+                put(x, x.length)
+            }
+        }
+    }
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/kt41308.fir.kt

@@ -1,9 +0,0 @@
-// ISSUE: KT-41308
-
-fun main() {
-    sequence {
-        val list: List<String>? = null
-        val outputList = <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.List<kotlin.String>")!>list ?: listOf()<!>
-        <!NONE_APPLICABLE!>yieldAll<!>(outputList)
-    }
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // ISSUE: KT-41308
 
 fun main() {

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

@@ -31720,6 +31720,22 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
             }
         }
 
+        @Nested
+        @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/builderInference")
+        @TestDataPath("$PROJECT_ROOT")
+        public class BuilderInference extends AbstractDiagnosticTest {
+            @Test
+            public void testAllFilesPresentInBuilderInference() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/builderInference"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+            }
+
+            @Test
+            @TestMetadata("completeIrrelevantCalls.kt")
+            public void testCompleteIrrelevantCalls() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/completeIrrelevantCalls.kt");
+            }
+        }
+
         @Nested
         @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/builtins")
         @TestDataPath("$PROJECT_ROOT")