Quick fix "Add remaining when branches" with some tests

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

@@ -37,13 +37,19 @@ import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils.isEnumClass
 import org.jetbrains.kotlin.resolve.DescriptorUtils.isEnumEntry
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import java.util.*
 
-interface WhenMissingCase
+interface WhenMissingCase {
+
+    val branchConditionText: String
+}
 
 // Always must be first in the list
 private object UnknownMissingCase : WhenMissingCase {
     override fun toString() = "unknown"
+
+    override val branchConditionText = "else"
 }
 
 val List<WhenMissingCase>.hasUnknown: Boolean
@@ -61,7 +67,9 @@ private interface WhenExhaustivenessChecker {
 }
 
 private object NullMissingCase : WhenMissingCase {
-    override fun toString() = "null"
+    override fun toString() = branchConditionText
+
+    override val branchConditionText = "null"
 }
 
 private object WhenOnNullableExhaustivenessChecker : WhenExhaustivenessChecker {
@@ -88,7 +96,9 @@ private object WhenOnNullableExhaustivenessChecker : WhenExhaustivenessChecker {
 }
 
 private class BooleanMissingCase(val b: Boolean) : WhenMissingCase {
-    override fun toString() = b.toString()
+    override fun toString() = branchConditionText
+
+    override val branchConditionText = b.toString()
 }
 
 private object WhenOnBooleanExhaustivenessChecker : WhenExhaustivenessChecker {
@@ -121,6 +131,8 @@ private object WhenOnBooleanExhaustivenessChecker : WhenExhaustivenessChecker {
 
 private class ClassMissingCase(val descriptor: ClassDescriptor): WhenMissingCase {
     override fun toString() = descriptor.name.identifier.let { if (descriptor.kind.isSingleton) it else "is $it" }
+
+    override val branchConditionText = descriptor.fqNameSafe.asString().let { if (descriptor.kind.isSingleton) it else "is $it" }
 }
 
 private abstract class WhenOnClassExhaustivenessChecker : WhenExhaustivenessChecker {

idea/idea-analysis/src/org/jetbrains/kotlin/idea/KotlinBundle.properties

@@ -150,6 +150,8 @@ remove.val.var.from.parameter=Remove ''{0}'' from parameter
 add.override.to.equals.hashCode.toString=Add 'override' to equals, hashCode, toString in project
 add.when.else.branch.action.family.name=Add else branch
 add.when.else.branch.action=Add else branch
+add.when.remaining.branches.action.family.name=Add remaining branches
+add.when.remaining.branches.action=Add remaining branches
 move.when.else.branch.to.the.end.action=Move else branch to the end
 move.when.else.branch.to.the.end.family.name=Move else branch to the end
 change.to.property.name.family.name=Change to property name

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

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.quickfix
+
+import com.intellij.codeInsight.CodeInsightUtilCore
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.cfg.WhenChecker
+import org.jetbrains.kotlin.cfg.hasUnknown
+import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.psi.KtWhenEntry
+import org.jetbrains.kotlin.psi.KtWhenExpression
+
+class AddWhenRemainingBranchesFix(expression: KtWhenExpression) : KotlinQuickFixAction<KtWhenExpression>(expression) {
+
+    override fun getFamilyName() = "Add remaining branches"
+
+    override fun getText() = "Add remaining branches"
+
+    override fun isAvailable(project: Project, editor: Editor?, file: PsiFile): Boolean {
+        return super.isAvailable(project, editor, file) && element.closeBrace != null &&
+               !WhenChecker.getNecessaryCases(element, element.analyze()).hasUnknown
+    }
+
+    override fun invoke(project: Project, editor: Editor?, file: KtFile) {
+        val necessaryCases = WhenChecker.getNecessaryCases(element, element.analyze())
+
+        val whenCloseBrace = element.closeBrace
+        assert(whenCloseBrace != null) { "isAvailable should check if close brace exist" }
+        val psiFactory = KtPsiFactory(file)
+        var insertedBranch: PsiElement? = null
+
+        for (case in necessaryCases) {
+            val entry = psiFactory.createWhenEntry("${case.branchConditionText} -> throw AssertionError(\"\")")
+            insertedBranch = element.addBefore(entry, whenCloseBrace)
+            element.addAfter(psiFactory.createNewLine(), insertedBranch)
+
+        }
+
+        if (insertedBranch != null) {
+            val insertedWhenEntry = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(insertedBranch) as KtWhenEntry
+            val textRange = insertedWhenEntry.textRange
+            val indexOfOpenBrace = insertedWhenEntry.text.indexOf('{')
+            editor!!.caretModel.moveToOffset(textRange.startOffset + indexOfOpenBrace + 1)
+        }
+    }
+
+    companion object {
+        @JvmStatic
+        fun createFactory(): KotlinSingleIntentionActionFactory {
+            return object : KotlinSingleIntentionActionFactory() {
+                public override fun createAction(diagnostic: Diagnostic): KotlinQuickFixAction<out PsiElement>? {
+                    val element = diagnostic.psiElement
+                    val whenExpression = PsiTreeUtil.getParentOfType(element, KtWhenExpression::class.java, false) ?: return null
+                    return AddWhenRemainingBranchesFix(whenExpression)
+                }
+            }
+        }
+    }
+}

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

@@ -189,6 +189,7 @@ class QuickFixRegistrar : QuickFixContributor {
 
         ELSE_MISPLACED_IN_WHEN.registerFactory(MoveWhenElseBranchFix.createFactory())
         NO_ELSE_IN_WHEN.registerFactory(AddWhenElseBranchFix.createFactory())
+        NO_ELSE_IN_WHEN.registerFactory(AddWhenRemainingBranchesFix.createFactory())
         BREAK_OR_CONTINUE_IN_WHEN.registerFactory(AddLoopLabelFix)
 
         NO_TYPE_ARGUMENTS_ON_RHS.registerFactory(AddStarProjectionsFix.createFactoryForIsExpression())

idea/testData/quickfix/when/addRemainingBranchesBoolean.kt

@@ -0,0 +1,4 @@
+// "Add remaining branches" "true"
+fun test(b: Boolean) = wh<caret>en(b) {
+    false -> 0
+}

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

@@ -0,0 +1,5 @@
+// "Add remaining branches" "true"
+fun test(b: Boolean) = when(b) {
+    false -> 0
+    true -> throw AssertionError("")
+}

idea/testData/quickfix/when/addRemainingBranchesEnum.kt

@@ -0,0 +1,5 @@
+// "Add remaining branches" "true"
+enum class Color { R, G, B }
+fun test(c: Color) = wh<caret>en(c) {
+    Color.B -> 0xff
+}

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

@@ -0,0 +1,7 @@
+// "Add remaining branches" "true"
+enum class Color { R, G, B }
+fun test(c: Color) = when(c) {
+    Color.B -> 0xff
+    Color.R -> throw AssertionError("")
+    Color.G -> throw AssertionError("")
+}

idea/testData/quickfix/when/addRemainingBranchesSealed.kt

@@ -0,0 +1,11 @@
+// "Add remaining branches" "true"
+sealed class Variant {
+    object Singleton : Variant()
+
+    class Something(val x: Int) : Variant()
+
+    object Another : Variant()
+}
+fun test(v: Variant?) = wh<caret>en(v) {
+    Variant.Singleton -> "s"
+}

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

@@ -0,0 +1,14 @@
+// "Add remaining branches" "true"
+sealed class Variant {
+    object Singleton : Variant()
+
+    class Something(val x: Int) : Variant()
+
+    object Another : Variant()
+}
+fun test(v: Variant?) = when(v) {
+    Variant.Singleton -> "s"
+    is Variant.Something -> throw AssertionError("")
+    Variant.Another -> throw AssertionError("")
+    null -> throw AssertionError("")
+}

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

@@ -7568,6 +7568,24 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
     public static class When extends AbstractQuickFixTest {
+        @TestMetadata("addRemainingBranchesBoolean.kt")
+        public void testAddRemainingBranchesBoolean() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/addRemainingBranchesBoolean.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("addRemainingBranchesEnum.kt")
+        public void testAddRemainingBranchesEnum() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/addRemainingBranchesEnum.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("addRemainingBranchesSealed.kt")
+        public void testAddRemainingBranchesSealed() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/addRemainingBranchesSealed.kt");
+            doTest(fileName);
+        }
+
         public void testAllFilesPresentInWhen() throws Exception {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/when"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }