FIR IDE: Add diagnosticFixFactories() API to share quickfix-creation

across multiple diagnostics with the same PSI.

idea/idea-fir/src/org/jetbrains/kotlin/idea/core/overrideImplement/KtImplementMembersHandler.kt

@@ -120,19 +120,19 @@ internal class KtImplementAsConstructorParameterQuickfix(private val members: Co
 
 object MemberNotImplementedQuickfixFactories {
 
-    val abstractMemberNotImplemented = diagnosticFixFactory<KtFirDiagnostic.AbstractMemberNotImplemented> { diagnostic ->
+    val abstractMemberNotImplemented = diagnosticFixFactory(KtFirDiagnostic.AbstractMemberNotImplemented::class) { diagnostic ->
         getUnimplementedMemberFixes(diagnostic.psi)
     }
 
-    val abstractClassMemberNotImplemented = diagnosticFixFactory<KtFirDiagnostic.AbstractClassMemberNotImplemented> { diagnostic ->
+    val abstractClassMemberNotImplemented = diagnosticFixFactory(KtFirDiagnostic.AbstractClassMemberNotImplemented::class) { diagnostic ->
         getUnimplementedMemberFixes(diagnostic.psi)
     }
 
-    val manyInterfacesMemberNotImplemented = diagnosticFixFactory<KtFirDiagnostic.ManyInterfacesMemberNotImplemented> { diagnostic ->
+    val manyInterfacesMemberNotImplemented = diagnosticFixFactory(KtFirDiagnostic.ManyInterfacesMemberNotImplemented::class) { diagnostic ->
         getUnimplementedMemberFixes(diagnostic.psi)
     }
 
-    val manyImplMemberNotImplemented = diagnosticFixFactory<KtFirDiagnostic.ManyImplMemberNotImplemented> { diagnostic ->
+    val manyImplMemberNotImplemented = diagnosticFixFactory(KtFirDiagnostic.ManyImplMemberNotImplemented::class) { diagnostic ->
         getUnimplementedMemberFixes(diagnostic.psi, false)
     }
 

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/api/fixes/HLDiagnosticFixFactory.kt

@@ -11,20 +11,24 @@ import org.jetbrains.kotlin.idea.api.applicator.HLApplicator
 import org.jetbrains.kotlin.idea.api.applicator.HLApplicatorInput
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.diagnostics.KtDiagnosticWithPsi
+import kotlin.reflect.KClass
 
-sealed class HLDiagnosticFixFactory<in DIAGNOSTIC : KtDiagnosticWithPsi<*>> {
+sealed class HLDiagnosticFixFactory<DIAGNOSTIC : KtDiagnosticWithPsi<*>> {
     abstract fun KtAnalysisSession.createQuickFixes(diagnostic: DIAGNOSTIC): List<IntentionAction>
+    abstract val diagnosticClass: KClass<DIAGNOSTIC>
 }
 
 private class HLDiagnosticFixFactoryWithFixedApplicator<DIAGNOSTIC : KtDiagnosticWithPsi<*>, TARGET_PSI : PsiElement, INPUT : HLApplicatorInput>(
+    override val diagnosticClass: KClass<DIAGNOSTIC>,
     private val applicator: HLApplicator<TARGET_PSI, INPUT>,
-    private val createTargets: KtAnalysisSession.(DIAGNOSTIC) -> List<HLApplicatorTargetWithInput<TARGET_PSI, INPUT>>
+    private val createTargets: KtAnalysisSession.(DIAGNOSTIC) -> List<HLApplicatorTargetWithInput<TARGET_PSI, INPUT>>,
 ) : HLDiagnosticFixFactory<DIAGNOSTIC>() {
     override fun KtAnalysisSession.createQuickFixes(diagnostic: DIAGNOSTIC): List<HLQuickFix<TARGET_PSI, INPUT>> =
         createTargets.invoke(this, diagnostic).map { (target, input) -> HLQuickFix(target, input, applicator) }
 }
 
 private class HLDiagnosticFixFactoryWithVariableApplicator<DIAGNOSTIC : KtDiagnosticWithPsi<*>>(
+    override val diagnosticClass: KClass<DIAGNOSTIC>,
     private val createQuickFixes: KtAnalysisSession.(DIAGNOSTIC) -> List<IntentionAction>
 ) : HLDiagnosticFixFactory<DIAGNOSTIC>() {
     override fun KtAnalysisSession.createQuickFixes(diagnostic: DIAGNOSTIC): List<IntentionAction> =
@@ -41,15 +45,26 @@ internal fun <DIAGNOSTIC : KtDiagnosticWithPsi<PsiElement>> KtAnalysisSession.cr
  * The targets and inputs are consumed by the given applicator to apply fixes.
  */
 fun <DIAGNOSTIC : KtDiagnosticWithPsi<*>, TARGET_PSI : PsiElement, INPUT : HLApplicatorInput> diagnosticFixFactory(
+    diagnosticClass: KClass<DIAGNOSTIC>,
     applicator: HLApplicator<TARGET_PSI, INPUT>,
     createTargets: KtAnalysisSession.(DIAGNOSTIC) -> List<HLApplicatorTargetWithInput<TARGET_PSI, INPUT>>
 ): HLDiagnosticFixFactory<DIAGNOSTIC> =
-    HLDiagnosticFixFactoryWithFixedApplicator(applicator, createTargets)
+    HLDiagnosticFixFactoryWithFixedApplicator(diagnosticClass, applicator, createTargets)
 
 /**
  * Returns a [HLDiagnosticFixFactory] that creates [HLQuickFix]es from a diagnostic.
  */
 fun <DIAGNOSTIC : KtDiagnosticWithPsi<*>> diagnosticFixFactory(
+    diagnosticClass: KClass<DIAGNOSTIC>,
     createQuickFixes: KtAnalysisSession.(DIAGNOSTIC) -> List<IntentionAction>
 ): HLDiagnosticFixFactory<DIAGNOSTIC> =
-    HLDiagnosticFixFactoryWithVariableApplicator(createQuickFixes)
+    HLDiagnosticFixFactoryWithVariableApplicator(diagnosticClass, createQuickFixes)
+
+/**
+ * Returns a [Collection] of [HLDiagnosticFixFactory] that creates [HLQuickFix]es from diagnostics that have the same type of [PsiElement].
+ */
+fun <DIAGNOSTIC_PSI : PsiElement, DIAGNOSTIC : KtDiagnosticWithPsi<DIAGNOSTIC_PSI>> diagnosticFixFactories(
+    vararg diagnosticClasses: KClass<out DIAGNOSTIC>,
+    createQuickFixes: KtAnalysisSession.(DIAGNOSTIC) -> List<IntentionAction>
+): Collection<HLDiagnosticFixFactory<out DIAGNOSTIC>> =
+    diagnosticClasses.map { HLDiagnosticFixFactoryWithVariableApplicator(it, createQuickFixes) }

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/api/fixes/KtQuickFixesList.kt

@@ -65,13 +65,6 @@ class KtQuickFixesListBuilder private constructor() {
         }
     }
 
-    @OptIn(PrivateForInline::class)
-    inline fun <reified DIAGNOSTIC : KtDiagnosticWithPsi<*>> registerApplicator(
-        quickFixFactory: HLDiagnosticFixFactory<DIAGNOSTIC>
-    ) {
-        registerApplicator(DIAGNOSTIC::class, quickFixFactory)
-    }
-
     @PrivateForInline
     fun <DIAGNOSTIC_PSI : PsiElement, DIAGNOSTIC : KtDiagnosticWithPsi<DIAGNOSTIC_PSI>> registerPsiQuickFix(
         diagnosticClass: KClass<DIAGNOSTIC>,
@@ -80,13 +73,18 @@ class KtQuickFixesListBuilder private constructor() {
         quickFixes.getOrPut(diagnosticClass) { mutableListOf() }.add(HLQuickFixFactory.HLQuickFixesPsiBasedFactory(quickFixFactory))
     }
 
-
+    @OptIn(PrivateForInline::class)
-    @PrivateForInline
+    fun <DIAGNOSTIC : KtDiagnosticWithPsi<*>> registerApplicators(
-    fun <DIAGNOSTIC : KtDiagnosticWithPsi<*>> registerApplicator(
+        quickFixFactories: Collection<HLDiagnosticFixFactory<out DIAGNOSTIC>>
-        diagnosticClass: KClass<DIAGNOSTIC>,
-        quickFixFactory: HLDiagnosticFixFactory<DIAGNOSTIC>
     ) {
-        quickFixes.getOrPut(diagnosticClass) { mutableListOf() }
+        quickFixFactories.forEach(::registerApplicator)
+    }
+
+    @OptIn(PrivateForInline::class)
+    fun <DIAGNOSTIC : KtDiagnosticWithPsi<*>> registerApplicator(
+        quickFixFactory: HLDiagnosticFixFactory<out DIAGNOSTIC>
+    ) {
+        quickFixes.getOrPut(quickFixFactory.diagnosticClass) { mutableListOf() }
             .add(HLQuickFixFactory.HLApplicatorBasedFactory(quickFixFactory))
     }
 

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

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.idea.intentions.fir.AddAccessorsIntention
 import org.jetbrains.kotlin.psi.KtProperty
 
 object AddAccessorsFactories {
-    val addAccessorsToUninitializedProperty = diagnosticFixFactory<KtFirDiagnostic.MustBeInitializedOrBeAbstract> { diagnostic ->
+    val addAccessorsToUninitializedProperty = diagnosticFixFactory(KtFirDiagnostic.MustBeInitializedOrBeAbstract::class) { diagnostic ->
         val property: KtProperty = diagnostic.psi
         val addGetter = property.getter == null
         val addSetter = property.isVar && property.setter == null

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

@@ -18,15 +18,15 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 object AddExclExclCallFixFactories {
-    val unsafeCallFactory = diagnosticFixFactory<KtFirDiagnostic.UnsafeCall> { diagnostic ->
+    val unsafeCallFactory = diagnosticFixFactory(KtFirDiagnostic.UnsafeCall::class) { diagnostic ->
         getFixForUnsafeCall(diagnostic.psi)
     }
 
-    val unsafeInfixCallFactory = diagnosticFixFactory<KtFirDiagnostic.UnsafeInfixCall> { diagnostic ->
+    val unsafeInfixCallFactory = diagnosticFixFactory(KtFirDiagnostic.UnsafeInfixCall::class) { diagnostic ->
         getFixForUnsafeCall(diagnostic.psi)
     }
 
-    val unsafeOperatorCallFactory = diagnosticFixFactory<KtFirDiagnostic.UnsafeOperatorCall> { diagnostic ->
+    val unsafeOperatorCallFactory = diagnosticFixFactory(KtFirDiagnostic.UnsafeOperatorCall::class) { diagnostic ->
         getFixForUnsafeCall(diagnostic.psi)
     }
 
@@ -92,7 +92,7 @@ object AddExclExclCallFixFactories {
         return listOfNotNull(target.asAddExclExclCallFix(hasImplicitReceiver = hasImplicitReceiver))
     }
 
-    val iteratorOnNullableFactory = diagnosticFixFactory<KtFirDiagnostic.IteratorOnNullable> { diagnostic ->
+    val iteratorOnNullableFactory = diagnosticFixFactory(KtFirDiagnostic.IteratorOnNullable::class) { diagnostic ->
         val expression = diagnostic.psi as? KtExpression ?: return@diagnosticFixFactory emptyList()
         val type = expression.getKtType()
         if (!type.canBeNull) return@diagnosticFixFactory emptyList()

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

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtProperty
 
 object AddLateInitFactory {
-    val addLateInitFactory = diagnosticFixFactory<KtFirDiagnostic.MustBeInitializedOrBeAbstract> { diagnostic ->
+    val addLateInitFactory = diagnosticFixFactory(KtFirDiagnostic.MustBeInitializedOrBeAbstract::class) { diagnostic ->
         val property: KtProperty = diagnostic.psi
         if (!property.isVar) return@diagnosticFixFactory emptyList()
 

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

@@ -118,14 +118,14 @@ object ChangeTypeQuickFixFactories {
         }
 
     val returnTypeMismatch =
-        diagnosticFixFactory<KtFirDiagnostic.ReturnTypeMismatch, KtCallableDeclaration, Input>(applicator) { diagnostic ->
+        diagnosticFixFactory(KtFirDiagnostic.ReturnTypeMismatch::class, applicator) { diagnostic ->
             val function = diagnostic.targetFunction.psi as? KtCallableDeclaration ?: return@diagnosticFixFactory emptyList()
             listOf(function withInput Input(TargetType.ENCLOSING_DECLARATION, createTypeInfo(diagnostic.actualType)))
         }
 
     @OptIn(ExperimentalStdlibApi::class)
     val componentFunctionReturnTypeMismatch =
-        diagnosticFixFactory<KtFirDiagnostic.ComponentFunctionReturnTypeMismatch, KtCallableDeclaration, Input>(applicator) { diagnostic ->
+        diagnosticFixFactory(KtFirDiagnostic.ComponentFunctionReturnTypeMismatch::class, applicator) { diagnostic ->
             val entryWithWrongType =
                 getDestructuringDeclarationEntryThatTypeMismatchComponentFunction(
                     diagnostic.componentFunctionName,
@@ -143,9 +143,9 @@ object ChangeTypeQuickFixFactories {
             }
         }
 
-    private inline fun <DIAGNOSTIC : KtDiagnosticWithPsi<KtNamedDeclaration>> changeReturnTypeOnOverride(
+    private inline fun <reified DIAGNOSTIC : KtDiagnosticWithPsi<KtNamedDeclaration>> changeReturnTypeOnOverride(
         crossinline getCallableSymbol: (DIAGNOSTIC) -> KtCallableSymbol?
-    ) = diagnosticFixFactory<DIAGNOSTIC, KtCallableDeclaration, Input>(applicator) { diagnostic ->
+    ) = diagnosticFixFactory(DIAGNOSTIC::class, applicator) { diagnostic ->
         val declaration = diagnostic.psi as? KtCallableDeclaration ?: return@diagnosticFixFactory emptyList()
         val callable = getCallableSymbol(diagnostic) ?: return@diagnosticFixFactory emptyList()
         listOfNotNull(

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

@@ -134,7 +134,7 @@ internal class ImportQuickFix(
     }
 
     internal companion object {
-        val FACTORY = diagnosticFixFactory<KtFirDiagnostic.UnresolvedReference> { diagnostic ->
+        val FACTORY = diagnosticFixFactory(KtFirDiagnostic.UnresolvedReference::class) { diagnostic ->
             val element = diagnostic.psi
 
             val indexHelper = IndexHelper(element.project, createSearchScope(element))

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

@@ -38,7 +38,7 @@ object InitializePropertyQuickFixFactory {
 
     @OptIn(ExperimentalStdlibApi::class)
     val initializePropertyFactory =
-        diagnosticFixFactory<KtFirDiagnostic.MustBeInitializedOrBeAbstract> { diagnostic ->
+        diagnosticFixFactory(KtFirDiagnostic.MustBeInitializedOrBeAbstract::class) { diagnostic ->
             val property: KtProperty = diagnostic.psi
             buildList {
                 add(

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

@@ -22,7 +22,7 @@ import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 object ReplaceCallFixFactories {
     val unsafeCallFactory =
-        diagnosticFixFactory<KtFirDiagnostic.UnsafeCall> { diagnostic ->
+        diagnosticFixFactory(KtFirDiagnostic.UnsafeCall::class) { diagnostic ->
             val psi = diagnostic.psi
             val target = if (psi is KtBinaryExpression && psi.operationToken in KtTokens.ALL_ASSIGNMENTS) {
                 // UNSAFE_CALL for assignments (e.g., `foo.bar = value`) is reported on the entire statement (KtBinaryExpression).
@@ -48,14 +48,14 @@ object ReplaceCallFixFactories {
         }
 
     val unsafeInfixCallFactory =
-        diagnosticFixFactory<KtFirDiagnostic.UnsafeInfixCall> { diagnostic ->
+        diagnosticFixFactory(KtFirDiagnostic.UnsafeInfixCall::class) { diagnostic ->
             val psi = diagnostic.psi
             val target = psi.parent as? KtBinaryExpression ?: return@diagnosticFixFactory emptyList()
             listOf(ReplaceInfixOrOperatorCallFix(target, shouldHaveNotNullType(target), diagnostic.operator))
         }
 
     val unsafeOperatorCallFactory =
-        diagnosticFixFactory<KtFirDiagnostic.UnsafeOperatorCall> { diagnostic ->
+        diagnosticFixFactory(KtFirDiagnostic.UnsafeOperatorCall::class) { diagnostic ->
             val psi = diagnostic.psi
             val operationToken = psi.safeAs<KtOperationReferenceExpression>()?.getReferencedNameElementType()
             if (operationToken == KtTokens.EQ || operationToken in OperatorConventions.COMPARISON_OPERATIONS) {
@@ -68,7 +68,7 @@ object ReplaceCallFixFactories {
         }
 
     val unsafeImplicitInvokeCallFactory =
-        diagnosticFixFactory<KtFirDiagnostic.UnsafeImplicitInvokeCall> { diagnostic ->
+        diagnosticFixFactory(KtFirDiagnostic.UnsafeImplicitInvokeCall::class) { diagnostic ->
             val target = diagnostic.psi as? KtNameReferenceExpression ?: return@diagnosticFixFactory emptyList()
 
             val callExpression = target.parent as? KtCallExpression ?: return@diagnosticFixFactory emptyList()

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

@@ -16,11 +16,11 @@ import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 object TypeMismatchFactories {
-    val argumentTypeMismatchFactory = diagnosticFixFactory<KtFirDiagnostic.ArgumentTypeMismatch> { diagnostic ->
+    val argumentTypeMismatchFactory = diagnosticFixFactory(KtFirDiagnostic.ArgumentTypeMismatch::class) { diagnostic ->
         getFixesForTypeMismatch(diagnostic.psi, diagnostic.expectedType, diagnostic.actualType)
     }
 
-    val returnTypeMismatchFactory = diagnosticFixFactory<KtFirDiagnostic.ReturnTypeMismatch> { diagnostic ->
+    val returnTypeMismatchFactory = diagnosticFixFactory(KtFirDiagnostic.ReturnTypeMismatch::class) { diagnostic ->
         getFixesForTypeMismatch(diagnostic.psi, diagnostic.expectedType, diagnostic.actualType)
     }
 

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

@@ -140,12 +140,12 @@ object WrapWithSafeLetCallFixFactories {
         }
     }
 
-    val forUnsafeCall = diagnosticFixFactory<KtFirDiagnostic.UnsafeCall> { diagnostic ->
+    val forUnsafeCall = diagnosticFixFactory(KtFirDiagnostic.UnsafeCall::class) { diagnostic ->
         val nullableExpression = diagnostic.receiverExpression
         createWrapWithSafeLetCallInputForNullableExpressionIfMoreThanImmediateParentIsWrapped(nullableExpression)
     }
 
-    val forUnsafeImplicitInvokeCall = diagnosticFixFactory<KtFirDiagnostic.UnsafeImplicitInvokeCall> { diagnostic ->
+    val forUnsafeImplicitInvokeCall = diagnosticFixFactory(KtFirDiagnostic.UnsafeImplicitInvokeCall::class) { diagnostic ->
         val callExpression = diagnostic.psi.parentOfType<KtCallExpression>(withSelf = true) ?: return@diagnosticFixFactory emptyList()
         val callingFunctionalVariableInLocalScope =
             isCallingFunctionalTypeVariableInLocalScope(callExpression) ?: return@diagnosticFixFactory emptyList()
@@ -164,15 +164,15 @@ object WrapWithSafeLetCallFixFactories {
         return localScope.scopes.getCallableSymbols { it.identifierOrNullIfSpecial == calleeName }.any { it == functionalVariableSymbol }
     }
 
-    val forUnsafeInfixCall = diagnosticFixFactory<KtFirDiagnostic.UnsafeInfixCall> { diagnostic ->
+    val forUnsafeInfixCall = diagnosticFixFactory(KtFirDiagnostic.UnsafeInfixCall::class) { diagnostic ->
         createWrapWithSafeLetCallInputForNullableExpressionIfMoreThanImmediateParentIsWrapped(diagnostic.receiverExpression)
     }
 
-    val forUnsafeOperatorCall = diagnosticFixFactory<KtFirDiagnostic.UnsafeOperatorCall> { diagnostic ->
+    val forUnsafeOperatorCall = diagnosticFixFactory(KtFirDiagnostic.UnsafeOperatorCall::class) { diagnostic ->
         createWrapWithSafeLetCallInputForNullableExpressionIfMoreThanImmediateParentIsWrapped(diagnostic.receiverExpression)
     }
 
-    val forArgumentTypeMismatch = diagnosticFixFactory<KtFirDiagnostic.ArgumentTypeMismatch> { diagnostic ->
+    val forArgumentTypeMismatch = diagnosticFixFactory(KtFirDiagnostic.ArgumentTypeMismatch::class) { diagnostic ->
         if (diagnostic.isMismatchDueToNullability) createWrapWithSafeLetCallInputForNullableExpression(diagnostic.psi.wrappingExpressionOrSelf)
         else emptyList()
     }