Introduce action to add missing when branches on sealed class

Made via diagnostics NON_EXHAUSTIVE_WHEN_FOR_SEALED_CLASS
with INFO severity and quick-fix
So #KT-17580 Fixed

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

@@ -808,6 +808,13 @@ class ControlFlowInformationProvider private constructor(
                                 trace.report(NON_EXHAUSTIVE_WHEN.on(element, enumMissingCases))
                             }
                         }
+                        val sealedClassDescriptor = WhenChecker.getClassDescriptorOfTypeIfSealed(subjectType)
+                        if (sealedClassDescriptor != null) {
+                            val sealedMissingCases = WhenChecker.getSealedMissingCases(element, context, sealedClassDescriptor)
+                            if (!sealedMissingCases.isEmpty()) {
+                                trace.report(NON_EXHAUSTIVE_WHEN_ON_SEALED_CLASS.on(element, sealedMissingCases))
+                            }
+                        }
                     }
                 }
             }

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

@@ -32,7 +32,6 @@ import org.jetbrains.kotlin.resolve.CompileTimeConstantUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 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.constants.CompileTimeConstant
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.types.KotlinType
@@ -275,6 +274,11 @@ object WhenChecker {
         return classDescriptor
     }
 
+    @JvmStatic
+    fun getClassDescriptorOfTypeIfSealed(type: KotlinType?): ClassDescriptor?
+            = type?.let { TypeUtils.getClassDescriptor(it) }?.takeIf { DescriptorUtils.isSealedClass(it) }
+
+
     @JvmStatic
     fun whenSubjectType(expression: KtWhenExpression, context: BindingContext) =
             expression.subjectExpression?.let { context.get(SMARTCAST, it)?.defaultType ?: context.getType(it) }
@@ -286,6 +290,13 @@ object WhenChecker {
             enumClassDescriptor: ClassDescriptor
     ) = WhenOnEnumExhaustivenessChecker.getMissingCases(expression, context, enumClassDescriptor, false)
 
+    @JvmStatic
+    fun getSealedMissingCases(
+            expression: KtWhenExpression,
+            context: BindingContext,
+            sealedClassDescriptor: ClassDescriptor
+    ) = WhenOnSealedExhaustivenessChecker.getMissingCases(expression, context, sealedClassDescriptor, false)
+
     fun getMissingCases(expression: KtWhenExpression, context: BindingContext): List<WhenMissingCase> {
         val type = whenSubjectType(expression, context) ?: return listOf(UnknownMissingCase)
         val nullable = type.isMarkedNullable

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/Errors.java

@@ -50,6 +50,7 @@ import java.util.Map;
 
 import static org.jetbrains.kotlin.diagnostics.PositioningStrategies.*;
 import static org.jetbrains.kotlin.diagnostics.Severity.ERROR;
+import static org.jetbrains.kotlin.diagnostics.Severity.INFO;
 import static org.jetbrains.kotlin.diagnostics.Severity.WARNING;
 
 /**
@@ -849,6 +850,9 @@ public interface Errors {
     DiagnosticFactory0<KtWhenEntry> REDUNDANT_ELSE_IN_WHEN = DiagnosticFactory0.create(WARNING, ELSE_ENTRY);
     DiagnosticFactory1<KtWhenExpression, List<WhenMissingCase>> NO_ELSE_IN_WHEN = DiagnosticFactory1.create(ERROR, WHEN_EXPRESSION);
     DiagnosticFactory1<KtWhenExpression, List<WhenMissingCase>> NON_EXHAUSTIVE_WHEN = DiagnosticFactory1.create(WARNING, WHEN_EXPRESSION);
+    DiagnosticFactory1<KtWhenExpression, List<WhenMissingCase>>
+            NON_EXHAUSTIVE_WHEN_ON_SEALED_CLASS = DiagnosticFactory1.create(INFO, WHEN_EXPRESSION);
+
     DiagnosticFactory0<PsiElement> COMMA_IN_WHEN_CONDITION_WITHOUT_ARGUMENT = DiagnosticFactory0.create(ERROR);
     DiagnosticFactory0<PsiElement> DUPLICATE_LABEL_IN_WHEN = DiagnosticFactory0.create(WARNING);
 

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/rendering/DefaultErrorMessages.java

@@ -520,6 +520,7 @@ public class DefaultErrorMessages {
 
         MAP.put(NO_ELSE_IN_WHEN, "''when'' expression must be exhaustive, add necessary {0}", RENDER_WHEN_MISSING_CASES);
         MAP.put(NON_EXHAUSTIVE_WHEN, "''when'' expression on enum is recommended to be exhaustive, add {0}", RENDER_WHEN_MISSING_CASES);
+        MAP.put(NON_EXHAUSTIVE_WHEN_ON_SEALED_CLASS, "''when'' expression on sealed classes is recommended to be exhaustive, add {0}", RENDER_WHEN_MISSING_CASES);
 
         MAP.put(TYPE_MISMATCH_IN_RANGE, "Type mismatch: incompatible types of range and element checked in it");
         MAP.put(CYCLIC_INHERITANCE_HIERARCHY, "There's a cycle in the inheritance hierarchy for this type");

compiler/testData/diagnostics/tests/when/NonExhaustiveWarningForSealedClass.kt

@@ -0,0 +1,16 @@
+sealed class S
+
+object First : S()
+
+class Derived(val s: String) : S()
+
+object Last : S()
+
+fun use(s: String) = s
+
+fun foo(s: S) {
+    <!NON_EXHAUSTIVE_WHEN_ON_SEALED_CLASS!>when<!> (s) {
+        First -> {}
+        is Derived -> use(<!DEBUG_INFO_SMARTCAST!>s<!>.s)
+    }
+}

compiler/testData/diagnostics/tests/when/NonExhaustiveWarningForSealedClass.txt

@@ -0,0 +1,33 @@
+package
+
+public fun foo(/*0*/ s: S): kotlin.Unit
+public fun use(/*0*/ s: kotlin.String): kotlin.String
+
+public final class Derived : S {
+    public constructor Derived(/*0*/ s: kotlin.String)
+    public final val s: kotlin.String
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public object First : S {
+    private constructor First()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public object Last : S {
+    private constructor Last()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public sealed class S {
+    private constructor S()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -23458,6 +23458,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("NonExhaustiveWarningForSealedClass.kt")
+            public void testNonExhaustiveWarningForSealedClass() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/when/NonExhaustiveWarningForSealedClass.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("NonExhaustiveWarningNull.kt")
             public void testNonExhaustiveWarningNull() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/when/NonExhaustiveWarningNull.kt");

idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/KotlinPsiChecker.kt

@@ -250,7 +250,7 @@ private class ElementAnnotator(private val element: PsiElement,
                         }
                 )
             }
-            Severity.INFO -> return // Do nothing
+            Severity.INFO -> AnnotationPresentationInfo(ranges, highlightType = ProblemHighlightType.INFORMATION)
         }
 
         setUpAnnotations(diagnostics, presentationInfo)
@@ -298,7 +298,7 @@ private class AnnotationPresentationInfo(
                     holder.createWarningAnnotation(range, defaultMessage)
                 }
             }
