Replace with loop over elements: don't suggest when index is used for other array

#KT-40592 Fixed

GitOrigin-RevId: bb94e46006a93d23dded396a7406146501c5f628

7 files changed, 73 insertions, 7 deletions

diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/ReplaceManualRangeWithIndicesCallsInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/ReplaceManualRangeWithIndicesCallsInspection.kt
index d7a6ae8a1836..d0647c7ec026 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/ReplaceManualRangeWithIndicesCallsInspection.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/inspections/ReplaceManualRangeWithIndicesCallsInspection.kt
@@ -14,6 +14,7 @@ import com.intellij.psi.search.searches.ReferencesSearch
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.intentions.getArguments
 import org.jetbrains.kotlin.idea.intentions.isSizeOrLength
+import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
@@ -40,20 +41,25 @@ class ReplaceManualRangeWithIndicesCallsInspection : AbstractKotlinInspection()
     }
 
     private fun visitRange(holder: ProblemsHolder, expression: KtExpression, left: KtExpression, right: KtExpression, method: String) {
-        if ((method == "until" && left.toIntConstant() == 0 && right.receiverIfIsSizeOrLengthCall() != null) ||
-            (method == "rangeTo" && left.toIntConstant() == 0 && right.receiverIfIsSizeOrLengthMinusOneCall() != null)
-        ) {
-            visitIndicesRange(holder, expression)
-        }
+        if (left.toIntConstant() != 0) return
+        val collection = when(method) {
+            "until" -> right.receiverIfIsSizeOrLengthCall()
+            "rangeTo" -> right.receiverIfIsSizeOrLengthMinusOneCall()
+            else -> null
+        } as? KtSimpleNameExpression ?: return
+        visitIndicesRange(holder, expression, collection)
     }
 
-    private fun visitIndicesRange(holder: ProblemsHolder, range: KtExpression) {
+    private fun visitIndicesRange(holder: ProblemsHolder, range: KtExpression, collection: KtSimpleNameExpression) {
         val parent = range.parent.parent
         if (parent is KtForExpression) {
             val paramElement = parent.loopParameter?.originalElement ?: return
             val usageElement = ReferencesSearch.search(paramElement).singleOrNull()?.element
             val arrayAccess = usageElement?.parent?.parent as? KtArrayAccessExpression
-            if (arrayAccess != null && arrayAccess.indexExpressions.singleOrNull() == usageElement) {
+            if (arrayAccess != null &&
+                arrayAccess.indexExpressions.singleOrNull() == usageElement &&
+                (arrayAccess.arrayExpression as? KtSimpleNameExpression)?.mainReference?.resolve() == collection.mainReference.resolve()
+            ) {
                 val arrayAccessParent = arrayAccess.parent
                 if (arrayAccessParent !is KtBinaryExpression ||
                     arrayAccessParent.left != arrayAccess ||
diff --git a/idea/test/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java b/idea/test/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
index 53e00e011992..17c79d95ca07 100644
--- a/idea/test/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
+++ b/idea/test/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
@@ -10865,6 +10865,16 @@ public abstract class LocalInspectionTestGenerated extends AbstractLocalInspecti
             runTest("testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/indexInLvalue.kt");
         }
 
+        @TestMetadata("multidimensionalArray.kt")
+        public void testMultidimensionalArray() throws Exception {
+            runTest("testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray.kt");
+        }
+
+        @TestMetadata("multidimensionalArray2.kt")
+        public void testMultidimensionalArray2() throws Exception {
+            runTest("testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray2.kt");
+        }
+
         @TestMetadata("notUsedAsIndex.kt")
         public void testNotUsedAsIndex() throws Exception {
             runTest("testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/notUsedAsIndex.kt");
@@ -10884,6 +10894,11 @@ public abstract class LocalInspectionTestGenerated extends AbstractLocalInspecti
         public void testSimpleFor() throws Exception {
             runTest("testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/simpleFor.kt");
         }
+
+        @TestMetadata("usedForOtherArray.kt")
+        public void testUsedForOtherArray() throws Exception {
+            runTest("testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/usedForOtherArray.kt");
+        }
     }
 
     @RunWith(JUnit3RunnerWithInners.class)
diff --git a/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray.kt b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray.kt
new file mode 100644
index 000000000000..8a89fa5791a5
--- /dev/null
+++ b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray.kt
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// PROBLEM: none
+fun foo(matrix: Array<Array<Int>>) {
+    for (i in <caret>0 until matrix[0].size) {
+        for (j in 0 until matrix.size) {
+            print(matrix[j][i])
+        }
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray2.kt b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray2.kt
new file mode 100644
index 000000000000..6bf5d6cd1244
--- /dev/null
+++ b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray2.kt
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// FIX: Replace with loop over elements
+fun foo(matrix: Array<Array<Int>>) {
+    for (i in 0 until matrix[0].size) {
+        for (j in <caret>0 until matrix.size) {
+            print(matrix[j][i])
+        }
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray2.kt.after b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray2.kt.after
new file mode 100644
index 000000000000..7ae0cf69c6f8
--- /dev/null
+++ b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/multidimensionalArray2.kt.after
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// FIX: Replace with loop over elements
+fun foo(matrix: Array<Array<Int>>) {
+    for (i in 0 until matrix[0].size) {
+        for (element in matrix) {
+            print(element[i])
+        }
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/usedForOtherArray.kt b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/usedForOtherArray.kt
new file mode 100644
index 000000000000..3c4e5074795b
--- /dev/null
+++ b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/usedForOtherArray.kt
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// FIX: Replace with indices
+fun foo() {
+    val a = listOf(1, 2, 3)
+    val b = listOf(11, 22, 33, 44, 55)
+    for (i in <caret>0 until a.size) {
+        println(b[i])
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/usedForOtherArray.kt.after b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/usedForOtherArray.kt.after
new file mode 100644
index 000000000000..fd0fe8a70019
--- /dev/null
+++ b/idea/testData/inspectionsLocal/replaceManualRangeWithIndicesCalls/usedForOtherArray.kt.after
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// FIX: Replace with indices
+fun foo() {
+    val a = listOf(1, 2, 3)
+    val b = listOf(11, 22, 33, 44, 55)
+    for (i in a.indices) {
+        println(b[i])
+    }
+}
\ No newline at end of file