FIR IDE: introduce ChangeReturnTypeOnOverrideQuickFix

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt

@@ -1106,6 +1106,7 @@ fun main(args: Array<String>) {
             testClass<AbstractHighLevelQuickFixTest> {
                 val pattern = "^([\\w\\-_]+)\\.kt$"
                 model("quickfix/modifiers", pattern = pattern, filenameStartsLowerCase = true, recursive = false)
+                model("quickfix/override/typeMismatchOnOverride", pattern = pattern, filenameStartsLowerCase = true, recursive = false)
             }
         }
 

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

@@ -18,8 +18,16 @@ annotation class ForKtQuickFixesListBuilder()
 class KtQuickFixesListBuilder private constructor() {
     val quickFixes = mutableMapOf<KClass<out KtDiagnosticWithPsi<*>>, MutableList<QuickFixFactory>>()
 
-    inline fun <reified PSI : PsiElement, reified D : KtDiagnosticWithPsi<PSI>> register(quickFixFactory: QuickFixFactory) {
-        quickFixes.getOrPut(D::class) { mutableListOf() }.add(quickFixFactory)
+    inline fun <reified PSI : PsiElement, reified DIAGNOSTIC : KtDiagnosticWithPsi<PSI>> register(
+        quickFixFactory: QuickFixesPsiBasedFactory<PSI>
+    ) {
+        quickFixes.getOrPut(DIAGNOSTIC::class) { mutableListOf() }.add(quickFixFactory)
+    }
+
+    inline fun <reified PSI : PsiElement, reified DIAGNOSTIC : KtDiagnosticWithPsi<PSI>> register(
+        quickFixFactory: QuickFixesHLApiBasedFactory<PSI, DIAGNOSTIC>
+    ) {
+        quickFixes.getOrPut(DIAGNOSTIC::class) { mutableListOf() }.add(quickFixFactory)
     }
 
     @OptIn(ForKtQuickFixesListBuilder::class)

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtCallableDeclaration
 import org.jetbrains.kotlin.psi.KtModifierListOwner
+import org.jetbrains.kotlin.psi.KtNamedDeclaration
 
 class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
     private val modifiers = KtQuickFixesListBuilder.register {
@@ -28,7 +29,15 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
         )
     }
 
+    private val overrides = KtQuickFixesListBuilder.register {
+        register(ChangeTypeQuickFix.changeFunctionReturnTypeOnOverride)
+        register(ChangeTypeQuickFix.changePropertyReturnTypeOnOverride)
+        register(ChangeTypeQuickFix.changeVariableReturnTypeOnOverride)
+
+    }
+
     override val list: KtQuickFixesList = KtQuickFixesList.createCombined(
         modifiers,
+        overrides,
     )
 }
\ No newline at end of file

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

