"Change type" quick fix: fix false negative in 'when' branch #KT-8820 Fixed

idea/ide-common/src/org/jetbrains/kotlin/idea/core/quickfix/QuickFixUtil.java

@@ -105,9 +105,9 @@ public class QuickFixUtil {
         return declaration instanceof KtParameter ? (KtParameter) declaration : null;
     }
 
-    private static boolean equalOrLastInThenOrElse(KtExpression thenOrElse, KtExpression expression) {
-        if (thenOrElse == expression) return true;
-        return thenOrElse instanceof KtBlockExpression && expression.getParent() == thenOrElse &&
+    private static boolean equalOrLastInBlock(KtExpression block, KtExpression expression) {
+        if (block == expression) return true;
+        return block instanceof KtBlockExpression && expression.getParent() == block &&
                PsiTreeUtil.getNextSiblingOfType(expression, KtExpression.class) == null;
     }
 
@@ -115,13 +115,30 @@ public class QuickFixUtil {
     public static KtIfExpression getParentIfForBranch(@Nullable KtExpression expression) {
         KtIfExpression ifExpression = PsiTreeUtil.getParentOfType(expression, KtIfExpression.class, true);
         if (ifExpression == null) return null;
-        if (equalOrLastInThenOrElse(ifExpression.getThen(), expression)
-            || equalOrLastInThenOrElse(ifExpression.getElse(), expression)) {
+        if (equalOrLastInBlock(ifExpression.getThen(), expression)
+            || equalOrLastInBlock(ifExpression.getElse(), expression)) {
             return ifExpression;
         }
         return null;
     }
 
+    @Nullable
+    private static KtWhenExpression getParentWhenForBranch(@Nullable KtExpression expression) {
+        KtWhenEntry whenEntry = PsiTreeUtil.getParentOfType(expression, KtWhenEntry.class, true);
+        if (whenEntry == null) return null;
+        KtExpression whenEntryExpression = whenEntry.getExpression();
+        if (whenEntryExpression == null) return null;
+        if (!equalOrLastInBlock(whenEntryExpression, expression)) return null;
+        return PsiTreeUtil.getParentOfType(whenEntry, KtWhenExpression.class, true);
+    }
+
+    @Nullable
+    private static KtExpression getParentForBranch(@Nullable KtExpression expression) {
+        KtExpression parent = getParentIfForBranch(expression);
+        if (parent != null) return parent;
+        return getParentWhenForBranch(expression);
+    }
+
     // Returns true iff parent's value always or sometimes is evaluable to child's value, e.g.
     // parent = (x), child = x;
     // parent = if (...) x else y, child = x;
@@ -141,7 +158,7 @@ public class QuickFixUtil {
                 child = (KtExpression) childParent;
                 continue;
             }
-            child = getParentIfForBranch(child);
+            child = getParentForBranch(child);
             if (child == null) return false;
         }
         return true;

idea/testData/quickfix/typeMismatch/when.kt

@@ -0,0 +1,7 @@
+// "Change return type of enclosing function 'test' to 'String'" "true"
+fun test(i: Int) {
+    return when (i) {
+        0 -> ""<caret>
+        else -> ""
+    }
+}

idea/testData/quickfix/typeMismatch/when.kt.after

@@ -0,0 +1,7 @@
+// "Change return type of enclosing function 'test' to 'String'" "true"
+fun test(i: Int): String {
+    return when (i) {
+        0 -> ""<caret>
+        else -> ""
+    }
+}

idea/testData/quickfix/typeMismatch/when2.kt

@@ -0,0 +1,11 @@
+// "Change return type of enclosing function 'test' to 'String'" "true"
+fun test(i: Int) {
+    return when (i) {
+        0 -> {
+            ""
+        }
+        else -> {
+            ""<caret>
+        }
+    }
+}

idea/testData/quickfix/typeMismatch/when2.kt.after

@@ -0,0 +1,11 @@
+// "Change return type of enclosing function 'test' to 'String'" "true"
+fun test(i: Int): String {
+    return when (i) {
+        0 -> {
+            ""
+        }
+        else -> {
+            ""<caret>
+        }
+    }
+}

idea/testData/quickfix/typeMismatch/when3.kt

@@ -0,0 +1,7 @@
+// "Change type of 's' to 'String'" "true"
+fun test(i: Int) {
+    val s: Int = when (i) {
+        0 -> ""<caret>
+        else -> ""
+    }
+}

idea/testData/quickfix/typeMismatch/when3.kt.after

@@ -0,0 +1,7 @@
+// "Change type of 's' to 'String'" "true"
+fun test(i: Int) {
+    val s: String = when (i) {
+        0 -> ""<caret>
+        else -> ""
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -11843,6 +11843,21 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/typeMismatch/unresolvableTypeParams.kt");
         }
 
+        @TestMetadata("when.kt")
+        public void testWhen() throws Exception {
+            runTest("idea/testData/quickfix/typeMismatch/when.kt");
+        }
+
+        @TestMetadata("when2.kt")
+        public void testWhen2() throws Exception {
+            runTest("idea/testData/quickfix/typeMismatch/when2.kt");
+        }
+
+        @TestMetadata("when3.kt")
+        public void testWhen3() throws Exception {
+            runTest("idea/testData/quickfix/typeMismatch/when3.kt");
+        }
+
         @TestMetadata("idea/testData/quickfix/typeMismatch/casts")
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)