From 4d2e3a2379e82b3403fc8849fbf16e30af12fd10 Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Fri, 9 Apr 2021 14:01:19 -0700
Subject: [PATCH] FIR: Fix resolving += to ignore + if it doesn't return
 correct type

This is the FIR fix of the FE change tracked by
https://youtrack.jetbrains.com/issue/KT-45503

Consider the following code:

```
fun foo() {
  var l = mutableListOf("")
  l += ""
}
```

The above code used to be considered invalid due to ambiguous `+=`. But
after KT-45503, it's now considered valid code because `plus` is not
really a valid option for this code. This is because the return type
of `plus` for `List` is `List`, which does not match the type of
variable `l`.

As for now, FIR rejects the above code due to ambiguous `+=`. This
change fixes that by also consider the return type when resolving `+=`.
---
 .../FirExpressionsResolveTransformer.kt       | 15 +++++++++
 .../plusAssignOnVarAndCollections.fir.kt      | 31 -------------------
 .../plusAssignOnVarAndCollections.kt          |  1 +
 3 files changed, 16 insertions(+), 31 deletions(-)
 delete mode 100644 compiler/testData/diagnostics/tests/operatorsOverloading/plusAssignOnVarAndCollections.fir.kt

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
index 12b040e5b0f..2469c4ec843 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
@@ -388,6 +388,20 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         val operatorCallReference = resolvedOperatorCall.calleeReference as? FirNamedReferenceWithCandidate
         val operatorIsResolvable = operatorCallReference?.isError == false
 
+        fun operatorReturnTypeMatches(candidate: Candidate): Boolean {
+            // After KT-45503, non-assign flavor of operator is checked more strictly: the return type must be assignable to the variable.
+            val operatorCallReturnType = resolvedOperatorCall.typeRef.coneType
+            val substitutor = candidate.system.currentStorage()
+                .buildAbstractResultingSubstitutor(candidate.system.typeSystemContext) as ConeSubstitutor
+            return AbstractTypeChecker.isSubtypeOf(
+                session.typeContext,
+                substitutor.substituteOrSelf(operatorCallReturnType),
+                leftArgument.typeRef.coneType
+            )
+        }
+        // following `!!` is safe since `operatorIsResolvable = true` implies `operatorCallReference != null`
+        val operatorReturnTypeMatches = operatorIsResolvable && operatorReturnTypeMatches(operatorCallReference!!.candidate)
+
         val lhsReference = leftArgument.toResolvedCallableReference()
         val lhsSymbol = lhsReference?.resolvedSymbol as? FirVariableSymbol<*>
         val lhsVariable = lhsSymbol?.fir
@@ -444,6 +458,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             !assignIsResolvable && !operatorIsResolvable -> chooseAssign()
             !assignIsResolvable && operatorIsResolvable -> chooseOperator()
             assignIsResolvable && !operatorIsResolvable -> chooseAssign()
+            assignIsResolvable && operatorIsResolvable && !operatorReturnTypeMatches -> chooseAssign()
             else -> reportAmbiguity()
         }
     }
diff --git a/compiler/testData/diagnostics/tests/operatorsOverloading/plusAssignOnVarAndCollections.fir.kt b/compiler/testData/diagnostics/tests/operatorsOverloading/plusAssignOnVarAndCollections.fir.kt
deleted file mode 100644
index 3be6956ba69..00000000000
--- a/compiler/testData/diagnostics/tests/operatorsOverloading/plusAssignOnVarAndCollections.fir.kt
+++ /dev/null
@@ -1,31 +0,0 @@
-// WITH_RUNTIME
-
-fun test1() {
-    var list = ArrayList<Int>()
-    <!ASSIGN_OPERATOR_AMBIGUITY!>list -= 2<!>
-}
-
-fun test2() {
-    var set = HashMap<Int, Int>()
-    <!ASSIGN_OPERATOR_AMBIGUITY!>set += 2 to 2<!>
-}
-
-fun test3() {
-    var set = HashSet<Int>()
-    <!ASSIGN_OPERATOR_AMBIGUITY!>set -= 2<!>
-}
-
-fun test4() {
-    var list = mutableListOf(1)
-    <!ASSIGN_OPERATOR_AMBIGUITY!>list += 2<!>
-}
-
-fun test5() {
-    var map = mutableMapOf(1 to 1)
-    <!ASSIGN_OPERATOR_AMBIGUITY!>map += 2 to 2<!>
-}
-
-fun test6() {
-    var set = mutableSetOf(1)
-    <!ASSIGN_OPERATOR_AMBIGUITY!>set += 2<!>
-}
\ No newline at end of file
diff --git a/compiler/testData/diagnostics/tests/operatorsOverloading/plusAssignOnVarAndCollections.kt b/compiler/testData/diagnostics/tests/operatorsOverloading/plusAssignOnVarAndCollections.kt
index 5eff7fc7f4e..a4aef632975 100644
--- a/compiler/testData/diagnostics/tests/operatorsOverloading/plusAssignOnVarAndCollections.kt
+++ b/compiler/testData/diagnostics/tests/operatorsOverloading/plusAssignOnVarAndCollections.kt
@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // WITH_RUNTIME
 
 fun test1() {