Attach "add remaining branches" and "add else branch" fixes to NON_EXHAUSTIVE_WHEN warning #KT-12503 Fixed

compiler/frontend/src/org/jetbrains/kotlin/cfg/WhenChecker.kt

@@ -321,7 +321,7 @@ object WhenChecker {
         return true
     }
 
-    private fun getMissingCases(expression: KtWhenExpression, context: BindingContext): List<WhenMissingCase> {
+    fun getMissingCases(expression: KtWhenExpression, context: BindingContext): List<WhenMissingCase> {
         val type = whenSubjectType(expression, context) ?: return listOf(UnknownMissingCase)
         val nullable = !type.isFlexible() && isNullableTypeWithoutPossibleSmartCast(expression.subjectExpression, type, context)
         val checkers = exhaustivenessCheckers.filter { it.isApplicable(type) }

idea/src/org/jetbrains/kotlin/idea/quickfix/AddWhenRemainingBranchesFix.kt

@@ -37,25 +37,24 @@ class AddWhenRemainingBranchesFix(expression: KtWhenExpression) : KotlinQuickFix
 
     override fun isAvailable(project: Project, editor: Editor?, file: PsiFile): Boolean =
             super.isAvailable(project, editor, file) && element.closeBrace != null &&
-            with(WhenChecker.getNecessaryCases(element, element.analyze())) {
+            with(WhenChecker.getMissingCases(element, element.analyze())) {
                 isNotEmpty() && !hasUnknown
             }
 
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
-        val necessaryCases = WhenChecker.getNecessaryCases(element, element.analyze())
+        val missingCases = WhenChecker.getMissingCases(element, element.analyze())
 
-        val whenCloseBrace = element.closeBrace
-        assert(whenCloseBrace != null) { "isAvailable should check if close brace exist" }
+        val whenCloseBrace = element.closeBrace ?: throw AssertionError("isAvailable should check if close brace exist")
         val psiFactory = KtPsiFactory(file)
 
-        for (case in necessaryCases) {
+        for (case in missingCases) {
             val entry = psiFactory.createWhenEntry("${case.branchConditionText} -> TODO()")
             element.addBefore(entry, whenCloseBrace)
         }
     }
 
     companion object : KotlinSingleIntentionActionFactory() {
-        override fun createAction(diagnostic: Diagnostic): KotlinQuickFixAction<out PsiElement>? {
+        override fun createAction(diagnostic: Diagnostic): KotlinQuickFixAction<PsiElement>? {
             val whenExpression = diagnostic.psiElement.getNonStrictParentOfType<KtWhenExpression>() ?: return null
             return AddWhenRemainingBranchesFix(whenExpression)
         }

idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixRegistrar.kt

@@ -201,6 +201,8 @@ class QuickFixRegistrar : QuickFixContributor {
         ELSE_MISPLACED_IN_WHEN.registerFactory(MoveWhenElseBranchFix)
         NO_ELSE_IN_WHEN.registerFactory(AddWhenElseBranchFix)
         NO_ELSE_IN_WHEN.registerFactory(AddWhenRemainingBranchesFix)
+        NON_EXHAUSTIVE_WHEN.registerFactory(AddWhenElseBranchFix)
+        NON_EXHAUSTIVE_WHEN.registerFactory(AddWhenRemainingBranchesFix)
         BREAK_OR_CONTINUE_IN_WHEN.registerFactory(AddLoopLabelFix)
 
         NO_TYPE_ARGUMENTS_ON_RHS.registerFactory(AddStarProjectionsFix.IsExpressionFactory)

idea/testData/quickfix/when/addElseBranchEnumStatement.kt

@@ -0,0 +1,9 @@
+// "Add else branch" "true"
+enum class Color { R, G, B }
+fun use(c: Color) {
+    <caret>when (c) {
+        Color.R -> red()
+    }
+}
+
+fun red() {}

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

@@ -0,0 +1,11 @@
+// "Add else branch" "true"
+enum class Color { R, G, B }
+fun use(c: Color) {
+    when (c) {
+        Color.R -> red()
+        else -> {
+        }
+    }
+}
+
+fun red() {}

idea/testData/quickfix/when/addRemainingBranchesEnumStatement.kt

@@ -0,0 +1,11 @@
+// "Add remaining branches" "true"
+// ERROR: Unresolved reference: TODO
+// ERROR: Unresolved reference: TODO
+enum class Color { R, G, B }
+fun use(c: Color) {
+    <caret>when (c) {
+        Color.R -> red()
+    }
+}
+
+fun red() {}

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

@@ -0,0 +1,13 @@
+// "Add remaining branches" "true"
+// ERROR: Unresolved reference: TODO
+// ERROR: Unresolved reference: TODO
+enum class Color { R, G, B }
+fun use(c: Color) {
+    when (c) {
+        Color.R -> red()
+        Color.G -> TODO()
+        Color.B -> TODO()
+    }
+}
+
+fun red() {}

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

@@ -8532,6 +8532,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
     public static class When extends AbstractQuickFixTest {
+        @TestMetadata("addElseBranchEnumStatement.kt")
+        public void testAddElseBranchEnumStatement() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/addElseBranchEnumStatement.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("addRemainingBranchesBoolean.kt")
         public void testAddRemainingBranchesBoolean() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/addRemainingBranchesBoolean.kt");
@@ -8544,6 +8550,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("addRemainingBranchesEnumStatement.kt")
+        public void testAddRemainingBranchesEnumStatement() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/addRemainingBranchesEnumStatement.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("addRemainingBranchesInNonDefaultPackage.kt")
         public void testAddRemainingBranchesInNonDefaultPackage() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/addRemainingBranchesInNonDefaultPackage.kt");