From 9b5f754e759bd907fe4cca5fb8dfc620cc18cefa Mon Sep 17 00:00:00 2001
From: Mark Punzalan <markpun@google.com>
Date: Thu, 10 Jun 2021 08:02:50 +0000
Subject: [PATCH] FIR IDE: Enable add accessors quickfix for
 MUST_BE_INITIALIZED.

Also added diagnosticFixFactoriesFromIntentionActions() to share
quickfix-creation across multiple diagnostics with the same PSI.
---
 .../idea/fir/api/fixes/HLDiagnosticFixFactory.kt       | 10 ++++++++++
 .../kotlin/idea/quickfix/MainKtQuickFixRegistrar.kt    |  2 +-
 .../idea/quickfix/fixes/AddAccessorsFactories.kt       |  7 +++++--
 .../quickfix/addPropertyAccessors/varHasGetter.kt      |  3 +--
 .../addPropertyAccessors/varHasGetter.kt.after         |  3 +--
 .../quickfix/addPropertyAccessors/varHasSetter.kt      |  3 +--
 .../addPropertyAccessors/varHasSetter.kt.after         |  3 +--
 7 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/api/fixes/HLDiagnosticFixFactory.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/api/fixes/HLDiagnosticFixFactory.kt
index 00453e67f43..4e1b875ab52 100644
--- a/idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/api/fixes/HLDiagnosticFixFactory.kt
+++ b/idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/api/fixes/HLDiagnosticFixFactory.kt
@@ -89,6 +89,16 @@ fun <DIAGNOSTIC : KtDiagnosticWithPsi<*>> diagnosticFixFactoryFromIntentionActio
     return HLDiagnosticFixFactoryUsingQuickFixActionBase(diagnosticClass, createQuickFixes)
 }
 
+/**
+ * Returns a [Collection] of [HLDiagnosticFixFactory] that creates [IntentionAction]s from diagnostics that have the same type of
+ * [PsiElement].
+ */
+fun <DIAGNOSTIC_PSI : PsiElement, DIAGNOSTIC : KtDiagnosticWithPsi<DIAGNOSTIC_PSI>> diagnosticFixFactoriesFromIntentionActions(
+    vararg diagnosticClasses: KClass<out DIAGNOSTIC>,
+    createIntentionActions: KtAnalysisSession.(DIAGNOSTIC) -> List<IntentionAction>
+): Collection<HLDiagnosticFixFactory<out DIAGNOSTIC>> =
+    diagnosticClasses.map { diagnosticFixFactoryFromIntentionActions(it, createIntentionActions) }
+
 private class IntentionActionAsQuickFixWrapper<T : PsiElement>(val intentionAction: IntentionAction, element: T) :
     QuickFixActionBase<T>(element) {
     override fun getText(): String = intentionAction.text
diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/MainKtQuickFixRegistrar.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/MainKtQuickFixRegistrar.kt
index ffb1fcdefba..8eb0db29064 100644
--- a/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/MainKtQuickFixRegistrar.kt
+++ b/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/MainKtQuickFixRegistrar.kt
@@ -79,7 +79,7 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
         )
         registerApplicators(InitializePropertyQuickFixFactories.initializePropertyFactory)
         registerApplicator(AddLateInitFactory.addLateInitFactory)
-        registerApplicator(AddAccessorsFactories.addAccessorsToUninitializedProperty)
+        registerApplicators(AddAccessorsFactories.addAccessorsToUninitializedProperty)
     }
 
     private val overrides = KtQuickFixesListBuilder.registerPsiQuickFix {
diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddAccessorsFactories.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddAccessorsFactories.kt
index e8273f19973..9b7770e5632 100644
--- a/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddAccessorsFactories.kt
+++ b/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddAccessorsFactories.kt
@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.idea.quickfix.fixes
 
-import org.jetbrains.kotlin.idea.fir.api.fixes.diagnosticFixFactoryFromIntentionActions
+import org.jetbrains.kotlin.idea.fir.api.fixes.diagnosticFixFactoriesFromIntentionActions
 import org.jetbrains.kotlin.idea.fir.intentions.HLAddGetterAndSetterIntention
 import org.jetbrains.kotlin.idea.fir.intentions.HLAddGetterIntention
 import org.jetbrains.kotlin.idea.fir.intentions.HLAddSetterIntention
@@ -14,7 +14,10 @@ import org.jetbrains.kotlin.psi.KtProperty
 
 object AddAccessorsFactories {
     val addAccessorsToUninitializedProperty =
-        diagnosticFixFactoryFromIntentionActions(KtFirDiagnostic.MustBeInitializedOrBeAbstract::class) { diagnostic ->
+        diagnosticFixFactoriesFromIntentionActions(
+            KtFirDiagnostic.MustBeInitialized::class,
+            KtFirDiagnostic.MustBeInitializedOrBeAbstract::class
+        ) { diagnostic ->
             val property: KtProperty = diagnostic.psi
             val addGetter = property.getter == null
             val addSetter = property.isVar && property.setter == null
diff --git a/idea/testData/quickfix/addPropertyAccessors/varHasGetter.kt b/idea/testData/quickfix/addPropertyAccessors/varHasGetter.kt
index 23c9f8c080f..ead42a039e3 100644
--- a/idea/testData/quickfix/addPropertyAccessors/varHasGetter.kt
+++ b/idea/testData/quickfix/addPropertyAccessors/varHasGetter.kt
@@ -4,5 +4,4 @@ class Test {
         get() {
             return 1
         }
-}
-/* IGNORE_FIR */
\ No newline at end of file
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/addPropertyAccessors/varHasGetter.kt.after b/idea/testData/quickfix/addPropertyAccessors/varHasGetter.kt.after
index b0171ce9c96..8153f5e1411 100644
--- a/idea/testData/quickfix/addPropertyAccessors/varHasGetter.kt.after
+++ b/idea/testData/quickfix/addPropertyAccessors/varHasGetter.kt.after
@@ -5,5 +5,4 @@ class Test {
             return 1
         }
         set(value) {<caret>}
-}
-/* IGNORE_FIR */
\ No newline at end of file
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/addPropertyAccessors/varHasSetter.kt b/idea/testData/quickfix/addPropertyAccessors/varHasSetter.kt
index 411de17a597..1b4835daa07 100644
--- a/idea/testData/quickfix/addPropertyAccessors/varHasSetter.kt
+++ b/idea/testData/quickfix/addPropertyAccessors/varHasSetter.kt
@@ -3,5 +3,4 @@
 class Test {
     var x: Int<caret>
         set(value) {}
-}
-/* IGNORE_FIR */
\ No newline at end of file
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/addPropertyAccessors/varHasSetter.kt.after b/idea/testData/quickfix/addPropertyAccessors/varHasSetter.kt.after
index df74bb2924b..bfc6afd8468 100644
--- a/idea/testData/quickfix/addPropertyAccessors/varHasSetter.kt.after
+++ b/idea/testData/quickfix/addPropertyAccessors/varHasSetter.kt.after
@@ -6,5 +6,4 @@ class Test {
             <caret>TODO()
         }
         set(value) {}
-}
-/* IGNORE_FIR */
\ No newline at end of file
+}
\ No newline at end of file