FIR: Support captured types for sub-arguments

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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.addSubtypeConstraintIfCompatible
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
+import org.jetbrains.kotlin.types.model.CaptureStatus
 
 
 fun Candidate.resolveArgumentExpression(
@@ -144,7 +145,7 @@ private fun Candidate.resolveBlockArgument(
     }
 }
 
-fun resolveSubCallArgument(
+fun Candidate.resolveSubCallArgument(
     csBuilder: ConstraintSystemBuilder,
     argument: FirResolvable,
     expectedType: ConeKotlinType,
@@ -169,7 +170,7 @@ fun resolveSubCallArgument(
     resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isDispatch, isSafeCall)
 }
 
-fun resolvePlainExpressionArgument(
+fun Candidate.resolvePlainExpressionArgument(
     csBuilder: ConstraintSystemBuilder,
     argument: FirExpression,
     expectedType: ConeKotlinType?,
@@ -184,7 +185,7 @@ fun resolvePlainExpressionArgument(
     resolvePlainArgumentType(csBuilder, argumentType, expectedType, sink, isReceiver, isDispatch, isSafeCall)
 }
 
-fun resolvePlainArgumentType(
+fun Candidate.resolvePlainArgumentType(
     csBuilder: ConstraintSystemBuilder,
     argumentType: ConeKotlinType,
     expectedType: ConeKotlinType,
@@ -195,15 +196,25 @@ fun resolvePlainArgumentType(
 ) {
     val position = SimpleConstraintSystemConstraintPosition //TODO
 
+    val capturedType = prepareCapturedType(argumentType, sink.components.session)
+
     val nullableExpectedType = expectedType.withNullability(ConeNullability.NULLABLE)
     if (isReceiver && isSafeCall) {
-        if (!isDispatch && !csBuilder.addSubtypeConstraintIfCompatible(argumentType, nullableExpectedType, position)) {
+        if (!isDispatch && !csBuilder.addSubtypeConstraintIfCompatible(capturedType, nullableExpectedType, position)) {
             sink.reportApplicability(CandidateApplicability.WRONG_RECEIVER) // TODO
         }
         return
     }
 
-    checkApplicabilityForArgumentType(csBuilder, argumentType, expectedType, position, isReceiver, isDispatch, nullableExpectedType, sink)
+    checkApplicabilityForArgumentType(csBuilder, capturedType, expectedType, position, isReceiver, isDispatch, nullableExpectedType, sink)
+}
+
+fun Candidate.prepareCapturedType(argumentType: ConeKotlinType, session: FirSession): ConeKotlinType {
+    if (argumentType.typeArguments.isEmpty() || argumentType !is ConeClassLikeType) return argumentType
+
+    return bodyResolveComponents.inferenceComponents.ctx.captureFromArguments(
+        argumentType, CaptureStatus.FROM_EXPRESSION
+    ) as? ConeKotlinType ?: argumentType
 }
 
 private fun checkApplicabilityForArgumentType(

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

@@ -9,10 +9,8 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
-import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintIncorporator
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintInjector
 import org.jetbrains.kotlin.resolve.calls.inference.components.ResultTypeResolver
@@ -22,8 +20,6 @@ import org.jetbrains.kotlin.types.AbstractTypeApproximator
 import org.jetbrains.kotlin.types.checker.KotlinTypeRefiner
 import org.jetbrains.kotlin.types.model.SimpleTypeMarker
 import org.jetbrains.kotlin.types.model.TypeSystemInferenceExtensionContextDelegate
-import org.jetbrains.kotlin.types.model.TypeVariableMarker
-import org.jetbrains.kotlin.types.model.TypeVariableTypeConstructorMarker
 
 fun ConeTypeContext.hasNullableSuperType(type: ConeKotlinType): Boolean {
     if (type is ConeClassLikeType) return false
@@ -50,7 +46,7 @@ class InferenceComponents(
     val returnTypeCalculator: ReturnTypeCalculator,
     val scopeSession: ScopeSession
 ) {
-    private val approximator = object : AbstractTypeApproximator(ctx) {
+    val approximator: AbstractTypeApproximator = object : AbstractTypeApproximator(ctx) {
         override fun createErrorType(message: String): SimpleTypeMarker {
             return ConeClassErrorType(message)
         }

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

@@ -116,7 +116,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
             } else {
                 val argumentExtensionReceiverValue = candidate.implicitExtensionReceiverValue
                 if (argumentExtensionReceiverValue != null && explicitReceiverKind.shouldBeCheckedAgainstImplicit()) {
-                    resolvePlainArgumentType(
+                    candidate.resolvePlainArgumentType(
                         candidate.csBuilder,
                         argumentType = argumentExtensionReceiverValue.type,
                         expectedType = candidate.substitutor.substituteOrSelf(expectedReceiverType.type),

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

@@ -22,11 +22,9 @@ import org.jetbrains.kotlin.fir.resolve.transformers.MapArguments
 import org.jetbrains.kotlin.fir.resolve.transformers.StoreType
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolve.typeFromCallee
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
-import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -39,6 +37,7 @@ import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
+import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.model.StubTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
@@ -97,7 +96,10 @@ class FirCallCompleter(
             val finalSubstitutor =
                 candidate.system.asReadOnlyStorage().buildAbstractResultingSubstitutor(inferenceComponents.ctx) as ConeSubstitutor
             return call.transformSingle(
-                FirCallCompletionResultsWriterTransformer(session, finalSubstitutor, returnTypeCalculator),
+                FirCallCompletionResultsWriterTransformer(
+                    session, finalSubstitutor, returnTypeCalculator,
+                    inferenceComponents.approximator
+                ),
                 null
             )
         }
@@ -126,7 +128,7 @@ class FirCallCompleter(
                     FirValueParameterImpl(
                         null,
                         session,
-                        FirResolvedTypeRefImpl(null, itType),
+                        FirResolvedTypeRefImpl(null, itType.approximateLambdaInputType()),
                         name,
                         FirVariableSymbol(name),
                         defaultValue = null,
@@ -141,9 +143,14 @@ class FirCallCompleter(
             val expectedReturnTypeRef = expectedReturnType?.let { lambdaArgument.returnTypeRef.resolvedTypeFromPrototype(it) }
 
             val newLambdaExpression = lambdaArgument.copy(
-                receiverTypeRef = receiverType?.let { lambdaArgument.receiverTypeRef?.resolvedTypeFromPrototype(it) },
+                receiverTypeRef = receiverType?.let {
+                    lambdaArgument.receiverTypeRef?.resolvedTypeFromPrototype(it.approximateLambdaInputType())
+                },
                 valueParameters = lambdaArgument.valueParameters.mapIndexed { index, parameter ->
-                    parameter.transformReturnTypeRef(StoreType, parameter.returnTypeRef.resolvedTypeFromPrototype(parameters[index]))
+                    parameter.transformReturnTypeRef(
+                        StoreType,
+                        parameter.returnTypeRef.resolvedTypeFromPrototype(parameters[index].approximateLambdaInputType())
+                    )
                     parameter
                 } + listOfNotNull(itParam),
                 returnTypeRef = expectedReturnTypeRef ?: noExpectedType
@@ -156,4 +163,9 @@ class FirCallCompleter(
             return returnArguments to InferenceSession.default
         }
     }
+
+    private fun ConeKotlinType.approximateLambdaInputType(): ConeKotlinType =
+        inferenceComponents.approximator.approximateToSuperType(
+            this, TypeApproximatorConfiguration.FinalApproximationAfterResolutionAndInference
+        ) as ConeKotlinType? ?: this
 }

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

@@ -29,12 +29,15 @@ import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.fir.types.impl.FirTypeProjectionWithVarianceImpl
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.types.AbstractTypeApproximator
+import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.Variance
 
 class FirCallCompletionResultsWriterTransformer(
     override val session: FirSession,
     private val finalSubstitutor: ConeSubstitutor,
-    private val typeCalculator: ReturnTypeCalculator
+    private val typeCalculator: ReturnTypeCalculator,
+    private val typeApproximator: AbstractTypeApproximator
 ) : FirAbstractTreeTransformer<Nothing?>(phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
 
     override fun transformQualifiedAccessExpression(
@@ -78,7 +81,11 @@ class FirCallCompletionResultsWriterTransformer(
         candidate: Candidate
     ): FirResolvedTypeRef {
         val initialType = candidate.substitutor.substituteOrNull(type)
-        val finalType = finalSubstitutor.substituteOrNull(initialType)
+        val finalType = finalSubstitutor.substituteOrNull(initialType)?.let { substitutedType ->
+            typeApproximator.approximateToSuperType(
+                substitutedType, TypeApproximatorConfiguration.FinalApproximationAfterResolutionAndInference
+            ) as ConeKotlinType? ?: substitutedType
+        }
 
         return withReplacedConeType(finalType)
     }

compiler/fir/resolve/testData/resolve/inference/receiverWithCapturedType.kt

@@ -12,7 +12,7 @@ fun <D : CallableDescriptor> ResolvedCall<D>.updateD(d: D): D {
 }
 
 fun test_1_1(resolvedCall: ResolvedCall<out CallableDescriptor>) {
-    resolvedCall.<!INAPPLICABLE_CANDIDATE!>getParameterForArgument<!>() // should be ok
+    resolvedCall.getParameterForArgument() // should be ok
 }
 
 fun test_1_2(resolvedCall: ResolvedCall<in CallableDescriptor>) {
@@ -33,4 +33,4 @@ fun test_2_2(resolvedCall: ResolvedCall<in CallableDescriptor>, d: CallableDescr
 
 fun test_2_3(resolvedCall: ResolvedCall<CallableDescriptor>, d: CallableDescriptor) {
     val x = resolvedCall.updateD(d) // should be ok
-}
+}

compiler/fir/resolve/testData/resolve/inference/receiverWithCapturedType.txt

@@ -15,19 +15,19 @@ FILE: receiverWithCapturedType.kt
         ^updateD R|<local>/d|
     }
     public final fun test_1_1(resolvedCall: R|ResolvedCall<out CallableDescriptor>|): R|kotlin/Unit| {
-        R|<local>/resolvedCall|.<Inapplicable(WRONG_RECEIVER): [/getParameterForArgument]>#()
+        R|<local>/resolvedCall|.R|/getParameterForArgument|<R|captured type: lowerType = null|>()
     }
     public final fun test_1_2(resolvedCall: R|ResolvedCall<in CallableDescriptor>|): R|kotlin/Unit| {
-        R|<local>/resolvedCall|.R|/getParameterForArgument|<R|CallableDescriptor|>()
+        R|<local>/resolvedCall|.R|/getParameterForArgument|<R|captured type: lowerType = CallableDescriptor|>()
     }
     public final fun test_1_3(resolvedCall: R|ResolvedCall<CallableDescriptor>|): R|kotlin/Unit| {
         R|<local>/resolvedCall|.R|/getParameterForArgument|<R|CallableDescriptor|>()
     }
     public final fun test_2_1(resolvedCall: R|ResolvedCall<out CallableDescriptor>|, d: R|CallableDescriptor|): R|kotlin/Unit| {
-        lval x: <ERROR TYPE REF: Inapplicable(WRONG_RECEIVER): [/updateD]> = R|<local>/resolvedCall|.<Inapplicable(WRONG_RECEIVER): [/updateD]>#(R|<local>/d|)
+        lval x: <ERROR TYPE REF: Inapplicable(INAPPLICABLE): [/updateD]> = R|<local>/resolvedCall|.<Inapplicable(INAPPLICABLE): [/updateD]>#(R|<local>/d|)
     }
     public final fun test_2_2(resolvedCall: R|ResolvedCall<in CallableDescriptor>|, d: R|CallableDescriptor|): R|kotlin/Unit| {
-        lval x: R|CallableDescriptor| = R|<local>/resolvedCall|.R|/updateD|<R|CallableDescriptor|>(R|<local>/d|)
+        lval x: R|CallableDescriptor| = R|<local>/resolvedCall|.R|/updateD|<R|captured type: lowerType = CallableDescriptor|>(R|<local>/d|)
     }
     public final fun test_2_3(resolvedCall: R|ResolvedCall<CallableDescriptor>|, d: R|CallableDescriptor|): R|kotlin/Unit| {
         lval x: R|CallableDescriptor| = R|<local>/resolvedCall|.R|/updateD|<R|CallableDescriptor|>(R|<local>/d|)

compiler/fir/resolve/testData/resolve/inference/simpleCapturedTypes.kt

@@ -0,0 +1,21 @@
+interface Ann {
+    fun foo()
+}
+
+interface KC<T> {
+    val x: T
+}
+fun <T> id(x: KC<T>): KC<T> = x
+fun <T> KC<T>.idR(): KC<T> = this
+val <T> KC<T>.idP: KC<T> get() = this
+
+private fun getSetterInfos(kc: KC<out Ann>) {
+    id(kc).x.foo()
+
+    kc.idR().x.foo()
+    kc.idP.x.foo()
+
+    val x1 = id(kc)
+    val x2 = kc.idR()
+    val x3 = kc.idP
+}

compiler/fir/resolve/testData/resolve/inference/simpleCapturedTypes.txt

@@ -0,0 +1,28 @@
+FILE: simpleCapturedTypes.kt
+    public abstract interface Ann : R|kotlin/Any| {
+        public abstract fun foo(): R|kotlin/Unit|
+
+    }
+    public abstract interface KC<T> : R|kotlin/Any| {
+        public abstract val x: R|T|
+            public get(): R|T|
+
+    }
+    public final fun <T> id(x: R|KC<T>|): R|KC<T>| {
+        ^id R|<local>/x|
+    }
+    public final fun <T> R|KC<T>|.idR(): R|KC<T>| {
+        ^idR this@R|/idR|
+    }
+    public final val <T> R|KC<T>|.idP: R|KC<T>|
+        public get(): R|KC<T>| {
+            ^ this@R|/idP|
+        }
+    private final fun getSetterInfos(kc: R|KC<out Ann>|): R|kotlin/Unit| {
+        R|/id|<R|captured type: lowerType = null|>(R|<local>/kc|).R|FakeOverride</KC.x: R|Ann|>|.R|/Ann.foo|()
+        R|<local>/kc|.R|/idR|<R|captured type: lowerType = null|>().R|FakeOverride</KC.x: R|Ann|>|.R|/Ann.foo|()
+        R|<local>/kc|.R|/idP|.R|FakeOverride</KC.x: R|Ann|>|.R|/Ann.foo|()
+        lval x1: R|KC<out Ann>| = R|/id|<R|captured type: lowerType = null|>(R|<local>/kc|)
+        lval x2: R|KC<out Ann>| = R|<local>/kc|.R|/idR|<R|captured type: lowerType = null|>()
+        lval x3: R|KC<out Ann>| = R|<local>/kc|.R|/idP|
+    }

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

@@ -722,6 +722,29 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         }
     }
 
+    @TestMetadata("compiler/fir/resolve/testData/resolve/inference")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Inference extends AbstractFirDiagnosticsTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInInference() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/fir/resolve/testData/resolve/inference"), Pattern.compile("^([^.]+)\\.kt$"), true);
+        }
+
+        @TestMetadata("receiverWithCapturedType.kt")
+        public void testReceiverWithCapturedType() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/inference/receiverWithCapturedType.kt");
+        }
+
+        @TestMetadata("simpleCapturedTypes.kt")
+        public void testSimpleCapturedTypes() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/inference/simpleCapturedTypes.kt");
+        }
+    }
+
     @TestMetadata("compiler/fir/resolve/testData/resolve/multifile")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
@@ -928,11 +951,6 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/problems/propertyFromJavaPlusAssign.kt");
         }
 
-        @TestMetadata("receiverWithCapturedType.kt")
-        public void testReceiverWithCapturedType() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolve/problems/receiverWithCapturedType.kt");
-        }
-
         @TestMetadata("samConversionInConstructorCall.kt")
         public void testSamConversionInConstructorCall() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/problems/samConversionInConstructorCall.kt");