FIR IDE: AddWhenRemainingBranchFix

The fix reuses logic that is already available from
FirWhenExhaustivenessTransformer to collect missing when branches. The
current logic unfortunately uses hackyAllowRunningOnEdt to shorten the
generated code.

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirWhenExhaustivenessTransformer.kt

@@ -33,6 +33,59 @@ class FirWhenExhaustivenessTransformer(private val bodyResolveComponents: BodyRe
             WhenOnEnumExhaustivenessChecker,
             WhenOnSealedClassExhaustivenessChecker
         )
+
+        @OptIn(ExperimentalStdlibApi::class)
+        fun computeAllMissingCases(session: FirSession, whenExpression: FirWhenExpression): List<WhenMissingCase> {
+            val subjectType = getSubjectType(session, whenExpression) ?: return emptyList()
+            return buildList {
+                for (type in subjectType.unwrapIntersectionType()) {
+                    val checkers = getCheckers(type, session)
+                    collectMissingCases(checkers, whenExpression, type, session)
+                }
+            }
+        }
+
+        private fun getSubjectType(session: FirSession, whenExpression: FirWhenExpression): ConeKotlinType? {
+            val subjectType = whenExpression.subjectVariable?.returnTypeRef?.coneType
+                ?: whenExpression.subject?.typeRef?.coneType
+                ?: return null
+
+            return subjectType.fullyExpandedType(session).lowerBoundIfFlexible()
+        }
+
+        private fun ConeKotlinType.unwrapIntersectionType(): Collection<ConeKotlinType> {
+            return (this as? ConeIntersectionType)?.intersectedTypes ?: listOf(this)
+        }
+
+
+        @OptIn(ExperimentalStdlibApi::class)
+        private fun getCheckers(
+            subjectType: ConeKotlinType,
+            session: FirSession
+        ): List<WhenExhaustivenessChecker> {
+            return buildList {
+                exhaustivenessCheckers.filterTo<WhenExhaustivenessChecker, MutableCollection<in WhenExhaustivenessChecker>>(this) {
+                    it.isApplicable(subjectType, session)
+                }
+                if (isNotEmpty() && subjectType.isMarkedNullable) {
+                    this.add(WhenOnNullableExhaustivenessChecker)
+                }
+            }
+        }
+
+        private fun MutableList<WhenMissingCase>.collectMissingCases(
+            checkers: List<WhenExhaustivenessChecker>,
+            whenExpression: FirWhenExpression,
+            subjectType: ConeKotlinType,
+            session: FirSession
+        ) {
+            for (checker in checkers) {
+                checker.computeMissingCases(whenExpression, subjectType, session, this)
+            }
+            if (isEmpty() && whenExpression.branches.isEmpty()) {
+                add(WhenMissingCase.Unknown)
+            }
+        }
     }
 
     override fun <E : FirElement> transformElement(element: E, data: Any?): E {
@@ -52,21 +105,17 @@ class FirWhenExhaustivenessTransformer(private val bodyResolveComponents: BodyRe
         }
 
         val session = bodyResolveComponents.session
-        val subjectType = (whenExpression.subjectVariable?.returnTypeRef?.coneType
-            ?: whenExpression.subject?.typeRef?.coneType)
-            ?.fullyExpandedType(session)?.lowerBoundIfFlexible()
-            ?: run {
-                whenExpression.replaceExhaustivenessStatus(ExhaustivenessStatus.NotExhaustive.NO_ELSE_BRANCH)
-                return
-            }
+        val subjectType = getSubjectType(session, whenExpression) ?: run {
+            whenExpression.replaceExhaustivenessStatus(ExhaustivenessStatus.NotExhaustive.NO_ELSE_BRANCH)
+            return
+        }
 
         if (whenExpression.branches.isEmpty() && subjectType.isNothing) {
             whenExpression.replaceExhaustivenessStatus(ExhaustivenessStatus.ExhaustiveAsNothing)
             return
         }
 
-        val unwrappedIntersectionTypes = (subjectType as? ConeIntersectionType)?.intersectedTypes ?: listOf(subjectType)
-
+        val unwrappedIntersectionTypes = subjectType.unwrapIntersectionType()
 
         var status: ExhaustivenessStatus = ExhaustivenessStatus.NotExhaustive.NO_ELSE_BRANCH
 
@@ -86,30 +135,18 @@ class FirWhenExhaustivenessTransformer(private val bodyResolveComponents: BodyRe
         whenExpression.replaceExhaustivenessStatus(status)
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
     private fun computeStatusForNonIntersectionType(
         unwrappedSubjectType: ConeKotlinType,
         session: FirSession,
         whenExpression: FirWhenExpression,
     ): ExhaustivenessStatus {
-        val checkers = buildList {
-            exhaustivenessCheckers.filterTo(this) { it.isApplicable(unwrappedSubjectType, session) }
-            if (isNotEmpty() && unwrappedSubjectType.isMarkedNullable) {
-                add(WhenOnNullableExhaustivenessChecker)
-            }
-        }
-
+        val checkers = getCheckers(unwrappedSubjectType, session)
         if (checkers.isEmpty()) {
             return ExhaustivenessStatus.NotExhaustive.NO_ELSE_BRANCH
         }
 
         val whenMissingCases = mutableListOf<WhenMissingCase>()
-        for (checker in checkers) {
-            checker.computeMissingCases(whenExpression, unwrappedSubjectType, session, whenMissingCases)
-        }
-        if (whenMissingCases.isEmpty() && whenExpression.branches.isEmpty()) {
-            whenMissingCases.add(WhenMissingCase.Unknown)
-        }
+        whenMissingCases.collectMissingCases(checkers, whenExpression, unwrappedSubjectType, session)
 
         return if (whenMissingCases.isEmpty()) {
             ExhaustivenessStatus.ProperlyExhaustive

idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/MainKtQuickFixRegistrar.kt

@@ -124,9 +124,6 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
         registerApplicator(TypeMismatchFactories.assignmentTypeMismatch)
         registerApplicator(TypeMismatchFactories.initializerTypeMismatch)
 
-        // TODO: NON_EXHAUSTIVE_WHEN[_ON_SEALED_CLASS] will be replaced in future. We need to register the fix for those diagnostics as well
-        registerPsiQuickFixes(KtFirDiagnostic.NoElseInWhen::class, AddWhenElseBranchFix)
-
         registerApplicator(WrapWithSafeLetCallFixFactories.forUnsafeCall)
         registerApplicator(WrapWithSafeLetCallFixFactories.forUnsafeImplicitInvokeCall)
         registerApplicator(WrapWithSafeLetCallFixFactories.forUnsafeInfixCall)
@@ -137,6 +134,12 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
         registerPsiQuickFixes(KtFirDiagnostic.InapplicableLateinitModifier::class, RemoveNullableFix.removeForLateInitProperty)
     }
 
+    private val whenStatements = KtQuickFixesListBuilder.registerPsiQuickFix {
+        // TODO: NON_EXHAUSTIVE_WHEN[_ON_SEALED_CLASS] will be replaced in future. We need to register the fix for those diagnostics as well
+        registerPsiQuickFixes(KtFirDiagnostic.NoElseInWhen::class, AddWhenElseBranchFix)
+        registerApplicator(AddWhenRemainingBranchFixFactories.noElseInWhen)
+    }
+
     private val typeMismatch = KtQuickFixesListBuilder.registerPsiQuickFix {
         registerApplicator(ChangeTypeQuickFixFactories.componentFunctionReturnTypeMismatch)
         registerApplicator(ChangeTypeQuickFixFactories.returnTypeMismatch)
@@ -156,6 +159,7 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
         imports,
         mutability,
         expressions,
+        whenStatements,
         typeMismatch
     )
 }

idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddWhenRemainingBranchFixFactories.kt

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.quickfix.fixes
+
+import org.jetbrains.kotlin.diagnostics.WhenMissingCase
+import org.jetbrains.kotlin.idea.KotlinBundle
+import org.jetbrains.kotlin.idea.api.applicator.HLApplicator
+import org.jetbrains.kotlin.idea.api.applicator.HLApplicatorInput
+import org.jetbrains.kotlin.idea.api.applicator.applicator
+import org.jetbrains.kotlin.idea.fir.api.fixes.HLQuickFix
+import org.jetbrains.kotlin.idea.fir.api.fixes.diagnosticFixFactory
+import org.jetbrains.kotlin.idea.frontend.api.analyse
+import org.jetbrains.kotlin.idea.frontend.api.components.ShortenOption
+import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassKind
+import org.jetbrains.kotlin.idea.frontend.api.tokens.HackToForceAllowRunningAnalyzeOnEDT
+import org.jetbrains.kotlin.idea.frontend.api.tokens.hackyAllowRunningOnEdt
+import org.jetbrains.kotlin.idea.util.generateWhenBranches
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.KtWhenExpression
+
+object AddWhenRemainingBranchFixFactories {
+    class Input(val whenMissingCases: List<WhenMissingCase>, val enumToStarImport: ClassId?) : HLApplicatorInput
+
+    val applicator: HLApplicator<KtWhenExpression, Input> = getApplicator(false)
+    val applicatorUsingStarImport: HLApplicator<KtWhenExpression, Input> = getApplicator(true)
+
+    @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class)
+    private fun getApplicator(useStarImport: Boolean = false) = applicator<KtWhenExpression, Input> {
+        familyAndActionName(
+            if (useStarImport) KotlinBundle.lazyMessage("fix.add.remaining.branches.with.star.import")
+            else KotlinBundle.lazyMessage("fix.add.remaining.branches")
+        )
+        applyTo { whenExpression, input ->
+            if (useStarImport) assert(input.enumToStarImport != null)
+            generateWhenBranches(whenExpression, input.whenMissingCases)
+            val shortenCommand = hackyAllowRunningOnEdt {
+                analyse(whenExpression) {
+                    collectPossibleReferenceShorteningsInElement(
+                        whenExpression,
+                        callableShortenOption = {
+                            if (useStarImport && it.callableIdIfNonLocal?.classId == input.enumToStarImport) {
+                                ShortenOption.SHORTEN_AND_STAR_IMPORT
+                            } else {
+                                ShortenOption.DO_NOT_SHORTEN
+                            }
+                        })
+
+                }
+            }
+            shortenCommand.invokeShortening()
+        }
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    val noElseInWhen = diagnosticFixFactory(KtFirDiagnostic.NoElseInWhen::class) { diagnostic ->
+        val whenExpression = diagnostic.psi
+        val subjectExpression = whenExpression.subjectExpression ?: return@diagnosticFixFactory emptyList()
+
+        buildList {
+            val missingCases = diagnostic.missingWhenCases.takeIf {
+                it.isNotEmpty() && it.singleOrNull() != WhenMissingCase.Unknown
+            } ?: return@buildList
+
+            add(HLQuickFix(whenExpression, Input(missingCases, null), applicator))
+            val baseClassSymbol = subjectExpression.getKtType().expandedClassSymbol ?: return@buildList
+            val enumToStarImport = baseClassSymbol.classIdIfNonLocal
+            if (baseClassSymbol.classKind == KtClassKind.ENUM_CLASS && enumToStarImport != null) {
+                add(HLQuickFix(whenExpression, Input(missingCases, enumToStarImport), applicatorUsingStarImport))
+            }
+        }
+    }
+}

idea/idea-fir/tests/org/jetbrains/kotlin/idea/fir/quickfix/HighLevelQuickFixTestGenerated.java

@@ -1946,6 +1946,21 @@ public class HighLevelQuickFixTestGenerated extends AbstractHighLevelQuickFixTes
             runTest("idea/testData/quickfix/when/addRemainingBranchesInNonDefaultPackage.kt");
         }
 
+        @TestMetadata("addRemainingBranchesMissingLeftBracket.kt")
+        public void testAddRemainingBranchesMissingLeftBracket() throws Exception {
+            runTest("idea/testData/quickfix/when/addRemainingBranchesMissingLeftBracket.kt");
+        }
+
+        @TestMetadata("addRemainingBranchesMissingRightBracket.kt")
+        public void testAddRemainingBranchesMissingRightBracket() throws Exception {
+            runTest("idea/testData/quickfix/when/addRemainingBranchesMissingRightBracket.kt");
+        }
+
+        @TestMetadata("addRemainingBranchesMissingRightParenthesis.kt")
+        public void testAddRemainingBranchesMissingRightParenthesis() throws Exception {
+            runTest("idea/testData/quickfix/when/addRemainingBranchesMissingRightParenthesis.kt");
+        }
+
         @TestMetadata("addRemainingBranchesSealed.kt")
         public void testAddRemainingBranchesSealed() throws Exception {
             runTest("idea/testData/quickfix/when/addRemainingBranchesSealed.kt");

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtExpressionInfoProvider.kt

@@ -5,14 +5,19 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.components
 
+import org.jetbrains.kotlin.diagnostics.WhenMissingCase
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.psi.KtReturnExpression
+import org.jetbrains.kotlin.psi.KtWhenExpression
 
 abstract class KtExpressionInfoProvider : KtAnalysisSessionComponent() {
     abstract fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtCallableSymbol?
+    abstract fun getWhenMissingCases(whenExpression: KtWhenExpression): List<WhenMissingCase>
 }
 
 interface KtExpressionInfoProviderMixIn : KtAnalysisSessionMixIn {
     fun KtReturnExpression.getReturnTargetSymbol(): KtCallableSymbol? =
         analysisSession.expressionInfoProvider.getReturnExpressionTargetSymbol(this)
+
+    fun KtWhenExpression.getMissingCases(): List<WhenMissingCase> = analysisSession.expressionInfoProvider.getWhenMissingCases(this)
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionInfoProvider.kt

@@ -5,13 +5,17 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
+import org.jetbrains.kotlin.diagnostics.WhenMissingCase
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
+import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
+import org.jetbrains.kotlin.fir.resolve.transformers.FirWhenExhaustivenessTransformer
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
-import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.components.KtExpressionInfoProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.psi.KtReturnExpression
+import org.jetbrains.kotlin.psi.KtWhenExpression
 
 internal class KtFirExpressionInfoProvider(
     override val analysisSession: KtFirAnalysisSession,
@@ -22,4 +26,9 @@ internal class KtFirExpressionInfoProvider(
         val firTargetSymbol = fir.target.labeledElement
         return firSymbolBuilder.callableBuilder.buildCallableSymbol(firTargetSymbol)
     }
+
+    override fun getWhenMissingCases(whenExpression: KtWhenExpression): List<WhenMissingCase> {
+        val firWhenExpression = whenExpression.getOrBuildFirSafe<FirWhenExpression>(analysisSession.firResolveState) ?: return emptyList()
+        return FirWhenExhaustivenessTransformer.computeAllMissingCases(analysisSession.firResolveState.rootModuleSession, firWhenExpression)
+    }
 }

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/util/psiUtils.kt

@@ -5,18 +5,17 @@
 
 package org.jetbrains.kotlin.idea.util
 
-import com.intellij.psi.PsiClass
-import com.intellij.psi.PsiElement
-import com.intellij.psi.PsiElementVisitor
-import com.intellij.psi.PsiJavaFile
+import com.intellij.psi.*
 import com.intellij.psi.impl.source.tree.LeafPsiElement
 import com.intellij.psi.util.parentOfType
 import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.cfg.containingDeclarationForPseudocode
+import org.jetbrains.kotlin.diagnostics.WhenMissingCase
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.renderer.render
 
 fun KtElement.getElementTextInContext(): String {
     val context = parentOfType<KtImportDirective>()
@@ -60,3 +59,44 @@ fun KtExpression.resultingWhens(): List<KtWhenExpression> = when (this) {
     is KtBlockExpression -> statements.lastOrNull()?.resultingWhens() ?: listOf()
     else -> listOf()
 }
+
+fun generateWhenBranches(element: KtWhenExpression, missingCases: List<WhenMissingCase>) {
+    val psiFactory = KtPsiFactory(element)
+    val whenCloseBrace = element.closeBrace ?: run {
+        val craftingMaterials = psiFactory.createExpression("when(1){}") as KtWhenExpression
+        if (element.rightParenthesis == null) {
+            element.addAfter(
+                craftingMaterials.rightParenthesis!!,
+                element.subjectExpression ?: throw AssertionError("caller should have checked the presence of subject expression.")
+            )
+        }
+        if (element.openBrace == null) {
+            element.addAfter(craftingMaterials.openBrace!!, element.rightParenthesis!!)
+        }
+        element.addAfter(craftingMaterials.closeBrace!!, element.entries.lastOrNull() ?: element.openBrace!!)
+        element.closeBrace!!
+    }
+    val elseBranch = element.entries.find { it.isElse }
+    (whenCloseBrace.prevSibling as? PsiWhiteSpace)?.replace(psiFactory.createNewLine())
+    for (case in missingCases) {
+        val branchConditionText = when (case) {
+            WhenMissingCase.Unknown,
+            WhenMissingCase.NullIsMissing,
+            is WhenMissingCase.BooleanIsMissing,
+            is WhenMissingCase.ConditionTypeIsExpect -> case.branchConditionText
+            is WhenMissingCase.IsTypeCheckIsMissing ->
+                if (case.isSingleton) {
+                    ""
+                } else {
+                    "is "
+                } + case.classId.asSingleFqName().render()
+            is WhenMissingCase.EnumCheckIsMissing -> case.callableId.asSingleFqName().render()
+        }
+        val entry = psiFactory.createWhenEntry("$branchConditionText -> TODO()")
+        if (elseBranch != null) {
+            element.addBefore(entry, elseBranch)
+        } else {
+            element.addBefore(entry, whenCloseBrace)
+        }
+    }
+}

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

@@ -19,18 +19,20 @@ package org.jetbrains.kotlin.idea.quickfix
 import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
-import com.intellij.psi.PsiWhiteSpace
-import org.jetbrains.kotlin.cfg.*
+import org.jetbrains.kotlin.cfg.WhenChecker
+import org.jetbrains.kotlin.cfg.hasUnknown
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.diagnostics.Diagnostic
-import org.jetbrains.kotlin.diagnostics.WhenMissingCase
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.ShortenReferences
-import org.jetbrains.kotlin.idea.core.quoteIfNeeded
 import org.jetbrains.kotlin.idea.intentions.ImportAllMembersIntention
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.idea.util.generateWhenBranches
+import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtWhenConditionWithExpression
+import org.jetbrains.kotlin.psi.KtWhenExpression
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 
 class AddWhenRemainingBranchesFix(
@@ -82,31 +84,7 @@ class AddWhenRemainingBranchesFix(
             if (element == null) return
             val missingCases = WhenChecker.getMissingCases(element, element.analyze())
 
-            val whenCloseBrace = element.closeBrace ?: throw AssertionError("isAvailable should check if close brace exist")
-            val elseBranch = element.entries.find { it.isElse }
-            val psiFactory = KtPsiFactory(element)
-            (whenCloseBrace.prevSibling as? PsiWhiteSpace)?.replace(psiFactory.createNewLine())
-            for (case in missingCases) {
-                val branchConditionText = when (case) {
-                    WhenMissingCase.Unknown,
-                    WhenMissingCase.NullIsMissing,
-                    is WhenMissingCase.BooleanIsMissing,
-                    is WhenMissingCase.ConditionTypeIsExpect -> case.branchConditionText
-                    is WhenMissingCase.IsTypeCheckIsMissing ->
-                        if (case.isSingleton) {
-                            ""
-                        } else {
-                            "is "
-                        } + case.classId.asSingleFqName().quoteIfNeeded().asString()
-                    is WhenMissingCase.EnumCheckIsMissing -> case.callableId.asSingleFqName().quoteIfNeeded().asString()
-                }
-                val entry = psiFactory.createWhenEntry("$branchConditionText -> TODO()")
-                if (elseBranch != null) {
-                    element.addBefore(entry, elseBranch)
-                } else {
-                    element.addBefore(entry, whenCloseBrace)
-                }
-            }
+            generateWhenBranches(element, missingCases)
 
             ShortenReferences.DEFAULT.process(element)
 

idea/testData/quickfix/when/addRemainingBranchesBlankLine.kt

@@ -9,4 +9,3 @@ fun test(a: A) {
 
     }
 }
-/* IGNORE_FIR */

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

@@ -9,4 +9,3 @@ fun test(a: A) {
         is B -> TODO()
     }
 }
-/* IGNORE_FIR */

idea/testData/quickfix/when/addRemainingBranchesBlankLineWithComment.kt

@@ -11,4 +11,3 @@ fun test(a: A) {
 
     }
 }
-/* IGNORE_FIR */

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

@@ -11,4 +11,3 @@ fun test(a: A) {
         is B -> TODO()
     }
 }
-/* IGNORE_FIR */

idea/testData/quickfix/when/addRemainingBranchesBoolean.kt

@@ -3,4 +3,3 @@
 fun test(b: Boolean) = wh<caret>en(b) {
     false -> 0
 }
-/* IGNORE_FIR */

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

@@ -4,4 +4,3 @@ fun test(b: Boolean) = when(b) {
     false -> 0
     true -> TODO()
 }
-/* IGNORE_FIR */

idea/testData/quickfix/when/addRemainingBranchesEnum.kt

@@ -5,4 +5,3 @@ enum class Color { R, G, B }
 fun test(c: Color) = wh<caret>en(c) {
     Color.B -> 0xff
 }
-/* IGNORE_FIR */

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

@@ -7,4 +7,3 @@ fun test(c: Color) = when(c) {
     Color.R -> TODO()
     Color.G -> TODO()
 }
-/* IGNORE_FIR */

idea/testData/quickfix/when/addRemainingBranchesEnumBackTicks.kt

@@ -8,4 +8,3 @@ enum class FooEnum {
 fun test(foo: FooEnum?) = <caret>when (foo) {
     FooEnum.A -> "A"
 }
-/* IGNORE_FIR */

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

@@ -14,4 +14,3 @@ fun test(foo: FooEnum?) = <caret>when (foo) {
     FooEnum.`null` -> TODO()
     null -> TODO()
 }
-/* IGNORE_FIR */

idea/testData/quickfix/when/addRemainingBranchesInNonDefaultPackage.kt

@@ -6,4 +6,3 @@ enum class Color { R, G, B }
 fun test(c: Color) = wh<caret>en(c) {
     Color.B -> 0xff
 }
-/* IGNORE_FIR */

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

@@ -8,4 +8,3 @@ fun test(c: Color) = when(c) {
     Color.R -> TODO()
     Color.G -> TODO()
 }
-/* IGNORE_FIR */

idea/testData/quickfix/when/addRemainingBranchesMissingLeftBracket.fir.kt

@@ -0,0 +1,8 @@
+// "Add remaining branches" "true"
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  val i = w<caret>hen (a)

idea/testData/quickfix/when/addRemainingBranchesMissingLeftBracket.fir.kt.after

@@ -0,0 +1,10 @@
+// "Add remaining branches" "true"
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  val i = w<caret>hen (a) {
+      is B -> TODO()
+  }

idea/testData/quickfix/when/addRemainingBranchesMissingLeftBracket.kt

@@ -0,0 +1,8 @@
+// "Add remaining branches" "false"
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  val i = w<caret>hen (a)

idea/testData/quickfix/when/addRemainingBranchesMissingRightBracket.fir.kt

@@ -0,0 +1,8 @@
+// "Add remaining branches" "true"
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  val i = w<caret>hen (a) {

idea/testData/quickfix/when/addRemainingBranchesMissingRightBracket.fir.kt.after

@@ -0,0 +1,10 @@
+// "Add remaining branches" "true"
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  val i = w<caret>hen (a) {
+      is B -> TODO()
+  }

idea/testData/quickfix/when/addRemainingBranchesMissingRightBracket.kt

@@ -0,0 +1,9 @@
+// "Add remaining branches" "false"
+// ACTION: Eliminate argument of 'when'
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  val i = w<caret>hen (a) {

idea/testData/quickfix/when/addRemainingBranchesMissingRightParenthesis.fir.kt

@@ -0,0 +1,8 @@
+// "Add remaining branches" "true"
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  val i = w<caret>hen (a

idea/testData/quickfix/when/addRemainingBranchesMissingRightParenthesis.fir.kt.after

@@ -0,0 +1,10 @@
+// "Add remaining branches" "true"
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  val i = w<caret>hen (a) {
+      is B -> TODO()
+  }

idea/testData/quickfix/when/addRemainingBranchesMissingRightParenthesis.kt

@@ -0,0 +1,8 @@
+// "Add remaining branches" "false"
+// WITH_RUNTIME
+
+sealed class A
+class B : A()
+
+fun test(a: A) {
+  <caret>when (a

idea/testData/quickfix/when/addRemainingBranchesSealed.kt

@@ -12,4 +12,3 @@ sealed class Variant {
 fun test(v: Variant?) = wh<caret>en(v) {
     Variant.Singleton -> "s"
 }
-/* IGNORE_FIR */

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

@@ -15,4 +15,3 @@ fun test(v: Variant?) = when(v) {
     is Variant.Something -> TODO()
     null -> TODO()
 }
-/* IGNORE_FIR */

idea/testData/quickfix/when/addRemainingBranchesSealedBackTicks.kt

@@ -13,4 +13,3 @@ object `null`: FooSealed()
 fun test(foo: FooSealed?) = <caret>when (foo) {
     A -> "A"
 }
-/* IGNORE_FIR */

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

@@ -20,4 +20,3 @@ fun test(foo: FooSealed?) = <caret>when (foo) {
     is `true` -> TODO()
     null -> TODO()
 }
-/* IGNORE_FIR */

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

@@ -15518,6 +15518,21 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/when/addRemainingBranchesInNonDefaultPackage.kt");
         }
 
+        @TestMetadata("addRemainingBranchesMissingLeftBracket.kt")
+        public void testAddRemainingBranchesMissingLeftBracket() throws Exception {
+            runTest("idea/testData/quickfix/when/addRemainingBranchesMissingLeftBracket.kt");
+        }
+
+        @TestMetadata("addRemainingBranchesMissingRightBracket.kt")
+        public void testAddRemainingBranchesMissingRightBracket() throws Exception {
+            runTest("idea/testData/quickfix/when/addRemainingBranchesMissingRightBracket.kt");
+        }
+
+        @TestMetadata("addRemainingBranchesMissingRightParenthesis.kt")
+        public void testAddRemainingBranchesMissingRightParenthesis() throws Exception {
+            runTest("idea/testData/quickfix/when/addRemainingBranchesMissingRightParenthesis.kt");
+        }
+
         @TestMetadata("addRemainingBranchesSealed.kt")
         public void testAddRemainingBranchesSealed() throws Exception {
             runTest("idea/testData/quickfix/when/addRemainingBranchesSealed.kt");