FIR. BI: Fix stub type leakage in builder inference

Previously, updateTypeInBuilderInference was calling
updateTypeFromSmartcast to actually perform type update in
implicit receiver after stub types was inferred
Such action results in creation of following FIR:

FirSmartCastExpression(
    original=FQAE(FirImplicitThisReference, typeRef=R|Inv<Stub>|)
    typeRef=R|Inv<String>|
)

in receiver position during completion of calls

However, it wasn't the case in general situation due to
action of FirStubTypeTransformer, which, in turn visits and updates
type ref inside original expression, but only if there was at
least one call (that was completed) using that implicit receiver
As after such type update updateTypeFromSmartcast function does
nothing

Yet in situation, when there was only partially resolved calls
referencing that implicit receiver we actually create smart-cast
expression and don't update type

The change just removes usage of updateTypeFromSmartcast and
replaces is with direct type update
We still mutate state of implicit receiver, potentially
improperly, it should be addressed in future

^KT-54708
^KT-58365 Fixed

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

@@ -36820,6 +36820,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                 runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/inferenceFromLambdaReturnType.kt");
             }
 
+            @Test
+            @TestMetadata("k2StubTypeLeak.kt")
+            public void testK2StubTypeLeak() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/k2StubTypeLeak.kt");
+            }
+
             @Test
             @TestMetadata("resolveUsualCallWithBuilderInference.kt")
             public void testResolveUsualCallWithBuilderInference() throws Exception {

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

@@ -36820,6 +36820,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                 runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/inferenceFromLambdaReturnType.kt");
             }
 
+            @Test
+            @TestMetadata("k2StubTypeLeak.kt")
+            public void testK2StubTypeLeak() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/k2StubTypeLeak.kt");
+            }
+
             @Test
             @TestMetadata("resolveUsualCallWithBuilderInference.kt")
             public void testResolveUsualCallWithBuilderInference() throws Exception {

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

@@ -36820,6 +36820,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                 runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/inferenceFromLambdaReturnType.kt");
             }
 
+            @Test
+            @TestMetadata("k2StubTypeLeak.kt")
+            public void testK2StubTypeLeak() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/k2StubTypeLeak.kt");
+            }
+
             @Test
             @TestMetadata("resolveUsualCallWithBuilderInference.kt")
             public void testResolveUsualCallWithBuilderInference() throws Exception {

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

@@ -36916,6 +36916,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                 runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/inferenceFromLambdaReturnType.kt");
             }
 
+            @Test
+            @TestMetadata("k2StubTypeLeak.kt")
+            public void testK2StubTypeLeak() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/k2StubTypeLeak.kt");
+            }
+
             @Test
             @TestMetadata("resolveUsualCallWithBuilderInference.kt")
             public void testResolveUsualCallWithBuilderInference() throws Exception {

compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/calls/FirReceivers.kt

@@ -105,10 +105,12 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
     @RequiresOptIn
     annotation class ImplicitReceiverInternals
 
-    @OptIn(ImplicitReceiverInternals::class)
     @Deprecated(level = DeprecationLevel.ERROR, message = "Builder inference should not modify implicit receivers. KT-54708")
     fun updateTypeInBuilderInference(type: ConeKotlinType) {
-        updateTypeFromSmartcast(type)
+        this.type = type
+        receiverExpression = receiverExpression(boundSymbol, type, contextReceiverNumber)
+        implicitScope =
+            type.scope(useSiteSession, scopeSession, FakeOverrideTypeCalculator.DoNothing, requiredPhase = FirResolvePhase.STATUS)
     }
 
     /*

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

@@ -241,6 +241,7 @@ class FirBuilderInferenceSession(
         val stubTypeSubstitutor = FirStubTypeTransformer(substitutor)
         lambda.transformSingle(stubTypeSubstitutor, null)
 
+        // TODO: Builder inference should not modify implicit receivers. KT-54708
         for (receiver in lambdaImplicitReceivers) {
             @Suppress("DEPRECATION_ERROR")
             receiver.updateTypeInBuilderInference(substitutor.substituteOrSelf(receiver.type))

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

@@ -1,4 +1,3 @@
-// IGNORE_LEAKED_INTERNAL_TYPES: KT-54708
 // WITH_STDLIB
 
 import kotlin.experimental.ExperimentalTypeInference

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

@@ -0,0 +1,20 @@
+// FIR_IDENTICAL
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+fun test_1() {
+    sequence {
+        constrain(this)
+        yieldAll(mk()) // Will be completed in partial mode, due to builder inference
+
+        Unit
+    }
+}
+
+fun constrain(t: Inv<String>) {}
+
+fun <U> sequence(block: Inv<U>.() -> Unit): U = null!!
+
+interface Inv<T> {
+    fun <S: T> yieldAll(seq: Inv<S>)
+}
+
+fun <K> mk(): K = TODO()

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

@@ -1,4 +1,3 @@
-// IGNORE_LEAKED_INTERNAL_TYPES: KT-54708
 // !OPT_IN: kotlin.RequiresOptIn
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // ISSUE: KT-35684

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

@@ -1,4 +1,3 @@
-// IGNORE_LEAKED_INTERNAL_TYPES: KT-54708
 // !OPT_IN: kotlin.RequiresOptIn
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // ISSUE: KT-35684

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

@@ -37688,6 +37688,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/inferenceFromLambdaReturnType.kt");
             }
 
+            @Test
+            @TestMetadata("k2StubTypeLeak.kt")
+            public void testK2StubTypeLeak() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/builderInference/k2StubTypeLeak.kt");
+            }
+
             @Test
             @TestMetadata("resolveUsualCallWithBuilderInference.kt")
             public void testResolveUsualCallWithBuilderInference() throws Exception {