FIR: substitute the whole lambda body after builder inference

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

@@ -7,6 +7,7 @@ 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.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.hasAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
@@ -28,6 +29,7 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.BUILDER_INFERENCE_ANNOTATION_
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 
 class FirBuilderInferenceSession(
+    private val lambda: FirAnonymousFunction,
     resolutionContext: ResolutionContext,
     private val stubsForPostponedVariables: Map<ConeTypeVariable, ConeStubType>,
 ) : AbstractManyCandidatesInferenceSession(resolutionContext) {
@@ -225,19 +227,16 @@ class FirBuilderInferenceSession(
         return introducedConstraint
     }
 
-    // TODO: besides calls, perhaps use the stub type substitutor for all top-level expressions inside the lambda
     private fun updateCalls(commonSystem: NewConstraintSystemImpl) {
         val nonFixedToVariablesSubstitutor = createNonFixedTypeToVariableSubstitutor()
         val commonSystemSubstitutor = commonSystem.buildCurrentSubstitutor() as ConeSubstitutor
         val nonFixedTypesToResultSubstitutor = ConeComposedSubstitutor(commonSystemSubstitutor, nonFixedToVariablesSubstitutor)
-        val completionResultsWriter = components.callCompleter.createCompletionResultsWriter(nonFixedTypesToResultSubstitutor)
 
         val stubTypeSubstitutor = FirStubTypeTransformer(nonFixedTypesToResultSubstitutor)
-        for ((completedCall, _) in commonCalls) {
+        lambda.transformSingle(stubTypeSubstitutor, null)
-            completedCall.transformSingle(stubTypeSubstitutor, null)
+        // TODO: support diagnostics, see [CoroutineInferenceSession#updateCalls]
-            // TODO: support diagnostics, see [CoroutineInferenceSession#updateCalls]
-        }
 
+        val completionResultsWriter = components.callCompleter.createCompletionResultsWriter(nonFixedTypesToResultSubstitutor)
         for ((call, _) in partiallyResolvedCalls) {
             call.transformSingle(completionResultsWriter, null)
             // TODO: support diagnostics, see [CoroutineInferenceSession#updateCalls]

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

@@ -272,7 +272,11 @@ class FirCallCompleter(
 
             val builderInferenceSession = runIf(stubsForPostponedVariables.isNotEmpty()) {
                 @Suppress("UNCHECKED_CAST")
-                FirBuilderInferenceSession(transformer.resolutionContext, stubsForPostponedVariables as Map<ConeTypeVariable, ConeStubType>)
+                FirBuilderInferenceSession(
+                    lambdaArgument,
+                    transformer.resolutionContext,
+                    stubsForPostponedVariables as Map<ConeTypeVariable, ConeStubType>
+                )
             }
 
             transformer.context.withAnonymousFunctionTowerDataContext(lambdaArgument.symbol) {

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

@@ -1,7 +1,6 @@
 // !LANGUAGE: +UnrestrictedBuilderInference
 // !DIAGNOSTICS: -DEPRECATION -EXPERIMENTAL_IS_NOT_ENABLED
 // WITH_RUNTIME
-// IGNORE_BACKEND_FIR: JVM_IR
 // DONT_TARGET_EXACT_BACKEND: WASM
 
 // FILE: main.kt

compiler/testData/diagnostics/tests/inference/builderInference/constraints/violating.fir.kt

@@ -10,7 +10,7 @@ fun main() {
     buildList {
         add(3)
         object : A {
-            override fun foo(): MutableList<String> = this@buildList
+            override fun foo(): MutableList<String> = <!RETURN_TYPE_MISMATCH!>this@buildList<!>
         }
     }
     buildList {
@@ -19,12 +19,12 @@ fun main() {
     }
     buildList {
         add("3")
-        val x: MutableList<Int> = this@buildList
+        val x: MutableList<Int> = <!INITIALIZER_TYPE_MISMATCH!>this@buildList<!>
     }
     buildList {
         val y: CharSequence = ""
         add(y)
-        val x: MutableList<String> = this@buildList
+        val x: MutableList<String> = <!INITIALIZER_TYPE_MISMATCH!>this@buildList<!>
     }
     buildList {
         add("")

compiler/testData/diagnostics/tests/inference/builderInference/stubTypes/memberScope.fir.kt

@@ -37,8 +37,8 @@ fun test() {
         get()?.hashCode()
         get()?.<!NONE_APPLICABLE!>equals<!>(1)
         val x = get()
-        x<!UNNECESSARY_SAFE_CALL!>?.<!>hashCode()
+        x?.hashCode()
-        x<!UNNECESSARY_SAFE_CALL!>?.<!><!NONE_APPLICABLE!>equals<!>(1)
+        x?.<!NONE_APPLICABLE!>equals<!>(1)
         x.equals("")
     }
     val ret3 = build {
@@ -50,20 +50,20 @@ fun test() {
         get()?.hashCode()
         get()?.<!NONE_APPLICABLE!>equals<!>(1)
         val x = get()
-        x<!UNNECESSARY_SAFE_CALL!>?.<!>hashCode()
+        x?.hashCode()
-        x<!UNNECESSARY_SAFE_CALL!>?.<!><!NONE_APPLICABLE!>equals<!>(1)
+        x?.<!NONE_APPLICABLE!>equals<!>(1)
 
         if (get() == null) {}
         if (get() === null) {}
 
         if (x != null) {
-            x<!UNNECESSARY_SAFE_CALL!>?.<!>hashCode()
+            x?.hashCode()
-            x<!UNNECESSARY_SAFE_CALL!>?.<!><!NONE_APPLICABLE!>equals<!>(1)
+            x?.<!NONE_APPLICABLE!>equals<!>(1)
             x.equals("")
             x.hashCode()
             x.toString()
             x.test()
-            x<!UNNECESSARY_SAFE_CALL!>?.<!>test2()
+            x?.test2()
             x.test2()
         }
 
@@ -166,8 +166,8 @@ fun test() {
         get()?.hashCode()
         get()?.<!NONE_APPLICABLE!>equals<!>(1)
         val x = get()
-        x<!UNNECESSARY_SAFE_CALL!>?.<!>hashCode()
+        x?.hashCode()
-        x<!UNNECESSARY_SAFE_CALL!>?.<!><!NONE_APPLICABLE!>equals<!>(1)
+        x?.<!NONE_APPLICABLE!>equals<!>(1)
 
         if (get() == null) {}
         if (get() === null) {}
@@ -296,8 +296,8 @@ fun test() {
         get()?.hashCode()
         get()?.<!NONE_APPLICABLE!>equals<!>(1)
         val x = get()
-        x<!UNNECESSARY_SAFE_CALL!>?.<!>hashCode()
+        x?.hashCode()
-        x<!UNNECESSARY_SAFE_CALL!>?.<!><!NONE_APPLICABLE!>equals<!>(1)
+        x?.<!NONE_APPLICABLE!>equals<!>(1)
 
         if (get() == null) {}
         if (get() === null) {}

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

@@ -19,7 +19,7 @@ fun <E : Any?, C : Base<E>> Receiver<E>.toChannel(destination: C): C {
 }
 
 fun <R : Any?> foo(r: Receiver<R>): R {
-  return produce<R>(block = local fun Derived<R>.<anonymous>() {
+  return produce<R>(block = local fun Derived<ErrorType>.<anonymous>() {
     r.toChannel<R, Derived<ErrorType>>(destination = <this>) /*~> Unit */
   }
 )
@@ -28,3 +28,4 @@ fun <R : Any?> foo(r: Receiver<R>): R {
 fun box(): String {
   return "OK"
 }
+

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

@@ -72,16 +72,16 @@ FILE fqName:<root> fileName:/kt47082.kt
       RETURN type=kotlin.Nothing from='public final fun foo <R> (r: <root>.Receiver<R of <root>.foo>): R of <root>.foo declared in <root>'
         CALL 'public final fun produce <E> (block: @[ExtensionFunctionType] kotlin.Function1<<root>.Derived<E of <root>.produce>, kotlin.Unit>): E of <root>.produce declared in <root>' type=R of <root>.foo origin=null
           <E>: R of <root>.foo
-          block: FUN_EXPR type=kotlin.Function1<<root>.Derived<R of <root>.foo>, kotlin.Unit> origin=LAMBDA
+          block: FUN_EXPR type=kotlin.Function1<<root>.Derived<IrErrorType(null)>, kotlin.Unit> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Derived<R of <root>.foo>) returnType:kotlin.Unit
+            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Derived<IrErrorType(null)>) returnType:kotlin.Unit
-              $receiver: VALUE_PARAMETER name:<this> type:<root>.Derived<R of <root>.foo>
+              $receiver: VALUE_PARAMETER name:<this> type:<root>.Derived<IrErrorType(null)>
               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<IrErrorType(null)> origin=null
                     <E>: 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<IrErrorType(null)> origin=null
+                    destination: GET_VAR '<this>: <root>.Derived<IrErrorType(null)> 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>'