Add quick-fix "add .getOrThrow()" for SuccessOrFailure inspection

#KT-25621 Fixed

idea/src/org/jetbrains/kotlin/idea/inspections/coroutines/ResultIsSuccessOrFailureInspection.kt

@@ -5,19 +5,29 @@
 
 package org.jetbrains.kotlin.idea.inspections.coroutines
 
+import com.intellij.codeInspection.LocalQuickFix
+import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
+import org.jetbrains.kotlin.idea.core.setType
 import org.jetbrains.kotlin.idea.inspections.AbstractKotlinInspection
+import org.jetbrains.kotlin.idea.intentions.branchedTransformations.BranchedFoldingUtils
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.callableBuilder.getReturnTypeReference
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.bindingContextUtil.getTargetFunctionDescriptor
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 
 class ResultIsSuccessOrFailureInspection : AbstractKotlinInspection() {
@@ -57,7 +67,8 @@ class ResultIsSuccessOrFailureInspection : AbstractKotlinInspection() {
             holder.registerProblem(
                 toReport,
                 "Function returning $SHORT_NAME with a name that does not end with $CATCHING",
-                ProblemHighlightType.GENERIC_ERROR_OR_WARNING
+                ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
+                AddGetOrThrowFix(withBody = function.hasBody())
             )
         }
     }
@@ -74,11 +85,54 @@ class ResultIsSuccessOrFailureInspection : AbstractKotlinInspection() {
         }
     }
 
+    private class AddGetOrThrowFix(val withBody: Boolean) : LocalQuickFix {
+        override fun getName(): String =
+            if (withBody) "Add '.$GET_OR_THROW()' to function result (breaks use-sites!)"
+            else "Unfold '$SHORT_NAME' return type (breaks use-sites!)"
+
+        override fun getFamilyName(): String = name
+
+        private fun KtExpression.addGetOrThrow(factory: KtPsiFactory) {
+            val getOrThrowExpression = factory.createExpressionByPattern("$0.$GET_OR_THROW()", this)
+            replace(getOrThrowExpression)
+        }
+
+        override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+            val function = descriptor.psiElement.getNonStrictParentOfType<KtFunction>() ?: return
+            val returnTypeReference = function.getReturnTypeReference()
+            val context = function.analyze()
+            val functionDescriptor = context[BindingContext.FUNCTION, function] ?: return
+            if (returnTypeReference != null) {
+                val returnType = functionDescriptor.returnType ?: return
+                val returnTypeArgument = returnType.arguments.firstOrNull()?.type ?: return
+                function.setType(returnTypeArgument)
+            }
+            if (!withBody) return
+            val factory = KtPsiFactory(project)
+            val bodyExpression = function.bodyExpression
+            bodyExpression?.forEachDescendantOfType<KtReturnExpression> {
+                if (it.getTargetFunctionDescriptor(context) == functionDescriptor) {
+                    it.returnedExpression?.addGetOrThrow(factory)
+                }
+            }
+            if (function is KtFunctionLiteral) {
+                val lastStatement = function.bodyExpression?.statements?.lastOrNull()
+                if (lastStatement != null && lastStatement !is KtReturnExpression) {
+                    lastStatement.addGetOrThrow(factory)
+                }
+            } else if (!function.hasBlockBody()) {
+                bodyExpression?.addGetOrThrow(factory)
+            }
+        }
+    }
+
     companion object {
         private const val SHORT_NAME = "SuccessOrFailure"
 
         private const val FULL_NAME = "kotlin.$SHORT_NAME"
 
         private const val CATCHING = "Catching"
+
+        private const val GET_OR_THROW = "getOrThrow"
     }
 }

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/.inspection

@@ -0,0 +1 @@
+org.jetbrains.kotlin.idea.inspections.coroutines.ResultIsSuccessOrFailureInspection

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/anonymous.kt

