NI: clear calls info in coroutine inference before the second analysis of += right side

^KT-39376 Fixed

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

@@ -5,13 +5,16 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import org.jetbrains.kotlin.psi.KtDoubleColonExpression
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtReferenceExpression
+import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 import org.jetbrains.kotlin.resolve.BindingTrace
@@ -337,6 +340,19 @@ class CoroutineInferenceSession(
             deprecationResolver, moduleDescriptor, context.dataFlowValueFactory, typeApproximator, missingSupertypesResolver
         )
     }
+
+    /*
+     * It's used only for `+=` resolve to clear calls info before the second analysis of right side.
+     * TODO: remove it after moving `+=` resolve into OR mechanism
+     */
+    fun clearCallsInfoByContainingElement(containingElement: KtElement) {
+        commonCalls.removeIf remove@{ callInfo ->
+            val atom = callInfo.callResolutionResult.resultCallAtom.atom
+            if (atom !is PSIKotlinCallImpl) return@remove false
+
+            containingElement.anyDescendantOfType<KtElement> { it == atom.psiCall.callElement }
+        }
+    }
 }
 
 class ComposedSubstitutor(val left: NewTypeSubstitutor, val right: NewTypeSubstitutor) : NewTypeSubstitutor {

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingVisitorForStatements.java

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver;
 import org.jetbrains.kotlin.resolve.calls.context.CallPosition;
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency;
 import org.jetbrains.kotlin.resolve.calls.context.TemporaryTraceAndCache;
+import org.jetbrains.kotlin.resolve.calls.inference.CoroutineInferenceSession;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResults;
 import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResultsImpl;
@@ -245,6 +246,10 @@ public class ExpressionTypingVisitorForStatements extends ExpressionTypingVisito
         boolean lhsAssignable = basic.checkLValue(ignoreReportsTrace, context, left, right, expression, false);
         if (assignmentOperationType == null || lhsAssignable) {
             // Check for '+'
+            // We should clear calls info for coroutine inference within right side as here we analyze it a second time in another context
+            if (context.inferenceSession instanceof CoroutineInferenceSession) {
+                ((CoroutineInferenceSession) context.inferenceSession).clearCallsInfoByContainingElement(expression);
+            }
             Name counterpartName = OperatorConventions.BINARY_OPERATION_NAMES.get(OperatorConventions.ASSIGNMENT_OPERATION_COUNTERPARTS.get(operationType));
             binaryOperationDescriptors = components.callResolver.resolveBinaryCall(
                     context.replaceTraceAndCache(temporaryForBinaryOperation).replaceScope(scope),

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

@@ -0,0 +1,37 @@
+// !DIAGNOSTICS: -UNCHECKED_CAST -EXPERIMENTAL_API_USAGE_ERROR -UNUSED_PARAMETER
+
+class Bar
+
+fun <T> materialize() = null as T
+
+interface FlowCollector<in T> {
+    suspend fun emit(value: T)
+}
+
+interface Flow<T>
+
+public fun <T> flow(@BuilderInference block: suspend FlowCollector<T>.() -> Unit) = materialize<Flow<T>>()
+
+fun foo(total: Int, next: Int) = 10
+fun foo(total: Int, next: Float) = 10
+fun foo(total: Float, next: Int) = 10
+
+fun call(x: String) {}
+
+suspend fun foo(x: Int) = flow {
+    var y = 1
+    y += if (x > 2) 1 else 2
+
+    var newValue = 1
+    newValue += listOf<Int>().asSequence().fold(0) { total, next ->
+        call(<!CONSTANT_EXPECTED_TYPE_MISMATCH!>11<!>)
+        total + next
+    }
+    newValue += listOf<Int>().asSequence().fold(0, fun(total, next): Int { return total + next })
+    newValue += listOf<Int>().asSequence().fold(0, fun(total, next) = total + next)
+    newValue += listOf<Int>().asSequence().fold(0, ::foo)
+
+    emit(materialize<Bar>())
+
+    newValue += listOf<Int>().asSequence().fold(0) { total, next -> total + next }
+}

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

@@ -0,0 +1,29 @@
+package
+
+public fun call(/*0*/ x: kotlin.String): kotlin.Unit
+public fun </*0*/ T> flow(/*0*/ @kotlin.BuilderInference block: suspend FlowCollector<T>.() -> kotlin.Unit): Flow<T>
+public fun foo(/*0*/ total: kotlin.Float, /*1*/ next: kotlin.Int): kotlin.Int
+public suspend fun foo(/*0*/ x: kotlin.Int): Flow<Bar>
+public fun foo(/*0*/ total: kotlin.Int, /*1*/ next: kotlin.Float): kotlin.Int
+public fun foo(/*0*/ total: kotlin.Int, /*1*/ next: kotlin.Int): kotlin.Int
+public fun </*0*/ T> materialize(): T
+
+public final class Bar {
+    public constructor Bar()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public interface Flow</*0*/ T> {
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public interface FlowCollector</*0*/ in T> {
+    public abstract suspend fun emit(/*0*/ value: T): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

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

@@ -0,0 +1,42 @@
+// !DIAGNOSTICS: -UNUSED_ANONYMOUS_PARAMETER -UNUSED_PARAMETER
+
+fun foo(total: Int, next: Int) = 10
+fun foo(total: Int, next: Float) = 10
+fun foo(total: Float, next: Int) = 10
+
+fun call(x: String) {}
+
+fun foo(x: Float, y: Float) = {
+    var newValue = 1
+    newValue += listOf<Int>().asSequence().fold(0) { total, next ->
+        call(<!CONSTANT_EXPECTED_TYPE_MISMATCH!>11<!>)
+        total + next
+    }
+    newValue += listOf<Int>().asSequence().fold(0, fun(total, next): Int { return total + next })
+    newValue += listOf<Int>().asSequence().fold(0, fun(total, next) = total + next)
+    newValue += listOf<Int>().asSequence().fold(0, ::foo)
+}
+
+class A {
+    operator fun plus(x: Int) = A()
+    operator fun plusAssign(x: Float) {}
+}
+
+fun <T> id(x: T) = x
+
+fun foo2() = {
+    var x = A()
+    x += <!TYPE_MISMATCH!>{ "" }<!>
+    var y = A()
+    y += 1
+}
+
+class B {
+    operator fun plus(x: () -> String) = A()
+    operator fun plusAssign(x: () -> Int) {}
+}
+
+fun foo3(x: B) = {
+    x += { <!TYPE_MISMATCH, TYPE_MISMATCH, TYPE_MISMATCH!>""<!> }
+    x += id { <!TYPE_MISMATCH, TYPE_MISMATCH!>""<!> }
+}

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

@@ -0,0 +1,28 @@
+package
+
+public fun call(/*0*/ x: kotlin.String): kotlin.Unit
+public fun foo(/*0*/ x: kotlin.Float, /*1*/ y: kotlin.Float): () -> kotlin.Unit
+public fun foo(/*0*/ total: kotlin.Float, /*1*/ next: kotlin.Int): kotlin.Int
+public fun foo(/*0*/ total: kotlin.Int, /*1*/ next: kotlin.Float): kotlin.Int
+public fun foo(/*0*/ total: kotlin.Int, /*1*/ next: kotlin.Int): kotlin.Int
+public fun foo2(): () -> kotlin.Unit
+public fun foo3(/*0*/ x: B): () -> kotlin.Unit
+public fun </*0*/ T> id(/*0*/ x: T): T
+
+public final class A {
+    public constructor A()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final operator fun plus(/*0*/ x: kotlin.Int): A
+    public final operator fun plusAssign(/*0*/ x: kotlin.Float): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class B {
+    public constructor B()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final operator fun plus(/*0*/ x: () -> kotlin.String): A
+    public final operator fun plusAssign(/*0*/ x: () -> kotlin.Int): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

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

@@ -0,0 +1,20 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_ANONYMOUS_PARAMETER
+
+open class A
+
+operator fun <T> T.plus(x: (T) -> Int) = A()
+operator fun <T> T.plusAssign(x: (Int) -> T) {}
+
+fun foo(total: Int) = 10
+fun foo(total: Float) = 10
+fun foo(total: String) = 10
+
+fun <T> id(x: T) = x
+
+fun main() {
+    var newValue = A()
+    newValue += id { total -> A() }
+    newValue += id(fun(total) = A())
+    newValue += id(fun(total): A { return A() })
+    newValue += id(::foo)
+}

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

@@ -0,0 +1,16 @@
+package
+
+public fun foo(/*0*/ total: kotlin.Float): kotlin.Int
+public fun foo(/*0*/ total: kotlin.Int): kotlin.Int
+public fun foo(/*0*/ total: kotlin.String): kotlin.Int
+public fun </*0*/ T> id(/*0*/ x: T): T
+public fun main(): kotlin.Unit
+public operator fun </*0*/ T> T.plus(/*0*/ x: (T) -> kotlin.Int): A
+public operator fun </*0*/ T> T.plusAssign(/*0*/ x: (kotlin.Int) -> T): kotlin.Unit
+
+public open class A {
+    public constructor A()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

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

@@ -2062,6 +2062,21 @@ public class DiagnosticsTestWithStdLibGenerated extends AbstractDiagnosticsTestW
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/nestedSuspendCallInsideLambda.kt");
             }
 
+            @TestMetadata("plusAssignInCoroutineContext.kt")
+            public void testPlusAssignInCoroutineContext() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignInCoroutineContext.kt");
+            }
+
+            @TestMetadata("plusAssignWithLambda.kt")
+            public void testPlusAssignWithLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignWithLambda.kt");
+            }
+
+            @TestMetadata("plusAssignWithLambda2.kt")
+            public void testPlusAssignWithLambda2() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignWithLambda2.kt");
+            }
+
             @TestMetadata("qualifiedResolvedExpressionInsideBuilderInference.kt")
             public void testQualifiedResolvedExpressionInsideBuilderInference() throws Exception {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/qualifiedResolvedExpressionInsideBuilderInference.kt");

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

@@ -2062,6 +2062,21 @@ public class DiagnosticsTestWithStdLibUsingJavacGenerated extends AbstractDiagno
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/nestedSuspendCallInsideLambda.kt");
             }
 
+            @TestMetadata("plusAssignInCoroutineContext.kt")
+            public void testPlusAssignInCoroutineContext() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignInCoroutineContext.kt");
+            }
+
+            @TestMetadata("plusAssignWithLambda.kt")
+            public void testPlusAssignWithLambda() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignWithLambda.kt");
+            }
+
+            @TestMetadata("plusAssignWithLambda2.kt")
+            public void testPlusAssignWithLambda2() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/plusAssignWithLambda2.kt");
+            }
+
             @TestMetadata("qualifiedResolvedExpressionInsideBuilderInference.kt")
             public void testQualifiedResolvedExpressionInsideBuilderInference() throws Exception {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/qualifiedResolvedExpressionInsideBuilderInference.kt");