From ca6499d0f35a4e9dda7ea6f4a3a14dceb83603db Mon Sep 17 00:00:00 2001
From: Andrey Cherkasov <andrey.cherkasov@jetbrains.com>
Date: Thu, 25 Jul 2024 20:28:46 +0400
Subject: [PATCH] [kotlin] Convert AddReturnToLastExpressionInFunctionFix to
 ModCommandAction

^KTIJ-29555

GitOrigin-RevId: ab55789f1308595894df7463e5f1491b3dc05d3c
---
 .../AddReturnToLastExpressionInFunctionFix.kt | 25 +++++++++++--------
 .../AddReturnToLastExpressionInFunctionFix.kt |  2 +-
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/plugins/kotlin/frontend-independent/src/org/jetbrains/kotlin/idea/quickfix/AddReturnToLastExpressionInFunctionFix.kt b/plugins/kotlin/frontend-independent/src/org/jetbrains/kotlin/idea/quickfix/AddReturnToLastExpressionInFunctionFix.kt
index 39bcdb82d989..5a913561fad2 100644
--- a/plugins/kotlin/frontend-independent/src/org/jetbrains/kotlin/idea/quickfix/AddReturnToLastExpressionInFunctionFix.kt
+++ b/plugins/kotlin/frontend-independent/src/org/jetbrains/kotlin/idea/quickfix/AddReturnToLastExpressionInFunctionFix.kt
@@ -1,25 +1,28 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.quickfix
 
-import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.project.Project
+import com.intellij.modcommand.ActionContext
+import com.intellij.modcommand.ModPsiUpdater
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
-import org.jetbrains.kotlin.idea.codeinsight.api.classic.quickfixes.KotlinQuickFixAction
+import org.jetbrains.kotlin.idea.codeinsight.api.applicable.intentions.KotlinPsiUpdateModCommandAction
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtPsiFactory
 
-class AddReturnToLastExpressionInFunctionFix(element: KtDeclarationWithBody) : KotlinQuickFixAction<KtDeclarationWithBody>(element) {
+class AddReturnToLastExpressionInFunctionFix(
+    element: KtNamedFunction,
+) : KotlinPsiUpdateModCommandAction.ElementBased<KtDeclarationWithBody, Unit>(element, Unit) {
 
-    override fun getText() = KotlinBundle.message("fix.add.return.last.expression")
-    override fun getFamilyName() = text
+    override fun getFamilyName(): String = KotlinBundle.message("fix.add.return.last.expression")
 
-    override fun isAvailable(project: Project, editor: Editor?, file: KtFile): Boolean = element is KtNamedFunction
-
-    override fun invoke(project: Project, editor: Editor?, file: KtFile) {
+    override fun invoke(
+        actionContext: ActionContext,
+        element: KtDeclarationWithBody,
+        elementContext: Unit,
+        updater: ModPsiUpdater,
+    ) {
         val element = element as? KtNamedFunction ?: return
         val last = element.bodyBlockExpression?.statements?.lastOrNull() ?: return
-        last.replace(KtPsiFactory(project).createExpression("return ${last.text}"))
+        last.replace(KtPsiFactory(actionContext.project).createExpression("return ${last.text}"))
     }
 }
diff --git a/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/quickfix/AddReturnToLastExpressionInFunctionFix.kt b/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/quickfix/AddReturnToLastExpressionInFunctionFix.kt
index e87e3ec604cc..db9b30d4e006 100644
--- a/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/quickfix/AddReturnToLastExpressionInFunctionFix.kt
+++ b/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/quickfix/AddReturnToLastExpressionInFunctionFix.kt
@@ -27,6 +27,6 @@ internal object AddReturnToLastExpressionInFunctionFixFactory : KotlinSingleInte
         val expectedType = namedFunction.resolveToDescriptorIfAny()?.returnType?.takeIf { !it.isError } ?: return null
         if (!lastType.isSubtypeOf(expectedType)) return null
 
-        return AddReturnToLastExpressionInFunctionFix(casted.psiElement)
+        return AddReturnToLastExpressionInFunctionFix(namedFunction).asIntention()
     }
 }
\ No newline at end of file