FIR: update unsubstituted return types in builder

#KT-43340 Fixed

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

@@ -5,15 +5,13 @@
 
 package org.jetbrains.kotlin.fir.resolve.inference
 
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 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.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.ConeStubType
-import org.jetbrains.kotlin.fir.types.ConeTypeVariable
-import org.jetbrains.kotlin.fir.types.ConeTypeVariableTypeConstructor
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
@@ -203,11 +201,33 @@ class FirBuilderInferenceSession(
         val commonSystemSubstitutor = commonSystem.buildCurrentSubstitutor() as ConeSubstitutor
         val nonFixedTypesToResultSubstitutor = ConeComposedSubstitutor(commonSystemSubstitutor, nonFixedToVariablesSubstitutor)
         val completionResultsWriter = components.callCompleter.createCompletionResultsWriter(nonFixedTypesToResultSubstitutor)
+
+        for ((completedCall, _) in commonCalls) {
+            // TODO: Only update return type? Should we need to visit all appearances of unsubstituted postponed variables in types?
+            //  [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]
+        }
+
         for ((call, _) in partiallyResolvedCalls) {
             call.transformSingle(completionResultsWriter, null)
-            // TODO: support diagnostics, see CoroutineInferenceSession.kt:286
+            // 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() {

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

@@ -552,27 +552,32 @@ FILE fqName:<root> fileName:/typeVariableAfterBuildMap.kt
               FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int>) returnType:kotlin.Unit
                 $receiver: VALUE_PARAMETER name:<this> type:kotlin.collections.MutableMap<<root>.Visibility, kotlin.Int>
                 BLOCK_BODY
-                  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=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<IrErrorType, IrErrorType> 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
-                  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=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<IrErrorType, IrErrorType> 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
-                  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=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<IrErrorType, IrErrorType> 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
-                  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=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<IrErrorType, IrErrorType> 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'
-                    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=IrErrorType origin=null
+                  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
-                      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
+                      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
+                      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
+                      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
+                      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
+                        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>
         correspondingProperty: PROPERTY name:ORDERED_VISIBILITIES visibility:private modality:FINAL [val]
         $this: VALUE_PARAMETER name:<this> type:<root>.Visibilities

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

@@ -39,7 +39,7 @@ FILE fqName:<root> fileName:/castsInsideCoroutineInference.kt
               $receiver: VALUE_PARAMETER name:<this> type:<root>.FlowCollector<T of <root>.onCompletion>
               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<IrErrorType> origin=null
+                  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
                     collector: GET_VAR '<this>: <root>.FlowCollector<T of <root>.onCompletion> declared in <root>.onCompletion.<anonymous>' type=<root>.FlowCollector<IrErrorType> 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