[kotlin] Reuse getUnreachableTextRanges to point unreachable elements

#KTIJ-32807

GitOrigin-RevId: a9ebbc27f23e4fd4890cc9bc71220f54a3bd7b7a

plugins/kotlin/code-insight/inspections-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/inspections/diagnosticBased/KotlinUnreachableCodeInspection.kt

@@ -13,6 +13,7 @@ import com.intellij.psi.createSmartPointer
 import com.intellij.psi.util.startOffset
 import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.analysis.api.fir.diagnostics.KaFirDiagnostic
+import org.jetbrains.kotlin.cfg.UnreachableCode
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.codeinsight.api.applicable.inspections.KotlinKtDiagnosticBasedInspectionBase
 import org.jetbrains.kotlin.idea.codeinsight.api.applicable.inspections.KotlinModCommandQuickFix
@@ -23,7 +24,7 @@ import kotlin.reflect.KClass
 
 class KotlinUnreachableCodeInspection : KotlinKtDiagnosticBasedInspectionBase<KtElement, KaFirDiagnostic.UnreachableCode, KotlinUnreachableCodeInspection.Context>() {
     data class Context(
-        val unreachable: Collection<SmartPsiElementPointer<*>>
+        val unreachableElementPointers: Collection<SmartPsiElementPointer<*>>
     )
 
     override val diagnosticType: KClass<KaFirDiagnostic.UnreachableCode>
@@ -33,45 +34,27 @@ class KotlinUnreachableCodeInspection : KotlinKtDiagnosticBasedInspectionBase<Kt
         element: KtElement,
         diagnostic: KaFirDiagnostic.UnreachableCode
     ): Context? {
-        val reachable = diagnostic.reachable
-        val unreachable = diagnostic.unreachable
-        val firstUnreachable = unreachable.firstOrNull()
-        val unreachableElements = if (firstUnreachable != null) {
-            val reachableTextRange = reachable.firstOrNull()
-                ?.let { reachable.fold(it.textRange) { range: TextRange, element: PsiElement -> range.union(element.textRange) } }
-                ?: return null
+        val reachable = diagnostic.reachable.mapNotNullTo(HashSet()) { it as? KtElement }
+        val unreachable = diagnostic.unreachable.mapNotNullTo(HashSet()) { it as? KtElement }
 
-            val unreachableRange = unreachable.fold(firstUnreachable.textRange) { it, element -> it.union(element.textRange) }
+        val unreachableTextRanges =
+            UnreachableCode
+                .getUnreachableTextRanges(element, reachable, unreachable)
+                .ifEmpty { return null }
 
-            val startOffset = unreachableRange.startOffset
-            val reachableStartOffset = reachableTextRange.startOffset
-
-            val unreachableTextRange = if (startOffset <= reachableStartOffset) {
-                TextRange(startOffset, unreachableRange.endOffset.coerceAtMost(reachableStartOffset))
-            } else {
-                TextRange(reachableStartOffset, unreachableRange.endOffset.coerceAtMost(reachableTextRange.endOffset))
+        val unreachableElementPointers = buildList {
+            val startOffset = element.startOffset
+            for (textRange in unreachableTextRanges) {
+                var offset = textRange.startOffset - startOffset
+                while (offset < textRange.endOffset - startOffset) {
+                    val findElementAt = element.findElementAt(offset) ?: break
+                    offset += findElementAt.textLength
+                    this += findElementAt.createSmartPointer()
+                }
             }
-
-            calculateUnreachableElements(element, unreachableTextRange, unreachableRange)
-        } else {
-            emptyList()
         }
-        return Context(unreachableElements)
-    }
 
-    private fun calculateUnreachableElements(
-        element: KtElement,
-        unreachableTextRange: TextRange,
-        singleUnreachableRange: TextRange
-    ): List<SmartPsiElementPointer<*>> = buildList {
-        var offset = unreachableTextRange.startOffset - element.startOffset
-        while (offset < singleUnreachableRange.endOffset - element.startOffset) {
-            val e = element.findElementAt(offset) ?: break
-            if (e !is PsiWhiteSpace && unreachableTextRange.contains(e.textRange)) {
-                this += e.createSmartPointer()
-            }
-            offset += e.textRange.length
-        }
+        return Context(unreachableElementPointers)
     }
 
     override fun getProblemDescription(
@@ -91,14 +74,10 @@ class KotlinUnreachableCodeInspection : KotlinKtDiagnosticBasedInspectionBase<Kt
             element: KtElement,
             updater: ModPsiUpdater,
         ) {
-            val unreachable = context.unreachable
-            if (unreachable.isEmpty()) {
-                element.delete()
-            } else {
-                unreachable.mapNotNull {
-                    it.element?.let(updater::getWritable)
+            context.unreachableElementPointers
+                .mapNotNull {
+                    it.element?.takeUnless { e -> e is PsiWhiteSpace }?.let(updater::getWritable)
                 }.forEach(PsiElement::delete)
-            }
         }
     }
 
@@ -118,22 +97,12 @@ class KotlinUnreachableCodeInspection : KotlinKtDiagnosticBasedInspectionBase<Kt
         context: Context,
         isOnTheFly: Boolean
     ) {
-        val unreachable = context.unreachable
-        val first = unreachable.firstOrNull()?.element
-        val last = unreachable.lastOrNull()?.element
-        if (first == null || last == null) {
+        context.unreachableElementPointers.forEach { pointer ->
+            val e = pointer.element ?: return@forEach
             val problemDescriptor = holder.manager.createProblemDescriptor(
                 element,
                 context,
-                TextRange(0, element.textRange.length),
-                isOnTheFly,
-            )
-            holder.registerProblem(problemDescriptor)
-        } else {
-            val problemDescriptor = holder.manager.createProblemDescriptor(
-                element,
-                context,
-                first.textRange.union(last.textRange).shiftLeft(element.startOffset),
+                e.textRange.shiftLeft(element.startOffset),
                 isOnTheFly,
             )
             holder.registerProblem(problemDescriptor)

plugins/kotlin/code-insight/inspections-k2/tests/test/org/jetbrains/kotlin/idea/k2/inspections/tests/K2LocalInspectionTestGenerated.java

@@ -14319,6 +14319,11 @@ public abstract class K2LocalInspectionTestGenerated extends AbstractK2LocalInsp
             runTest("../../../idea/tests/testData/inspectionsLocal/kotlinUnreachableCode/returnReturn.kt");
         }
 
+        @TestMetadata("returnReturnAndStatement.kt")
+        public void testReturnReturnAndStatement() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/kotlinUnreachableCode/returnReturnAndStatement.kt");
+        }
+
         @TestMetadata("returnTryReturn.kt")
         public void testReturnTryReturn() throws Exception {
             runTest("../../../idea/tests/testData/inspectionsLocal/kotlinUnreachableCode/returnTryReturn.kt");

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

@@ -6971,6 +6971,11 @@ public abstract class LocalInspectionTestGenerated extends AbstractLocalInspecti
             runTest("testData/inspectionsLocal/kotlinUnreachableCode/returnReturn.kt");
         }
 
+        @TestMetadata("returnReturnAndStatement.kt")
+        public void testReturnReturnAndStatement() throws Exception {
+            runTest("testData/inspectionsLocal/kotlinUnreachableCode/returnReturnAndStatement.kt");
+        }
+
         @TestMetadata("returnTryReturn.kt")
         public void testReturnTryReturn() throws Exception {
             runTest("testData/inspectionsLocal/kotlinUnreachableCode/returnTryReturn.kt");

plugins/kotlin/idea/tests/testData/inspectionsLocal/kotlinUnreachableCode/returnReturnAndStatement.kt

@@ -0,0 +1,6 @@
+// IGNORE_K1
+// FIX: Remove unreachable code
+fun f(): Int {
+    val x = 1
+    r<caret>eturn  return 42
+}

plugins/kotlin/idea/tests/testData/inspectionsLocal/kotlinUnreachableCode/returnReturnAndStatement.kt.after

@@ -0,0 +1,6 @@
+// IGNORE_K1
+// FIX: Remove unreachable code
+fun f(): Int {
+    val x = 1
+    return 42
+}