FIR: more comprehensive substitution of stub types after builder inference

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

@@ -5,14 +5,15 @@
 
 package org.jetbrains.kotlin.fir.resolve.inference
 
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 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.transformSingle
+import org.jetbrains.kotlin.fir.visitors.*
 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
@@ -196,21 +197,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) {
-            // TODO: Only update return type? Should we need to visit all appearances of unsubstituted postponed variables in types?
+            completedCall.transformSingle(stubTypeSubstitutor, null)
-            //  [transformSingle] bails out very early since the completed call is literally completed, and not a named reference.
-            (completedCall as? FirFunctionCall)?.let { call ->
-                val resultType = call.typeRef.substituteTypeRef(nonFixedTypesToResultSubstitutor)
-                if (resultType != call.typeRef) {
-                    call.replaceTypeRef(resultType)
-                }
-            }
             // TODO: support diagnostics, see [CoroutineInferenceSession#updateCalls]
         }
 
@@ -219,15 +215,6 @@ class FirBuilderInferenceSession(
             // TODO: support diagnostics, see [CoroutineInferenceSession#updateCalls]
         }
     }
-
-    private fun FirTypeRef.substituteTypeRef(
-        substitutor: ConeSubstitutor,
-    ): FirTypeRef =
-        (this as? FirResolvedTypeRef)?.let {
-            substitutor.substituteOrNull(this.type)?.let {
-                this.withReplacedConeType(it)
-            }
-        } ?: this
 }
 
 class ConeComposedSubstitutor(val left: ConeSubstitutor, val right: ConeSubstitutor) : ConeSubstitutor() {
@@ -236,3 +223,21 @@ class ConeComposedSubstitutor(val left: ConeSubstitutor, val right: ConeSubstitu
         return left.substituteOrNull(rightSubstitution ?: type)
     }
 }
+
+class FirStubTypeTransformer(
+    private val substitutor: ConeSubstitutor
+) : FirDefaultTransformer<Nothing?>() {
+
+    override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+        @Suppress("UNCHECKED_CAST")
+        return (element.transformChildren(this, data) as E).compose()
+    }
+
+    override fun transformResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef, data: Nothing?): CompositeTransformResult<FirTypeRef> =
+        substitutor.substituteOrNull(resolvedTypeRef.type)?.let {
+            resolvedTypeRef.withReplacedConeType(it).compose()
+        } ?: resolvedTypeRef.compose()
+
+    override fun transformArgumentList(argumentList: FirArgumentList, data: Nothing?): CompositeTransformResult<FirArgumentList> =
+        argumentList.transformArguments(this, data).compose()
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirResolvedArgumentList.kt

@@ -7,12 +7,19 @@ package org.jetbrains.kotlin.fir.expressions.impl
 
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirAbstractArgumentList
+import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 class FirResolvedArgumentList internal constructor(
-    val mapping: LinkedHashMap<FirExpression, FirValueParameter>
+    mapping: LinkedHashMap<FirExpression, FirValueParameter>
 ) : FirAbstractArgumentList() {
+
+    var mapping: LinkedHashMap<FirExpression, FirValueParameter> = mapping
+        private set
+
     override val arguments: List<FirExpression>
         get() = mapping.keys.toList()
 
@@ -21,4 +28,9 @@ class FirResolvedArgumentList internal constructor(
             argument.accept(visitor, data)
         }
     }
+
+    override fun <D> transformArguments(transformer: FirTransformer<D>, data: D): FirArgumentList {
+        mapping = mapping.mapKeys { (k, _) -> k.transformSingle(transformer, data) } as LinkedHashMap<FirExpression, FirValueParameter>
+        return this
+    }
 }

compiler/testData/codegen/box/builderInference/lackOfNullCheckOnNullableInsideBuild.kt

@@ -1,6 +1,5 @@
 // !LANGUAGE: +NewInference
 // WITH_RUNTIME
-// IGNORE_BACKEND_FIR: JVM_IR
 // Issue: KT-36371
 
 import kotlin.experimental.ExperimentalTypeInference

compiler/testData/codegen/box/builderInference/substituteStubTypeIntoCR.kt

@@ -1,6 +1,5 @@
 // DONT_TARGET_EXACT_BACKEND: WASM
 // WASM_MUTE_REASON: COROUTINES
-// IGNORE_BACKEND_FIR: JVM_IR
 // IGNORE_BACKEND: JVM_IR
 // IGNORE_BACKEND: JS_IR
 // IGNORE_BACKEND: NATIVE

compiler/testData/codegen/box/builderInference/substituteStubTypeIntolambdaParameterDescriptor.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 
 import kotlin.experimental.ExperimentalTypeInference

compiler/testData/codegen/box/builderInference/substituteTypeVariableIntolambdaParameterDescriptor.kt

@@ -1,5 +1,4 @@
 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 
 class Foo<A>

compiler/testData/ir/irText/firProblems/typeVariableAfterBuildMap.fir.txt

