FIR resolve (by semoro): support correct type inference for generic args

This commit includes additional test and fixes e.g.
resolve of listOf() + listOf()

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.resolve.transformers.firUnsafe
 import org.jetbrains.kotlin.fir.resolve.withNullability
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
@@ -37,7 +38,8 @@ fun resolveArgumentExpression(
     typeProvider: (FirExpression) -> FirTypeRef?
 ) {
     return when (argument) {
-        is FirQualifiedAccessExpression, is FirFunctionCall -> resolvePlainExpressionArgument(
+        is FirFunctionCall -> resolveSubCallArgument(csBuilder, argument, expectedType, sink, isReceiver, isSafeCall, typeProvider)
+        is FirQualifiedAccessExpression -> resolvePlainExpressionArgument(
             csBuilder,
             argument,
             expectedType,
@@ -78,6 +80,21 @@ fun resolveArgumentExpression(
     }
 }
 
+fun resolveSubCallArgument(
+    csBuilder: ConstraintSystemBuilder,
+    argument: FirFunctionCall,
+    expectedType: ConeKotlinType,
+    sink: CheckerSink,
+    isReceiver: Boolean,
+    isSafeCall: Boolean,
+    typeProvider: (FirExpression) -> FirTypeRef?
+) {
+    val candidate = argument.candidate() ?: return resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, isSafeCall, typeProvider)
+    val type = sink.components.returnTypeCalculator.tryCalculateReturnType(candidate.symbol.firUnsafe()).coneTypeUnsafe<ConeKotlinType>()
+    val argumentType = candidate.substitutor.substituteOrSelf(type)
+    resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isSafeCall)
+}
+
 fun resolvePlainExpressionArgument(
     csBuilder: ConstraintSystemBuilder,
     argument: FirExpression,

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirWrappedArgumentExpression
 import org.jetbrains.kotlin.fir.symbols.ConeSymbol
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzer
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
@@ -40,6 +41,7 @@ class CandidateFactory(
 fun PostponedArgumentsAnalyzer.Context.addSubsystemFromExpression(expression: FirExpression) {
     when (expression) {
         is FirFunctionCall -> expression.candidate()?.let { addOtherSystem(it.system.asReadOnlyStorage()) }
+        is FirWrappedArgumentExpression -> addSubsystemFromExpression(expression.expression)
     }
 }
 

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirWrappedArgumentExpression
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
@@ -36,23 +37,41 @@ fun Candidate.computeCompletionMode(
     return when {
         // Consider call foo(bar(x)), if return type of bar is a proper one, then we can complete resolve for bar => full completion mode
         // Otherwise, we shouldn't complete bar until we process call foo
-        system.getBuilder().isProperType(currentReturnType) -> KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL
+        csBuilder.isProperType(currentReturnType) -> KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL
 
         // Nested call is connected with the outer one through the UPPER constraint (returnType <: expectedOuterType)
         // This means that there will be no new LOWER constraints =>
         //   it's possible to complete call now if there are proper LOWER constraints
-        system.getBuilder().isTypeVariable(currentReturnType) ->
+        csBuilder.isTypeVariable(currentReturnType) ->
             if (hasProperNonTrivialLowerConstraints(components, currentReturnType))
                 KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL
             else
                 KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.PARTIAL
 
+        // Return type has proper equal constraints => there is no need in the outer call
+        containsTypeVariablesWithProperEqualConstraints(components, currentReturnType) -> KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL
+
         else -> KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.PARTIAL
     }
 }
 
 val Candidate.csBuilder get() = system.getBuilder()
 
+private fun Candidate.containsTypeVariablesWithProperEqualConstraints(components: InferenceComponents, type: ConeKotlinType): Boolean =
+    with(components.ctx){
+        for ((variableConstructor, variableWithConstraints) in csBuilder.currentStorage().notFixedTypeVariables) {
+            if (!type.contains { it.typeConstructor() == variableConstructor }) continue
+
+            val constraints = variableWithConstraints.constraints
+            val onlyProperEqualConstraints =
+                constraints.isNotEmpty() && constraints.all { it.kind.isEqual() && csBuilder.isProperType(it.type) }
+
+            if (!onlyProperEqualConstraints) return false
+        }
+
+        return true
+    }
+
 private fun Candidate.hasProperNonTrivialLowerConstraints(components: InferenceComponents, typeVariable: ConeKotlinType): Boolean {
     assert(csBuilder.isTypeVariable(typeVariable)) { "$typeVariable is not a type variable" }
 
@@ -162,6 +181,7 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
                     }
                     this.arguments.forEach { it.process(to) }
                 }
+                is FirWrappedArgumentExpression -> this.expression.process(to)
                 // TOOD: WTF?
             }
 //            if (analyzed) {

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.symbols.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
@@ -45,7 +46,7 @@ open class ConeTypeVariable(name: String) : TypeVariableMarker {
     val defaultType = ConeTypeVariableType(ConeNullability.NOT_NULL, typeConstructor)
 }
 
-class InferenceComponents(val ctx: TypeSystemInferenceExtensionContextDelegate, val session: FirSession) {
+class InferenceComponents(val ctx: TypeSystemInferenceExtensionContextDelegate, val session: FirSession, val returnTypeCalculator: ReturnTypeCalculator) {
     private val approximator = object : AbstractTypeApproximator(ctx) {
         override fun createErrorType(message: String): SimpleTypeMarker {
             return ConeClassErrorType(message)
@@ -59,5 +60,6 @@ class InferenceComponents(val ctx: TypeSystemInferenceExtensionContextDelegate,
     fun createConstraintSystem(): NewConstraintSystemImpl {
         return NewConstraintSystemImpl(injector, ctx)
     }
+
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt

@@ -255,7 +255,7 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
         override fun KotlinTypeMarker.removeExactAnnotation(): KotlinTypeMarker {
             return this
         }
-    }, session)
+    }, session, jump)
 
     private var qualifierStack = mutableListOf<Name>()
     private var qualifierPartsToDrop = 0

compiler/fir/resolve/testData/resolve/stdlib/topLevelResolve.kt

@@ -0,0 +1,16 @@
+fun testPlus() {
+    val x = 1 + 2
+    val y = 3.0 + 4.0
+    val z = 5 + 6.0
+    val w = 7.0 + 8
+    val c = 'a' + 1
+    val s = "." + ".."
+    val ss = "" + 1
+    val list = listOf(1, 2, 3) + 4
+    val listAndList = listOf(4, 5, 6) + listOf(7, 8)
+    val mutableList = mutableListOf(9, 10) + listOf(11, 12, 13)
+    val setAndList = setOf(0) + listOf(1, 2)
+    val stringAndList = "" + emptyList<Boolean>()
+    val map = mapOf("" to 1, "." to 2) + (".." to 3)
+    val mapAndMap = mapOf("-" to 4) + mapOf("_" to 5)
+}

compiler/fir/resolve/testData/resolve/stdlib/topLevelResolve.txt

@@ -0,0 +1,17 @@
+FILE: topLevelResolve.kt
+    public final fun testPlus(): R|kotlin/Unit| {
+        lval x: R|kotlin/Int| = Int(1).R|kotlin/Int.plus|(Int(2))
+        lval y: R|kotlin/Double| = Double(3.0).R|kotlin/Double.plus|(Double(4.0))
+        lval z: R|kotlin/Double| = Int(5).R|kotlin/Int.plus|(Double(6.0))
+        lval w: R|kotlin/Double| = Double(7.0).R|kotlin/Double.plus|(Int(8))
+        lval c: R|kotlin/Char| = Char(a).R|kotlin/Char.plus|(Int(1))
+        lval s: R|kotlin/String| = String(.).R|kotlin/String.plus|(String(..))
+        lval ss: R|kotlin/String| = String().R|kotlin/String.plus|(Int(1))
+        lval list: R|kotlin/collections/List<kotlin/Int>| = R|kotlin/collections/listOf|<R|kotlin/Int|>(Int(1), Int(2), Int(3)).R|kotlin/collections/plus|<R|kotlin/Int|>(Int(4))
+        lval listAndList: R|kotlin/collections/List<kotlin/Int>| = R|kotlin/collections/listOf|<R|kotlin/Int|>(Int(4), Int(5), Int(6)).R|kotlin/collections/plus|<R|kotlin/Int|>(R|kotlin/collections/listOf|<R|kotlin/Int|>(Int(7), Int(8)))
+        lval mutableList: R|kotlin/collections/List<kotlin/Int>| = R|kotlin/collections/mutableListOf|<R|kotlin/Int|>(Int(9), Int(10)).R|kotlin/collections/plus|<R|kotlin/Int|>(R|kotlin/collections/listOf|<R|kotlin/Int|>(Int(11), Int(12), Int(13)))
+        lval setAndList: R|kotlin/collections/Set<kotlin/Int>| = R|kotlin/collections/setOf|<R|kotlin/Int|>(Int(0)).R|kotlin/collections/plus|<R|kotlin/Int|>(R|kotlin/collections/listOf|<R|kotlin/Int|>(Int(1), Int(2)))
+        lval stringAndList: R|kotlin/String| = String().R|kotlin/String.plus|(R|kotlin/collections/emptyList|<R|kotlin/Boolean|>())
+        lval map: R|kotlin/collections/Map<kotlin/String, kotlin/Int>| = R|kotlin/collections/mapOf|<R|kotlin/String|, R|kotlin/Int|>(String().R|kotlin/to|<R|kotlin/String|, R|kotlin/Int|>(Int(1)), String(.).R|kotlin/to|<R|kotlin/String|, R|kotlin/Int|>(Int(2))).R|kotlin/collections/plus|<R|kotlin/String|, R|kotlin/Int|>(String(..).R|kotlin/to|<R|kotlin/String|, R|kotlin/Int|>(Int(3)))
+        lval mapAndMap: <ERROR TYPE REF: Ambiguity: plus, [kotlin/plus, kotlin/plus, kotlin/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/sequences/plus, kotlin/sequences/plus, kotlin/sequences/plus, kotlin/text/plus]> = <Ambiguity: mapOf, [kotlin/collections/mapOf, kotlin/collections/mapOf]>#(String(-).R|kotlin/to|(Int(4))).<Ambiguity: plus, [kotlin/plus, kotlin/plus, kotlin/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/collections/plus, kotlin/sequences/plus, kotlin/sequences/plus, kotlin/sequences/plus, kotlin/text/plus]>#(<Ambiguity: mapOf, [kotlin/collections/mapOf, kotlin/collections/mapOf]>#(String(_).R|kotlin/to|(Int(5))))
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirResolveTestCaseWithStdlibGenerated.java

@@ -79,6 +79,11 @@ public class FirResolveTestCaseWithStdlibGenerated extends AbstractFirResolveTes
         runTest("compiler/fir/resolve/testData/resolve/stdlib/reflectionClass.kt");
     }
 
+    @TestMetadata("topLevelResolve.kt")
+    public void testTopLevelResolve() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/stdlib/topLevelResolve.kt");
+    }
+
     @TestMetadata("typeAliasDeserialization.kt")
     public void testTypeAliasDeserialization() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/stdlib/typeAliasDeserialization.kt");

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/model/NewConstraintSystemImpl.kt

@@ -179,8 +179,8 @@ class NewConstraintSystemImpl(
         return !type.contains {
             val capturedType = it.asSimpleType()?.asCapturedType()
             // TODO: change NewCapturedType to markered one for FE-IR
-            val typeToCheck = if (capturedType is NewCapturedType && capturedType.captureStatus() == CaptureStatus.FROM_EXPRESSION)
-                capturedType.constructor.projection.type
+            val typeToCheck = if (capturedType is CapturedTypeMarker && capturedType.captureStatus() == CaptureStatus.FROM_EXPRESSION)
+                capturedType.typeConstructorProjection().getType()
             else
                 it
 

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

@@ -43,14 +43,15 @@ FILE fqName:<root> fileName:/intersectionType1.kt
           GET_VAR 'y: kotlin.Any? declared in <root>.ofType' type=kotlin.Any? origin=null
   FUN name:test visibility:public modality:FINAL <> () returnType:kotlin.Unit 
     BLOCK_BODY
-      VAR name:a1 type:kotlin.Array<kotlin.Any?> [val]
-        CALL 'public final fun arrayOf (elements: kotlin.Array<out T of <uninitialized parent>>): kotlin.Array<T of <uninitialized parent>> [inline] declared in kotlin' type=kotlin.Array<kotlin.Any?> origin=null
+      VAR name:a1 type:kotlin.Array<<root>.In<kotlin.Int>> [val]
+        CALL 'public final fun arrayOf (elements: kotlin.Array<out T of <uninitialized parent>>): kotlin.Array<T of <uninitialized parent>> [inline] declared in kotlin' type=kotlin.Array<<root>.In<kotlin.Int>> origin=null
           elements: CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in <root>.In' type=<root>.In<kotlin.Int> origin=null
             <class: I>: <none>
-      VAR name:a2 type:kotlin.Array<kotlin.Any?> [val]
-        CALL 'public final fun arrayOf (elements: kotlin.Array<out T of <uninitialized parent>>): kotlin.Array<T of <uninitialized parent>> [inline] declared in kotlin' type=kotlin.Array<kotlin.Any?> origin=null
+      VAR name:a2 type:kotlin.Array<<root>.In<kotlin.String>> [val]
+        CALL 'public final fun arrayOf (elements: kotlin.Array<out T of <uninitialized parent>>): kotlin.Array<T of <uninitialized parent>> [inline] declared in kotlin' type=kotlin.Array<<root>.In<kotlin.String>> origin=null
           elements: CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in <root>.In' type=<root>.In<kotlin.String> origin=null
             <class: I>: <none>
-      ERROR_CALL 'Unresolved reference: <Inapplicable(INAPPLICABLE): [/foo]>#' type=IrErrorType
-        GET_VAR 'val a1: kotlin.Array<kotlin.Any?> [val] declared in <root>.test' type=kotlin.Array<kotlin.Any?> origin=null
-        GET_VAR 'val a2: kotlin.Array<kotlin.Any?> [val] declared in <root>.test' type=kotlin.Array<kotlin.Any?> origin=null
+      CALL 'public final fun foo <T> (a: kotlin.Array<<root>.In<T of <root>.foo>>, b: kotlin.Array<<root>.In<kotlin.String>>): kotlin.Boolean declared in <root>' type=kotlin.Boolean origin=null
+        <T>: <none>
+        a: GET_VAR 'val a1: kotlin.Array<<root>.In<kotlin.Int>> [val] declared in <root>.test' type=kotlin.Array<<root>.In<kotlin.Int>> origin=null
+        b: GET_VAR 'val a2: kotlin.Array<<root>.In<kotlin.String>> [val] declared in <root>.test' type=kotlin.Array<<root>.In<kotlin.String>> origin=null