+/*
+ * 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
+
+import com.intellij.codeInsight.intention.IntentionAction
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.diagnostics.KtDiagnosticWithPsi
+
+abstract class QuickFixesHLApiBasedFactory<PSI : PsiElement, DIAGNOSTIC : KtDiagnosticWithPsi<PSI>> : QuickFixFactory {
+    final override fun asKotlinIntentionActionsFactory(): KotlinIntentionActionsFactory {
+        error("Should not be called. This function is not considered to bue used in FE10 plugin, from FIR plugin consider using createQuickFix")
+    }
+
+    abstract fun KtAnalysisSession.createQuickFix(diagnostic: DIAGNOSTIC): List<IntentionAction>
+}
+
+
+inline fun <PSI : PsiElement, DIAGNOSTIC : KtDiagnosticWithPsi<PSI>> quickFixesHLApiBasedFactory(
+    crossinline createQuickFix: KtAnalysisSession.(DIAGNOSTIC) -> List<IntentionAction>
+) = object : QuickFixesHLApiBasedFactory<PSI, DIAGNOSTIC>() {
+    override fun KtAnalysisSession.createQuickFix(diagnostic: DIAGNOSTIC): List<IntentionAction> =
+        createQuickFix(diagnostic)
+}
\ No newline at end of file

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

+/*
+ * 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 com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import com.intellij.psi.util.parentOfType
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.components.KtTypeRendererOptions
+import org.jetbrains.kotlin.idea.frontend.api.diagnostics.KtDiagnosticWithPsi
+import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
+import org.jetbrains.kotlin.idea.frontend.api.symbols.*
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
+import org.jetbrains.kotlin.idea.frontend.api.types.isUnit
+import org.jetbrains.kotlin.idea.quickfix.ChangeCallableReturnTypeFix
+import org.jetbrains.kotlin.idea.quickfix.KotlinQuickFixAction
+import org.jetbrains.kotlin.idea.quickfix.quickFixesHLApiBasedFactory
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
+
+
+class ChangeTypeQuickFix internal constructor(
+    declaration: KtCallableDeclaration,
+    private val typeInfo: TypeInfo,
+    private val updateBaseFunction: Boolean
+) : KotlinQuickFixAction<KtCallableDeclaration>(declaration) {
+    override fun getText(): String {
+        val element = element ?: return ""
+        val functionPresentation = when {
+            updateBaseFunction -> {
+                val containerName = element.parentOfType<KtNamedDeclaration>()?.nameAsName?.takeUnless { it.isSpecial }
+                ChangeCallableReturnTypeFix.StringPresentation.baseFunctionOrConstructorParameterPresentation(element, containerName)
+            }
+            else -> null
+        }
+        return ChangeCallableReturnTypeFix.StringPresentation.getTextForQuickFix(
+            element,
+            functionPresentation,
+            typeInfo.isUnit,
+            typeInfo.short
+        )
+    }
+
+    override fun getFamilyName(): String = ChangeCallableReturnTypeFix.StringPresentation.familyName()
+
+    override fun invoke(project: Project, editor: Editor?, file: KtFile) {
+        val element = element ?: return
+        if (!element.isProcedure()) {
+            val newTypeRef = KtPsiFactory(project).createType(typeInfo.short)
+            element.typeReference = newTypeRef
+        } else {
+            element.typeReference = null
+        }
+    }
+
+    private fun KtCallableDeclaration.isProcedure() =
+        typeInfo.isUnit && this is KtFunction && hasBlockBody()
+
+    internal data class TypeInfo(
+        val qualified: String,
+        val short: String,
+        val isUnit: Boolean,
+    )
+
+    companion object {
+        val changeFunctionReturnTypeOnOverride =
+            ChangeTypeQuickFixFactory.changeReturnTypeOnOverride<KtNamedDeclaration, KtFirDiagnostic.ReturnTypeMismatchOnOverride> {
+                it.function as? KtFunctionSymbol
+            }
+
+        val changePropertyReturnTypeOnOverride =
+            ChangeTypeQuickFixFactory.changeReturnTypeOnOverride<KtNamedDeclaration, KtFirDiagnostic.PropertyTypeMismatchOnOverride> {
+                it.property as? KtPropertySymbol
+            }
+
+        val changeVariableReturnTypeOnOverride =
+            ChangeTypeQuickFixFactory.changeReturnTypeOnOverride<KtNamedDeclaration, KtFirDiagnostic.VarTypeMismatchOnOverride> {
+                it.variable as? KtPropertySymbol
+            }
+    }
+}
+
+private object ChangeTypeQuickFixFactory {
+    inline fun <PSI : KtNamedDeclaration, DIAGNOSTIC : KtDiagnosticWithPsi<PSI>> changeReturnTypeOnOverride(
+        crossinline getCallableSymbol: (DIAGNOSTIC) -> KtCallableSymbol?
+    ) = quickFixesHLApiBasedFactory<PSI, DIAGNOSTIC> { diagnostic ->
+        val declaration = diagnostic.psi as? KtCallableDeclaration ?: return@quickFixesHLApiBasedFactory emptyList()
+        val callable = getCallableSymbol(diagnostic) ?: return@quickFixesHLApiBasedFactory emptyList()
+        listOfNotNull(
+            createChangeCurrentDeclarationQuickFix(callable, declaration),
+            createChangeOverriddenFunctionQuickFix(callable),
+        )
+    }
+
+    fun <PSI : KtCallableDeclaration> KtAnalysisSession.createChangeCurrentDeclarationQuickFix(
+        callable: KtCallableSymbol,
+        declaration: PSI
+    ): ChangeTypeQuickFix? {
+        val lowerSuperType = findLowerBoundOfOverriddenCallablesReturnTypes(callable) ?: return null
+        val changeToTypeInfo = createTypeInfo(lowerSuperType)
+        return ChangeTypeQuickFix(declaration, changeToTypeInfo, updateBaseFunction = false)
+    }
+
+    fun KtAnalysisSession.createChangeOverriddenFunctionQuickFix(
+        callable: KtCallableSymbol
+    ): ChangeTypeQuickFix? {
+        val type = callable.annotatedType.type
+        val singleNonMatchingOverriddenFunction = findSingleNonMatchingOverriddenFunction(callable, type) ?: return null
+        val singleMatchingOverriddenFunctionPsi = singleNonMatchingOverriddenFunction.psiSafe<KtCallableDeclaration>() ?: return null
+        val changeToTypeInfo = createTypeInfo(type)
+        return ChangeTypeQuickFix(singleMatchingOverriddenFunctionPsi, changeToTypeInfo, updateBaseFunction = true)
+    }
+
+    private fun KtAnalysisSession.findSingleNonMatchingOverriddenFunction(
+        callable: KtCallableSymbol,
+        type: KtType
+    ): KtCallableSymbol? {
+        val overriddenSymbols = callable.getOverriddenSymbols()
+        return overriddenSymbols
+            .singleOrNull { overridden ->
+                overridden.origin != KtSymbolOrigin.INTERSECTION_OVERRIDE && !type.isSubTypeOf(overridden.annotatedType.type)
+            }
+    }
+
+    fun KtAnalysisSession.createTypeInfo(ktType: KtType) = ChangeTypeQuickFix.TypeInfo(
+        qualified = ktType.render(),
+        short = ktType.render(KtTypeRendererOptions.SHORT_NAMES),
+        isUnit = ktType.isUnit
+    )
+
+    private fun KtAnalysisSession.findLowerBoundOfOverriddenCallablesReturnTypes(symbol: KtCallableSymbol): KtType? {
+        var lowestType: KtType? = null
+        for (overridden in symbol.getOverriddenSymbols()) {
+            val overriddenType = overridden.annotatedType.type
+            when {
+                lowestType == null || overriddenType isSubTypeOf lowestType -> {
+                    lowestType = overriddenType
+                }
+                lowestType isNotSubTypeOf overriddenType -> {
+                    return null
+                }
+            }
+        }
+        return lowestType
+    }
+}