[NI] Use context for relevant nested call to run call checkers

#KT-24808 Fixed

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.name.FqNameUnsafe;
 import org.jetbrains.kotlin.psi.*;
+import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext;
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemCompleter;
 import org.jetbrains.kotlin.resolve.calls.model.PartialCallResolutionResult;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
@@ -118,6 +119,7 @@ public interface BindingContext {
 
     WritableSlice<Call, ResolvedCall<?>> RESOLVED_CALL = new BasicWritableSlice<>(DO_NOTHING);
     WritableSlice<Call, PartialCallResolutionResult> ONLY_RESOLVED_CALL = new BasicWritableSlice<>(DO_NOTHING);
+    WritableSlice<Call, BasicCallResolutionContext> PARTIAL_CALL_RESOLUTION_CONTEXT = new BasicWritableSlice<>(DO_NOTHING);
     WritableSlice<KtExpression, Call> DELEGATE_EXPRESSION_TO_PROVIDE_DELEGATE_CALL = new BasicWritableSlice<>(DO_NOTHING);
     WritableSlice<Call, TailRecursionKind> TAIL_RECURSION_CALL = Slices.createSimpleSlice();
     WritableSlice<KtElement, ConstraintSystemCompleter> CONSTRAINT_SYSTEM_COMPLETER = new BasicWritableSlice<>(DO_NOTHING);

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/inference/CoroutineInferenceSession.kt

@@ -198,7 +198,7 @@ class CoroutineInferenceSession(
         context: BasicCallResolutionContext
     ): ResolvedAtomCompleter {
         return ResolvedAtomCompleter(
-            resultSubstitutor, context.trace, context, kotlinToResolvedCallTransformer,
+            resultSubstitutor, context, kotlinToResolvedCallTransformer,
             expressionTypingServices, argumentTypeResolver, doubleColonExpressionResolver, builtIns,
             deprecationResolver, moduleDescriptor, context.dataFlowValueFactory
         )

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt

@@ -89,6 +89,8 @@ class KotlinToResolvedCallTransformer(
                     psiKotlinCall.psiCall
 
                 context.trace.record(BindingContext.ONLY_RESOLVED_CALL, psiCall, baseResolvedCall)
+                context.trace.record(BindingContext.PARTIAL_CALL_RESOLUTION_CONTEXT, psiCall, context)
+
                 context.inferenceSession.addPartialCallInfo(
                     PSIPartialCallInfo(baseResolvedCall, context, tracingStrategy)
                 )
@@ -114,7 +116,7 @@ class KotlinToResolvedCallTransformer(
 
                 val resultSubstitutor = baseResolvedCall.constraintSystem.buildResultingSubstitutor()
                 val ktPrimitiveCompleter = ResolvedAtomCompleter(
-                    resultSubstitutor, context.trace, context, this, expressionTypingServices, argumentTypeResolver,
+                    resultSubstitutor, context, this, expressionTypingServices, argumentTypeResolver,
                     doubleColonExpressionResolver, builtIns, deprecationResolver, moduleDescriptor, dataFlowValueFactory
                 )
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt

@@ -40,18 +40,18 @@ import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 
 class ResolvedAtomCompleter(
     private val resultSubstitutor: NewTypeSubstitutor,
-    private val trace: BindingTrace,
     private val topLevelCallContext: BasicCallResolutionContext,
     private val kotlinToResolvedCallTransformer: KotlinToResolvedCallTransformer,
     private val expressionTypingServices: ExpressionTypingServices,
     private val argumentTypeResolver: ArgumentTypeResolver,
     private val doubleColonExpressionResolver: DoubleColonExpressionResolver,
     private val builtIns: KotlinBuiltIns,
-    deprecationResolver: DeprecationResolver,
+    private val deprecationResolver: DeprecationResolver,
-    moduleDescriptor: ModuleDescriptor,
+    private val moduleDescriptor: ModuleDescriptor,
     private val dataFlowValueFactory: DataFlowValueFactory
 ) {
-    private val callCheckerContext = CallCheckerContext(topLevelCallContext, deprecationResolver, moduleDescriptor)
+    private val topLevelCallCheckerContext = CallCheckerContext(topLevelCallContext, deprecationResolver, moduleDescriptor)
+    private val topLevelTrace = topLevelCallCheckerContext.trace
 
     private fun complete(resolvedAtom: ResolvedAtom) {
         when (resolvedAtom) {
@@ -75,15 +75,28 @@ class ResolvedAtomCompleter(
 
         val resolvedCall = kotlinToResolvedCallTransformer.transformToResolvedCall<CallableDescriptor>(
             resolvedCallAtom,
-            trace,
+            topLevelTrace,
             resultSubstitutor,
             diagnostics
         )
-        kotlinToResolvedCallTransformer.bindAndReport(topLevelCallContext, trace, resolvedCall, diagnostics)
+
+        val resolutionContextForPartialCall =
+            topLevelCallContext.trace[BindingContext.PARTIAL_CALL_RESOLUTION_CONTEXT, resolvedCallAtom.atom.psiKotlinCall.psiCall]
+
+        val callCheckerContext = if (resolutionContextForPartialCall != null)
+            CallCheckerContext(
+                resolutionContextForPartialCall.replaceBindingTrace(topLevelTrace),
+                deprecationResolver,
+                moduleDescriptor
+            )
+        else
+            topLevelCallCheckerContext
+
+        kotlinToResolvedCallTransformer.bindAndReport(topLevelCallContext, topLevelTrace, resolvedCall, diagnostics)
         kotlinToResolvedCallTransformer.runCallCheckers(resolvedCall, callCheckerContext)
 
         val lastCall = if (resolvedCall is VariableAsFunctionResolvedCall) resolvedCall.functionCall else resolvedCall
-        kotlinToResolvedCallTransformer.runArgumentsChecks(topLevelCallContext, trace, lastCall as NewResolvedCallImpl<*>)
+        kotlinToResolvedCallTransformer.runArgumentsChecks(topLevelCallContext, topLevelTrace, lastCall as NewResolvedCallImpl<*>)
 
         return resolvedCall
     }
@@ -91,14 +104,14 @@ class ResolvedAtomCompleter(
     private fun completeLambda(lambda: ResolvedLambdaAtom) {
         val returnType = resultSubstitutor.safeSubstitute(lambda.returnType)
 
-        updateTraceForLambda(lambda, trace, returnType)
+        updateTraceForLambda(lambda, topLevelTrace, returnType)
 
         for (lambdaResult in lambda.resultArguments) {
             val resultValueArgument = lambdaResult as? PSIKotlinCallArgument ?: continue
             val newContext =
                 topLevelCallContext.replaceDataFlowInfo(resultValueArgument.dataFlowInfoAfterThisArgument)
                     .replaceExpectedType(returnType)
-                    .replaceBindingTrace(trace)
+                    .replaceBindingTrace(topLevelTrace)
 
             val argumentExpression = resultValueArgument.valueArgument.getArgumentExpression() ?: continue
             kotlinToResolvedCallTransformer.updateRecordedType(argumentExpression, newContext, true)
@@ -179,7 +192,7 @@ class ResolvedAtomCompleter(
         // write down type for callable reference expression
         val resultType = resultSubstitutor.safeSubstitute(callableCandidate.reflectionCandidateType, Variance.INVARIANT)
         argumentTypeResolver.updateResultArgumentTypeIfNotDenotable(
-            trace, expressionTypingServices.statementFilter,
+            topLevelTrace, expressionTypingServices.statementFilter,
             resultType,
             callableReferenceExpression
         )
@@ -195,7 +208,7 @@ class ResolvedAtomCompleter(
         val psiCall = CallMaker.makeCall(reference, explicitReceiver?.receiverValue, null, reference, emptyList())
 
         val tracing = TracingStrategyImpl.create(reference, psiCall)
-        val temporaryTrace = TemporaryBindingTrace.create(trace, "callable reference fake call")
+        val temporaryTrace = TemporaryBindingTrace.create(topLevelTrace, "callable reference fake call")
 
         val resolvedCall = ResolvedCallImpl(
             psiCall, callableCandidate.candidate, callableCandidate.dispatchReceiver?.receiver?.receiverValue,
@@ -204,9 +217,9 @@ class ResolvedAtomCompleter(
         )
         resolvedCall.setResultingSubstitutor(resultSubstitutor)
 
-        tracing.bindCall(trace, psiCall)
+        tracing.bindCall(topLevelTrace, psiCall)
-        tracing.bindReference(trace, resolvedCall)
+        tracing.bindReference(topLevelTrace, resolvedCall)
-        tracing.bindResolvedCall(trace, resolvedCall)
+        tracing.bindResolvedCall(topLevelTrace, resolvedCall)
 
         resolvedCall.setStatusToSuccess()
         resolvedCall.markCallAsCompleted()
@@ -225,8 +238,8 @@ class ResolvedAtomCompleter(
         }
 
         // TODO: probably we should also record key 'DATA_FLOW_INFO_BEFORE', see ExpressionTypingVisitorDispatcher.getTypeInfo
-        trace.recordType(callableReferenceExpression, resultType)
+        topLevelTrace.recordType(callableReferenceExpression, resultType)
-        trace.record(BindingContext.PROCESSED, callableReferenceExpression)
+        topLevelTrace.record(BindingContext.PROCESSED, callableReferenceExpression)
 
         doubleColonExpressionResolver.checkReferenceIsToAllowedMember(
             callableCandidate.candidate,
@@ -243,7 +256,7 @@ class ResolvedAtomCompleter(
             collectionLiteralArgument.expectedType?.let { resultSubstitutor.safeSubstitute(it) } ?: TypeUtils.NO_EXPECTED_TYPE
 
         val actualContext = context
-            .replaceBindingTrace(trace)
+            .replaceBindingTrace(topLevelTrace)
             .replaceExpectedType(expectedType)
             .replaceContextDependency(ContextDependency.INDEPENDENT)
 

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/nestedSuspendCallInsideLambda.kt

@@ -0,0 +1,9 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+suspend fun wrapUp2() {
+    withContext<Unit> {
+        other()
+    }
+}
+suspend fun <T> withContext(block: suspend () -> T) {}
+suspend fun <R> other(): R = TODO()

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/nestedSuspendCallInsideLambda.txt

@@ -0,0 +1,5 @@
+package
+
+public suspend fun </*0*/ R> other(): R
+public suspend fun </*0*/ T> withContext(/*0*/ block: suspend () -> T): kotlin.Unit
+public suspend fun wrapUp2(): kotlin.Unit

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestWithStdLibGenerated.java

@@ -1586,6 +1586,11 @@ public class DiagnosticsTestWithStdLibGenerated extends AbstractDiagnosticsTestW
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/kt15516.kt");
             }
 
+            @TestMetadata("nestedSuspendCallInsideLambda.kt")
+            public void testNestedSuspendCallInsideLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/nestedSuspendCallInsideLambda.kt");
+            }
+
             @TestMetadata("recursiveGenerators.kt")
             public void testRecursiveGenerators() throws Exception {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/recursiveGenerators.kt");

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsTestWithStdLibUsingJavacGenerated.java

@@ -1586,6 +1586,11 @@ public class DiagnosticsTestWithStdLibUsingJavacGenerated extends AbstractDiagno
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/kt15516.kt");
             }
 
+            @TestMetadata("nestedSuspendCallInsideLambda.kt")
+            public void testNestedSuspendCallInsideLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/nestedSuspendCallInsideLambda.kt");
+            }
+
             @TestMetadata("recursiveGenerators.kt")
             public void testRecursiveGenerators() throws Exception {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/recursiveGenerators.kt");