@@ -0,0 +1,9 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun test() {
+    val x = <caret>fun() = SuccessOrFailure("123")
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/anonymous.kt.after

@@ -0,0 +1,9 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun test() {
+    val x = fun() = SuccessOrFailure("123").getOrThrow()
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/complexBlock.kt

@@ -0,0 +1,40 @@
+// WITH_RUNTIME
+
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun <caret>incorrectBlock(arg: Boolean, arg2: Boolean?): SuccessOrFailure<Int> {
+    if (arg) {
+        class Local {
+            fun foo(): SuccessOrFailure<String> {
+                return SuccessOrFailure("NO")
+            }
+        }
+        return SuccessOrFailure(1)
+    } else {
+        when (arg2) {
+            true -> {
+                val x = fun(): SuccessOrFailure<Boolean> {
+                    return SuccessOrFailure(false)
+                }
+                if (x().getOrThrow()) {
+                    return SuccessOrFailure(2)
+                } else {
+                    return SuccessOrFailure(0)
+                }
+            }
+            else -> {
+                if (arg2 == false) {
+                    listOf(1, 2, 3).forEach {
+                        if (it == 2) return@forEach
+                    }
+                    return SuccessOrFailure(3)
+                }
+                return SuccessOrFailure(4)
+            }
+        }
+    }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/complexBlock.kt.after

@@ -0,0 +1,40 @@
+// WITH_RUNTIME
+
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun incorrectBlock(arg: Boolean, arg2: Boolean?): Int {
+    if (arg) {
+        class Local {
+            fun foo(): SuccessOrFailure<String> {
+                return SuccessOrFailure("NO")
+            }
+        }
+        return SuccessOrFailure(1).getOrThrow()
+    } else {
+        when (arg2) {
+            true -> {
+                val x = fun(): SuccessOrFailure<Boolean> {
+                    return SuccessOrFailure(false)
+                }
+                if (x().getOrThrow()) {
+                    return SuccessOrFailure(2).getOrThrow()
+                } else {
+                    return SuccessOrFailure(0).getOrThrow()
+                }
+            }
+            else -> {
+                if (arg2 == false) {
+                    listOf(1, 2, 3).forEach {
+                        if (it == 2) return@forEach
+                    }
+                    return SuccessOrFailure(3).getOrThrow()
+                }
+                return SuccessOrFailure(4).getOrThrow()
+            }
+        }
+    }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaSimple.kt

@@ -0,0 +1,9 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun test() {
+    val x = <caret>{ SuccessOrFailure(true) }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaSimple.kt.after

@@ -0,0 +1,9 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun test() {
+    val x = { SuccessOrFailure(true).getOrThrow() }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaWithPartialReturn.kt

@@ -0,0 +1,15 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun test(arg: Boolean) {
+    val x = foo@<caret>{
+        if (!arg) {
+            return@foo SuccessOrFailure(true)
+        } else {
+            SuccessOrFailure(false)
+        }
+    }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaWithPartialReturn.kt.after

@@ -0,0 +1,15 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun test(arg: Boolean) {
+    val x = foo@{
+        (if (!arg) {
+            return@foo SuccessOrFailure(true).getOrThrow()
+        } else {
+            SuccessOrFailure(false)
+        }).getOrThrow()
+    }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaWithReturn.kt

@@ -0,0 +1,9 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun test() {
+    val x = foo@<caret>{ return@foo SuccessOrFailure(true) }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaWithReturn.kt.after

@@ -0,0 +1,9 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun test() {
+    val x = foo@<caret>{ return@foo SuccessOrFailure(true).getOrThrow() }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/needParentheses.kt

@@ -0,0 +1,9 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+
+    operator fun plus(other: SuccessOrFailure<T>) = other
+}
+
+fun <caret>incorrect() = SuccessOrFailure("123") + SuccessOrFailure("456")

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/needParentheses.kt.after

@@ -0,0 +1,9 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+
+    operator fun plus(other: SuccessOrFailure<T>) = other
+}
+
+fun incorrect() = (SuccessOrFailure("123") + SuccessOrFailure("456")).getOrThrow()

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/simple.kt

@@ -0,0 +1,7 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun <caret>incorrect() = SuccessOrFailure("123")

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/simple.kt.after

@@ -0,0 +1,7 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun incorrect() = SuccessOrFailure("123").getOrThrow()

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/simpleBlock.kt

@@ -0,0 +1,21 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun <caret>incorrectBlock(arg: Boolean, arg2: Boolean?): SuccessOrFailure<Int> {
+    if (arg) {
+        return SuccessOrFailure(1)
+    } else {
+        when (arg2) {
+            true -> return SuccessOrFailure(2)
+            else -> {
+                if (arg2 == false) {
+                    return SuccessOrFailure(3)
+                }
+                return SuccessOrFailure(4)
+            }
+        }
+    }
+}

idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/simpleBlock.kt.after

@@ -0,0 +1,21 @@
+package kotlin
+
+class SuccessOrFailure<T>(val value: T?) {
+    fun getOrThrow(): T = value ?: throw AssertionError("")
+}
+
+fun incorrectBlock(arg: Boolean, arg2: Boolean?): Int {
+    if (arg) {
+        return SuccessOrFailure(1).getOrThrow()
+    } else {
+        when (arg2) {
+            true -> return SuccessOrFailure(2).getOrThrow()
+            else -> {
+                if (arg2 == false) {
+                    return SuccessOrFailure(3).getOrThrow()
+                }
+                return SuccessOrFailure(4).getOrThrow()
+            }
+        }
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -1864,6 +1864,59 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 runTest("idea/testData/inspectionsLocal/coroutines/redundantAsync/withStartNoContext.kt");
             }
         }
+
+        @TestMetadata("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class ResultIsSuccessOrFailure extends AbstractLocalInspectionTest {
+            private void runTest(String testDataFilePath) throws Exception {
+                KotlinTestUtils.runTest(this::doTest, TargetBackend.ANY, testDataFilePath);
+            }
+
+            public void testAllFilesPresentInResultIsSuccessOrFailure() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure"), Pattern.compile("^([\\w\\-_]+)\\.(kt|kts)$"), TargetBackend.ANY, true);
+            }
+
+            @TestMetadata("anonymous.kt")
+            public void testAnonymous() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/anonymous.kt");
+            }
+
+            @TestMetadata("complexBlock.kt")
+            public void testComplexBlock() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/complexBlock.kt");
+            }
+
+            @TestMetadata("lambdaSimple.kt")
+            public void testLambdaSimple() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaSimple.kt");
+            }
+
+            @TestMetadata("lambdaWithPartialReturn.kt")
+            public void testLambdaWithPartialReturn() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaWithPartialReturn.kt");
+            }
+
+            @TestMetadata("lambdaWithReturn.kt")
+            public void testLambdaWithReturn() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/lambdaWithReturn.kt");
+            }
+
+            @TestMetadata("needParentheses.kt")
+            public void testNeedParentheses() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/needParentheses.kt");
+            }
+
+            @TestMetadata("simple.kt")
+            public void testSimple() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/simple.kt");
+            }
+
+            @TestMetadata("simpleBlock.kt")
+            public void testSimpleBlock() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/resultIsSuccessOrFailure/simpleBlock.kt");
+            }
+        }
     }
 
     @TestMetadata("idea/testData/inspectionsLocal/deprecatedCallableAddReplaceWith")