From 45ccec3b6472924bcf30c3e3fbdda2684a63e71b Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Tue, 15 Jun 2021 15:46:52 -0700
Subject: [PATCH] FIR IDE: Add HLAddWhenRemainingBranchesIntention

---
 .../generators/tests/idea/GenerateTests.kt    |  1 +
 .../HLAddWhenRemainingBranchesIntention.kt    | 34 +++++++++++++++++++
 .../after.kt.template                         | 12 +++++++
 .../before.kt.template                        | 10 ++++++
 .../description.html                          |  5 +++
 idea/resources-fir/META-INF/firIntentions.xml |  5 +++
 .../addWhenRemainingBranches/.firIntention    |  1 +
 7 files changed, 68 insertions(+)
 create mode 100644 idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/intentions/HLAddWhenRemainingBranchesIntention.kt
 create mode 100644 idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/after.kt.template
 create mode 100644 idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/before.kt.template
 create mode 100644 idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/description.html
 create mode 100644 idea/testData/intentions/addWhenRemainingBranches/.firIntention

diff --git a/generators/idea-generator/tests/org/jetbrains/kotlin/generators/tests/idea/GenerateTests.kt b/generators/idea-generator/tests/org/jetbrains/kotlin/generators/tests/idea/GenerateTests.kt
index 8ec6d22fd0b..b173c8a1b82 100644
--- a/generators/idea-generator/tests/org/jetbrains/kotlin/generators/tests/idea/GenerateTests.kt
+++ b/generators/idea-generator/tests/org/jetbrains/kotlin/generators/tests/idea/GenerateTests.kt
@@ -1092,6 +1092,7 @@ fun main(args: Array<String>) {
                 model("intentions/importAllMembers", pattern = pattern)
                 model("intentions/importMember", pattern = pattern)
                 model("intentions/convertToBlockBody", pattern = pattern)
+                model("intentions/addWhenRemainingBranches", pattern = pattern)
             }
 
             testClass<AbstractFirShortenRefsTest> {
diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/intentions/HLAddWhenRemainingBranchesIntention.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/intentions/HLAddWhenRemainingBranchesIntention.kt
new file mode 100644
index 00000000000..fa4c64eef98
--- /dev/null
+++ b/idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/intentions/HLAddWhenRemainingBranchesIntention.kt
@@ -0,0 +1,34 @@
+/*
+ * 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.fir.intentions
+
+import org.jetbrains.kotlin.diagnostics.WhenMissingCase
+import org.jetbrains.kotlin.idea.fir.api.AbstractHLIntention
+import org.jetbrains.kotlin.idea.fir.api.applicator.HLApplicabilityRange
+import org.jetbrains.kotlin.idea.fir.api.applicator.HLApplicatorInputProvider
+import org.jetbrains.kotlin.idea.fir.api.applicator.inputProvider
+import org.jetbrains.kotlin.idea.fir.applicators.ApplicabilityRanges
+import org.jetbrains.kotlin.idea.quickfix.fixes.AddWhenRemainingBranchFixFactories
+import org.jetbrains.kotlin.psi.KtWhenExpression
+
+class HLAddWhenRemainingBranchesIntention : AbstractHLIntention<KtWhenExpression, AddWhenRemainingBranchFixFactories.Input>(
+    KtWhenExpression::class,
+    AddWhenRemainingBranchFixFactories.applicator
+) {
+    override val applicabilityRange: HLApplicabilityRange<KtWhenExpression> get() = ApplicabilityRanges.SELF
+
+    override val inputProvider: HLApplicatorInputProvider<KtWhenExpression, AddWhenRemainingBranchFixFactories.Input>
+        get() = inputProvider { element ->
+            // TODO: consider removing the condition below so that this intention also works if there is no else. Currently we only offer
+            //  this intention if there is an `else` branch so that it behaves identically with FE1.0 (because FE1.0 reports warnings for
+            //  non-exhaustive when, which then results in a quickfix).
+            if (element.entries.none { it.isElse }) return@inputProvider null
+            val whenMissingCases = element.getMissingCases().takeIf {
+                it.isNotEmpty() && it.singleOrNull() != WhenMissingCase.Unknown
+            } ?: return@inputProvider null
+            AddWhenRemainingBranchFixFactories.Input(whenMissingCases, null)
+        }
+}
\ No newline at end of file
diff --git a/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/after.kt.template b/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/after.kt.template
new file mode 100644
index 00000000000..18050671af0
--- /dev/null
+++ b/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/after.kt.template
@@ -0,0 +1,12 @@
+enum class Entry {
+    FOO, BAR, BAZ
+}
+
+fun test(e: Entry): Int {
+    return when (e) {
+        <spot>Entry.FOO -> 1
+        Entry.BAR -> TODO()
+        Entry.BAZ -> TODO()
+        else -> 0</spot>
+    }
+}
\ No newline at end of file
diff --git a/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/before.kt.template b/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/before.kt.template
new file mode 100644
index 00000000000..1821a49ec2c
--- /dev/null
+++ b/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/before.kt.template
@@ -0,0 +1,10 @@
+enum class Entry {
+    FOO, BAR, BAZ
+}
+
+fun test(e: Entry): Int {
+    return when (e) {
+        <spot>Entry.FOO -> 1
+        else -> 0</spot>
+    }
+}
\ No newline at end of file
diff --git a/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/description.html b/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/description.html
new file mode 100644
index 00000000000..4680d1b1d19
--- /dev/null
+++ b/idea/resources-en/intentionDescriptions/HLAddWhenRemainingBranchesIntention/description.html
@@ -0,0 +1,5 @@
+<html>
+<body>
+This intention adds remaining branches on a <b>when</b> expression.
+</body>
+</html>
diff --git a/idea/resources-fir/META-INF/firIntentions.xml b/idea/resources-fir/META-INF/firIntentions.xml
index 202f474c8cd..6ef8a101747 100644
--- a/idea/resources-fir/META-INF/firIntentions.xml
+++ b/idea/resources-fir/META-INF/firIntentions.xml
@@ -35,5 +35,10 @@
       <className>org.jetbrains.kotlin.idea.fir.intentions.HLConvertToBlockBodyIntention</className>
       <category>Kotlin</category>
     </intentionAction>
+
+    <intentionAction>
+      <className>org.jetbrains.kotlin.idea.fir.intentions.HLAddWhenRemainingBranchesIntention</className>
+      <category>Kotlin</category>
+    </intentionAction>
   </extensions>
 </idea-plugin>
\ No newline at end of file
diff --git a/idea/testData/intentions/addWhenRemainingBranches/.firIntention b/idea/testData/intentions/addWhenRemainingBranches/.firIntention
new file mode 100644
index 00000000000..a7a4444043e
--- /dev/null
+++ b/idea/testData/intentions/addWhenRemainingBranches/.firIntention
@@ -0,0 +1 @@
+org.jetbrains.kotlin.idea.fir.intentions.HLAddWhenRemainingBranchesIntention
\ No newline at end of file