-            else -> throw IllegalArgumentException("Only ERROR and WARNING diagnostics are supported")
+            Severity.INFO -> holder.createInfoAnnotation(range, defaultMessage)
         }
 
         annotation.tooltip = getMessage(diagnostic)

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

@@ -241,6 +241,8 @@ class QuickFixRegistrar : QuickFixContributor {
         REDUNDANT_ELSE_IN_WHEN.registerFactory(RemoveWhenElseBranchFix)
         NON_EXHAUSTIVE_WHEN.registerFactory(AddWhenElseBranchFix)
         NON_EXHAUSTIVE_WHEN.registerFactory(AddWhenRemainingBranchesFix)
+        NON_EXHAUSTIVE_WHEN_ON_SEALED_CLASS.registerFactory(AddWhenElseBranchFix)
+        NON_EXHAUSTIVE_WHEN_ON_SEALED_CLASS.registerFactory(AddWhenRemainingBranchesFix)
         BREAK_OR_CONTINUE_IN_WHEN.registerFactory(AddLoopLabelFix)
 
         NO_TYPE_ARGUMENTS_ON_RHS.registerFactory(AddStarProjectionsFix.IsExpressionFactory)

idea/testData/quickfix/when/addRemainingBranchesSealedStatement.kt

@@ -0,0 +1,16 @@
+// "Add remaining branches" "true"
+// ERROR: Unresolved reference: TODO
+// ERROR: Unresolved reference: TODO
+// ERROR: Unresolved reference: TODO
+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/addRemainingBranchesSealedStatement.kt.after

@@ -0,0 +1,19 @@
+// "Add remaining branches" "true"
+// ERROR: Unresolved reference: TODO
+// ERROR: Unresolved reference: TODO
+// ERROR: Unresolved reference: TODO
+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 -> TODO()
+        Variant.Another -> TODO()
+        null -> TODO()
+    }
+}

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

@@ -10957,6 +10957,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("addRemainingBranchesSealedStatement.kt")
+        public void testAddRemainingBranchesSealedStatement() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/when/addRemainingBranchesSealedStatement.kt");
+            doTest(fileName);
+        }
+
         public void testAllFilesPresentInWhen() throws Exception {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/when"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
         }