@@ -554,28 +554,28 @@ FILE fqName:<root> fileName:/typeVariableAfterBuildMap.kt
                 BLOCK_BODY
                   TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
                     CALL 'public abstract fun put (key: K of kotlin.collections.MutableMap, value: V of kotlin.collections.MutableMap): V of kotlin.collections.MutableMap? declared in kotlin.collections.MutableMap' type=kotlin.Int? origin=null
-                      $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<IrErrorType, IrErrorType> origin=null
+                      $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> origin=null
                       key: GET_OBJECT 'CLASS OBJECT name:PrivateToThis modality:FINAL visibility:public superTypes:[<root>.Visibility]' type=<root>.Visibilities.PrivateToThis
                       value: CONST Int type=kotlin.Int value=0
                   TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
                     CALL 'public abstract fun put (key: K of kotlin.collections.MutableMap, value: V of kotlin.collections.MutableMap): V of kotlin.collections.MutableMap? declared in kotlin.collections.MutableMap' type=kotlin.Int? origin=null
-                      $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<IrErrorType, IrErrorType> origin=null
+                      $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> origin=null
                       key: GET_OBJECT 'CLASS OBJECT name:Private modality:FINAL visibility:public superTypes:[<root>.Visibility]' type=<root>.Visibilities.Private
                       value: CONST Int type=kotlin.Int value=0
                   TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
                     CALL 'public abstract fun put (key: K of kotlin.collections.MutableMap, value: V of kotlin.collections.MutableMap): V of kotlin.collections.MutableMap? declared in kotlin.collections.MutableMap' type=kotlin.Int? origin=null
-                      $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<IrErrorType, IrErrorType> origin=null
+                      $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> origin=null
                       key: GET_OBJECT 'CLASS OBJECT name:Internal modality:FINAL visibility:public superTypes:[<root>.Visibility]' type=<root>.Visibilities.Internal
                       value: CONST Int type=kotlin.Int value=1
                   TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
                     CALL 'public abstract fun put (key: K of kotlin.collections.MutableMap, value: V of kotlin.collections.MutableMap): V of kotlin.collections.MutableMap? declared in kotlin.collections.MutableMap' type=kotlin.Int? origin=null
-                      $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<IrErrorType, IrErrorType> origin=null
+                      $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> origin=null
                       key: GET_OBJECT 'CLASS OBJECT name:Protected modality:FINAL visibility:public superTypes:[<root>.Visibility]' type=<root>.Visibilities.Protected
                       value: CONST Int type=kotlin.Int value=1
                   RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.Visibilities.ORDERED_VISIBILITIES'
                     TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
                       CALL 'public abstract fun put (key: K of kotlin.collections.MutableMap, value: V of kotlin.collections.MutableMap): V of kotlin.collections.MutableMap? declared in kotlin.collections.MutableMap' type=kotlin.Int? origin=null
-                        $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<IrErrorType, IrErrorType> origin=null
+                        $this: GET_VAR '<this>: kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> declared in <root>.Visibilities.ORDERED_VISIBILITIES.<anonymous>' type=kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int> origin=null
                         key: GET_OBJECT 'CLASS OBJECT name:Public modality:FINAL visibility:public superTypes:[<root>.Visibility]' type=<root>.Visibilities.Public
                         value: CONST Int type=kotlin.Int value=2
       FUN DEFAULT_PROPERTY_ACCESSOR name:<get-ORDERED_VISIBILITIES> visibility:private modality:FINAL <> ($this:<root>.Visibilities) returnType:kotlin.collections.Map<<root>.Visibility, kotlin.Int>

compiler/testData/ir/irText/types/castsInsideCoroutineInference.fir.txt

@@ -25,7 +25,8 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
                         CALL 'public abstract fun invoke (p1: P1 of kotlin.coroutines.SuspendFunction2, p2: P2 of kotlin.coroutines.SuspendFunction2): R of kotlin.coroutines.SuspendFunction2 [suspend,operator] declared in kotlin.coroutines.SuspendFunction2' type=kotlin.Unit origin=null
                           $this: GET_VAR 'block: @[ExtensionFunctionType] @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction2<<root>.CoroutineScope, <root>.FlowCollector<R of <root>.scopedFlow>, kotlin.Unit> declared in <root>.scopedFlow' type=@[ExtensionFunctionType] @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction2<<root>.CoroutineScope, <root>.FlowCollector<R of <root>.scopedFlow>, kotlin.Unit> origin=VARIABLE_AS_FUNCTION
                           p1: GET_VAR '<this>: <root>.CoroutineScope declared in <root>.scopedFlow.<anonymous>.<anonymous>' type=<root>.CoroutineScope origin=null
-                          p2: GET_VAR 'val collector: <root>.FlowCollector<IrErrorType> [val] declared in <root>.scopedFlow.<anonymous>' type=<root>.FlowCollector<IrErrorType> origin=null
+                          p2: TYPE_OP type=<root>.FlowCollector<IrErrorType> origin=IMPLICIT_CAST typeOperand=<root>.FlowCollector<IrErrorType>
+                            GET_VAR 'val collector: <root>.FlowCollector<IrErrorType> [val] declared in <root>.scopedFlow.<anonymous>' type=<root>.FlowCollector<R of <root>.scopedFlow> origin=null
   FUN name:onCompletion visibility:public modality:FINAL <T> ($receiver:<root>.Flow<T of <root>.onCompletion>, action:@[ExtensionFunctionType] @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.onCompletion>, kotlin.Throwable?, kotlin.Unit>) returnType:<root>.Flow<T of <root>.onCompletion>
     TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?]
     $receiver: VALUE_PARAMETER name:<this> type:<root>.Flow<T of <root>.onCompletion>
@@ -40,11 +41,11 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
               BLOCK_BODY
                 VAR name:safeCollector type:<root>.SafeCollector<IrErrorType> [val]
                   CONSTRUCTOR_CALL 'public constructor <init> (collector: <root>.FlowCollector<T of <root>.SafeCollector>) [primary] declared in <root>.SafeCollector' type=<root>.SafeCollector<T of <root>.onCompletion> origin=null
-                    <class: T>: IrErrorType
+                    <class: T>: T of <root>.onCompletion
-                    collector: GET_VAR '<this>: <root>.FlowCollector<T of <root>.onCompletion> declared in <root>.onCompletion.<anonymous>' type=<root>.FlowCollector<IrErrorType> origin=null
+                    collector: GET_VAR '<this>: <root>.FlowCollector<T of <root>.onCompletion> declared in <root>.onCompletion.<anonymous>' type=<root>.FlowCollector<T of <root>.onCompletion> origin=null
                 CALL 'public final fun invokeSafely <T> (action: @[ExtensionFunctionType] @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.invokeSafely>, kotlin.Throwable?, kotlin.Unit>): kotlin.Unit [suspend] declared in <root>' type=kotlin.Unit origin=null
                   <T>: T of <root>.onCompletion
-                  $receiver: GET_VAR 'val safeCollector: <root>.SafeCollector<IrErrorType> [val] declared in <root>.onCompletion.<anonymous>' type=<root>.SafeCollector<IrErrorType> origin=null
+                  $receiver: GET_VAR 'val safeCollector: <root>.SafeCollector<IrErrorType> [val] declared in <root>.onCompletion.<anonymous>' type=<root>.SafeCollector<T of <root>.onCompletion> origin=null
                   action: GET_VAR 'action: @[ExtensionFunctionType] @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.onCompletion>, kotlin.Throwable?, kotlin.Unit> declared in <root>.onCompletion' type=@[ExtensionFunctionType] @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.onCompletion>, kotlin.Throwable?, kotlin.Unit> origin=null
   FUN name:invokeSafely visibility:public modality:FINAL <T> ($receiver:<root>.FlowCollector<T of <root>.invokeSafely>, action:@[ExtensionFunctionType] @[ExtensionFunctionType] kotlin.coroutines.SuspendFunction2<<root>.FlowCollector<T of <root>.invokeSafely>, kotlin.Throwable?, kotlin.Unit>) returnType:kotlin.Unit [suspend]
     TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?]
@@ -136,8 +137,8 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
                       BLOCK_BODY
                         RETURN type=kotlin.Nothing from='local final fun <anonymous> (value: kotlin.Any?): kotlin.Unit [suspend] declared in <root>.asChannel.<anonymous>'
                           CALL 'public abstract fun send (e: E of <root>.SendChannel): kotlin.Unit [suspend] declared in <root>.SendChannel' type=kotlin.Unit origin=null
-                            $this: CALL 'public abstract fun <get-channel> (): <root>.SendChannel<E of <root>.ProducerScope> declared in <root>.ProducerScope' type=<root>.SendChannel<IrErrorType> origin=GET_PROPERTY
+                            $this: CALL 'public abstract fun <get-channel> (): <root>.SendChannel<E of <root>.ProducerScope> declared in <root>.ProducerScope' type=<root>.SendChannel<kotlin.Any> origin=GET_PROPERTY
-                              $this: GET_VAR '<this>: <root>.ProducerScope<kotlin.Any> declared in <root>.asChannel.<anonymous>' type=<root>.ProducerScope<IrErrorType> origin=null
+                              $this: GET_VAR '<this>: <root>.ProducerScope<kotlin.Any> declared in <root>.asChannel.<anonymous>' type=<root>.ProducerScope<kotlin.Any> origin=null
                             e: BLOCK type=kotlin.Any origin=ELVIS
                               VAR IR_TEMPORARY_VARIABLE name:tmp_1 type:kotlin.Any? [val]
                                 GET_VAR 'value: kotlin.Any? declared in <root>.asChannel.<anonymous>.<anonymous>' type=kotlin.Any? origin=null