Lift assignment: do not report in cases where argument are of different types #KT-21520 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 
 import org.jetbrains.kotlin.cfg.WhenChecker
+import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.branches
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -25,6 +26,8 @@ import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
+import org.jetbrains.kotlin.resolve.calls.callUtil.getType
+import org.jetbrains.kotlin.types.TypeConstructor
 import org.jetbrains.kotlin.types.typeUtil.isNothing
 
 object BranchedFoldingUtils {
@@ -52,8 +55,20 @@ object BranchedFoldingUtils {
                     it.getTargetLabel() == null
         }
 
-    private fun checkAssignmentsMatch(a1: KtBinaryExpression, a2: KtBinaryExpression): Boolean =
-        a1.left?.text == a2.left?.text && a1.operationToken == a2.operationToken
+    private fun KtBinaryExpression.checkAssignmentsMatch(other: KtBinaryExpression, rightTypeConstructor: TypeConstructor): Boolean {
+        return left?.text == other.left?.text
+                && operationToken == other.operationToken
+                && rightTypeConstructor == other.rightTypeConstructor()
+    }
+
+    private fun KtBinaryExpression.rightTypeConstructor(): TypeConstructor? {
+        val right = this.right ?: return null
+        val context = this.analyze()
+        val diagnostics = context.diagnostics
+        fun hasTypeMismatchError(e: KtExpression) = diagnostics.forElement(e).any { it.factory == Errors.TYPE_MISMATCH }
+        if (hasTypeMismatchError(this) || hasTypeMismatchError(right)) return null
+        return right.getType(context)?.constructor
+    }
 
     internal fun getFoldableAssignmentNumber(expression: KtExpression?): Int {
         expression ?: return -1
@@ -89,14 +104,15 @@ object BranchedFoldingUtils {
         }
         if (!collectAssignmentsAndCheck(expression)) return -1
         val firstAssignment = assignments.firstOrNull() ?: return 0
-        if (assignments.any { !BranchedFoldingUtils.checkAssignmentsMatch(it, firstAssignment) }) {
+        val rightTypeConstructor = firstAssignment.rightTypeConstructor() ?: return -1
+        if (assignments.any { !firstAssignment.checkAssignmentsMatch(it, rightTypeConstructor) }) {
             return -1
         }
         if (expression.anyDescendantOfType<KtBinaryExpression>(
                 predicate = {
                     if (it.operationToken in KtTokens.ALL_ASSIGNMENTS)
                         if (it.getNonStrictParentOfType<KtFinallySection>() != null)
-                            BranchedFoldingUtils.checkAssignmentsMatch(it, firstAssignment)
+                            firstAssignment.checkAssignmentsMatch(it, rightTypeConstructor)
                         else
                             it !in assignments
                     else

idea/testData/inspectionsLocal/liftOut/ifToAssignment/nullableMutableList.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+fun test(b: Boolean, x: Int, y: Int?) {
+    val list = mutableListOf<Int?>()
+    <caret>if (b) {
+        list += x
+    } else {
+        list += y
+    }
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/nullableMutableList.kt.after

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+fun test(b: Boolean, x: Int, y: Int?) {
+    val list = mutableListOf<Int?>()
+    list += if (b) {
+        x
+    } else {
+        y
+    }
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatch.kt

@@ -0,0 +1,9 @@
+// PROBLEM: none
+fun test(b: Boolean, x: Long, y: Int) {
+    var num: Long = 0L
+    <caret>if (b) {
+        num += x
+    } else {
+        num += y
+    }
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatch2.kt

@@ -0,0 +1,10 @@
+// PROBLEM: none
+// ERROR: Type mismatch: inferred type is Long? but Long was expected
+fun test(b: Boolean, x: Long, y: Long?) {
+    var num: Long = 0L
+    <caret>if (b) {
+        num += x
+    } else {
+        num += y
+    }
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatch3.kt

@@ -0,0 +1,13 @@
+// PROBLEM: none
+class C {
+    operator fun plusAssign(i: Int) {}
+    operator fun plusAssign(b: Boolean) {}
+}
+
+fun f(b: Boolean) {
+    val c = C()
+    <caret>if (b)
+        c += 1
+    else
+        c += true
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatchMutableList.kt

@@ -0,0 +1,10 @@
+// PROBLEM: none
+// WITH_RUNTIME
+fun test(b: Boolean) {
+    val list = mutableListOf<Int>()
+    <caret>if (b) {
+        list += 1
+    } else {
+        list += mutableListOf(2)
+    }
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatchMutableList2.kt

@@ -0,0 +1,12 @@
+// PROBLEM: none
+// ERROR: Type mismatch: inferred type is List<Any> but MutableList<Int> was expected
+// ERROR: Val cannot be reassigned
+// WITH_RUNTIME
+fun test(b: Boolean) {
+    val list = mutableListOf<Int>()
+    <caret>if (b) {
+        list += mutableListOf(1)
+    } else {
+        list += mutableListOf(2L)
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -2752,6 +2752,11 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/multipleAssignments.kt");
             }
 
+            @TestMetadata("nullableMutableList.kt")
+            public void testNullableMutableList() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/nullableMutableList.kt");
+            }
+
             @TestMetadata("simpleIf.kt")
             public void testSimpleIf() throws Exception {
                 runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/simpleIf.kt");
@@ -2791,6 +2796,31 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             public void testSimpleIfWithoutTerminatingAssignment() throws Exception {
                 runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/simpleIfWithoutTerminatingAssignment.kt");
             }
+
+            @TestMetadata("typeMismatch.kt")
+            public void testTypeMismatch() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatch.kt");
+            }
+
+            @TestMetadata("typeMismatch2.kt")
+            public void testTypeMismatch2() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatch2.kt");
+            }
+
+            @TestMetadata("typeMismatch3.kt")
+            public void testTypeMismatch3() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatch3.kt");
+            }
+
+            @TestMetadata("typeMismatchMutableList.kt")
+            public void testTypeMismatchMutableList() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatchMutableList.kt");
+            }
+
+            @TestMetadata("typeMismatchMutableList2.kt")
+            public void testTypeMismatchMutableList2() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/typeMismatchMutableList2.kt");
+            }
         }
 
         @TestMetadata("idea/testData/inspectionsLocal/liftOut/ifToReturn")