KT-5405 J2K: convert for's through indices of some list or array into use of ".indices" + fixed a bug in for-statement conversion

#KT-5405 Fixed

annotations/com/intellij/psi/annotations.xml

@@ -249,6 +249,10 @@
   <item name='com.intellij.psi.PsiElement com.intellij.psi.PsiElement replace(com.intellij.psi.PsiElement)'>
     <annotation name='org.jetbrains.annotations.NotNull'/>
   </item>
+  <item
+      name='com.intellij.psi.PsiElementFactory.SERVICE com.intellij.psi.PsiElementFactory getInstance(com.intellij.openapi.project.Project)'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
   <item name='com.intellij.psi.PsiElementVisitor void visitBinaryFile(com.intellij.psi.PsiBinaryFile) 0'>
     <annotation name='org.jetbrains.annotations.NotNull'/>
   </item>

j2k/src/org/jetbrains/jet/j2k/ast/Expressions.kt

@@ -143,6 +143,12 @@ class StarExpression(val operand: Expression) : Expression() {
     }
 }
 
+class RangeExpression(val start: Expression, val end: Expression): Expression() {
+    override fun generateCode(builder: CodeBuilder) {
+        builder.appendOperand(this, start).append("..").appendOperand(this, end)
+    }
+}
+
 fun createArrayInitializerExpression(arrayType: ArrayType, initializers: List<Expression>, needExplicitType: Boolean) : MethodCallExpression {
     val elementType = arrayType.elementType
     val createArrayFunction = if (elementType is PrimitiveType)

j2k/src/org/jetbrains/jet/j2k/ast/Statements.kt

@@ -88,9 +88,10 @@ class DoWhileStatement(val condition: Expression, val body: Element, singleLine:
     }
 }
 
+//TODO: explicit type (if option)
 class ForeachStatement(
-        val variable: Parameter,
+        val variableName: Identifier,
-        val expression: Expression,
+        val collection: Expression,
         val body: Element,
         singleLine: Boolean
 ) : Statement() {
@@ -98,7 +99,7 @@ class ForeachStatement(
     private val br = if (singleLine) " " else "\n"
 
     override fun generateCode(builder: CodeBuilder) {
-        builder append "for (" append variable.identifier append " in " append expression append ")" append br append body
+        builder append "for (" append variableName append " in " append collection append ")" append br append body
     }
 }
 

j2k/src/org/jetbrains/jet/j2k/ast/Util.kt

@@ -46,19 +46,21 @@ private fun Expression.precedence(): Int? {
         is BinaryExpression -> when(op) {
             "*", "/", "%" -> 3
             "+", "-" -> 4
-            "?:" -> 6
+            "?:" -> 7
-            ">", "<", ">=", "<=" -> 8
+            ">", "<", ">=", "<=" -> 9
-            "==", "!=", "===", "!===" -> 9
+            "==", "!=", "===", "!===" -> 10
-            "&&" -> 10
+            "&&" -> 11
-            "||" -> 11
+            "||" -> 12
-            else -> 5 /* simple name */
+            else -> 6 /* simple name */
         }
 
-        is IsOperator -> 7
+        is RangeExpression -> 5
 
-        is IfStatement -> 12
+        is IsOperator -> 8
 
-        is AssignmentExpression -> 13
+        is IfStatement -> 13
+
+        is AssignmentExpression -> 14
 
         else -> null
     }

j2k/src/org/jetbrains/jet/j2k/visitors/StatementVisitor.kt

@@ -23,7 +23,7 @@ import org.jetbrains.jet.j2k.countWriteAccesses
 import java.util.ArrayList
 import org.jetbrains.jet.j2k.hasWriteAccesses
 import org.jetbrains.jet.j2k.isInSingleLine
-import org.jetbrains.jet.j2k.getContainingMethod
+import com.intellij.psi.tree.IElementType
 
 open class StatementVisitor(public val converter: Converter) : JavaElementVisitor() {
     public var result: Statement = Statement.Empty
@@ -102,83 +102,117 @@ open class StatementVisitor(public val converter: Converter) : JavaElementVisito
         val condition = statement.getCondition()
         val body = statement.getBody()
 
-        val loopVar = initialization?.getFirstChild() as? PsiLocalVariable
+        if (initialization is PsiDeclarationStatement && initialization.getFirstChild()!!.getNextSibling() == null) {
-        val onceWritableIterator = loopVar != null
+            val loopVar = initialization.getFirstChild() as? PsiLocalVariable
-                && !loopVar.hasWriteAccesses(body)
+            if (loopVar != null
-                && !loopVar.hasWriteAccesses(condition)
+                    && !loopVar.hasWriteAccesses(body)
-                && loopVar.countWriteAccesses(update) == 1
+                    && !loopVar.hasWriteAccesses(condition)
-
+                    && loopVar.countWriteAccesses(update) == 1
-        val operationTokenType = (condition as? PsiBinaryExpression)?.getOperationTokenType()
+                    && condition is PsiBinaryExpression) {
-        if (initialization is PsiDeclarationStatement
+                val operationTokenType = condition.getOperationTokenType()
-                && initialization.getFirstChild() == initialization.getLastChild()
+                val lowerBound = condition.getLOperand()
-                && condition != null
+                val upperBound = condition.getROperand()
-                && update != null
+                if ((operationTokenType == JavaTokenType.LT || operationTokenType == JavaTokenType.LE) &&
-                && update.getChildren().size == 1
+                        lowerBound is PsiReferenceExpression &&
-                && update.getChildren().single().isPlusPlusExpression()
+                        lowerBound.resolve() == loopVar &&
-                && (operationTokenType == JavaTokenType.LT || operationTokenType == JavaTokenType.LE)
+                        upperBound != null) {
-                && loopVar != null
+                    val start = loopVar.getInitializer()
-                && loopVar.getNameIdentifier() != null
+                    if (start != null
-                && onceWritableIterator) {
+                            && update != null
-            val end = converter.convertExpression((condition as PsiBinaryExpression).getROperand())
+                            && update.getChildren().size == 1
-            val endExpression = if (operationTokenType == JavaTokenType.LT)
+                            && update.getChildren().single().isPlusPlusExpression()) {
-                BinaryExpression(end, LiteralExpression("1").assignNoPrototype(), "-").assignNoPrototype()
+                        val range = forIterationRange(start, upperBound, operationTokenType).assignNoPrototype()
-            else
+                        result = ForeachStatement(loopVar.declarationIdentifier(), range, convertStatementOrBlock(body), statement.isInSingleLine())
-                end
+                        return
-            result = ForeachWithRangeStatement(loopVar.declarationIdentifier(),
-                                                 converter.convertExpression(loopVar.getInitializer()),
-                                                 endExpression,
-                                                 convertStatementOrBlock(body),
-                                                 statement.isInSingleLine())
-        }
-        else {
-            val initializationConverted = converter.convertStatement(initialization)
-            val updateConverted = converter.convertStatement(update)
-
-            val whileBody = if (updateConverted.isEmpty) {
-                convertStatementOrBlock(body)
-            }
-            else if (body is PsiBlockStatement) {
-                val nameConflict = initialization is PsiDeclarationStatement && initialization.getDeclaredElements().any { loopVar ->
-                    loopVar is PsiNamedElement && body.getCodeBlock().getStatements().any { statement ->
-                        statement is PsiDeclarationStatement && statement.getDeclaredElements().any {
-                            it is PsiNamedElement && it.getName() == loopVar.getName()
-                        }
                     }
                 }
+            }
+        }
 
-                if (nameConflict) {
+        val initializationConverted = converter.convertStatement(initialization)
-                    val statements = listOf(converter.convertStatement(body), updateConverted)
+        val updateConverted = converter.convertStatement(update)
-                    Block(statements, LBrace().assignNoPrototype(), RBrace().assignNoPrototype(), true).assignNoPrototype()
+
-                }
+        val whileBody = if (updateConverted.isEmpty) {
-                else {
+            convertStatementOrBlock(body)
-                    val block = converter.convertBlock(body.getCodeBlock(), true)
+        }
-                    Block(block.statements + listOf(updateConverted), block.lBrace, block.rBrace, true).assignPrototypesFrom(block)
+        else if (body is PsiBlockStatement) {
+            val nameConflict = initialization is PsiDeclarationStatement && initialization.getDeclaredElements().any { loopVar ->
+                loopVar is PsiNamedElement && body.getCodeBlock().getStatements().any { statement ->
+                    statement is PsiDeclarationStatement && statement.getDeclaredElements().any {
+                        it is PsiNamedElement && it.getName() == loopVar.getName()
+                    }
                 }
             }
-            else {
+
+            if (nameConflict) {
                 val statements = listOf(converter.convertStatement(body), updateConverted)
                 Block(statements, LBrace().assignNoPrototype(), RBrace().assignNoPrototype(), true).assignNoPrototype()
             }
-            val whileStatement = WhileStatement(
-                    if (condition != null) converter.convertExpression(condition) else LiteralExpression("true").assignNoPrototype(),
-                    whileBody,
-                    statement.isInSingleLine()).assignNoPrototype()
-
-            if (initializationConverted.isEmpty) {
-                result = whileStatement
-            }
             else {
-                val statements = listOf(initializationConverted, whileStatement)
+                val block = converter.convertBlock(body.getCodeBlock(), true)
-                val block = Block(statements, LBrace().assignNoPrototype(), RBrace().assignNoPrototype()).assignNoPrototype()
+                Block(block.statements + listOf(updateConverted), block.lBrace, block.rBrace, true).assignPrototypesFrom(block)
-                result = MethodCallExpression.build(null, "run", listOf(), listOf(), false, LambdaExpression(null, block))
             }
         }
+        else {
+            val statements = listOf(converter.convertStatement(body), updateConverted)
+            Block(statements, LBrace().assignNoPrototype(), RBrace().assignNoPrototype(), true).assignNoPrototype()
+        }
+        val whileStatement = WhileStatement(
+                if (condition != null) converter.convertExpression(condition) else LiteralExpression("true").assignNoPrototype(),
+                whileBody,
+                statement.isInSingleLine()).assignNoPrototype()
+
+        if (initializationConverted.isEmpty) {
+            result = whileStatement
+        }
+        else {
+            val statements = listOf(initializationConverted, whileStatement)
+            val block = Block(statements, LBrace().assignNoPrototype(), RBrace().assignNoPrototype()).assignNoPrototype()
+            result = MethodCallExpression.build(null, "run", listOf(), listOf(), false, LambdaExpression(null, block))
+        }
+    }
+
+    private fun forIterationRange(start: PsiExpression, upperBound: PsiExpression, comparisonTokenType: IElementType): Expression {
+        if (start is PsiLiteralExpression
+                && start.getValue() == 0
+                && comparisonTokenType == JavaTokenType.LT) {
+            // check if it's iteration through list indices
+            if (upperBound is PsiMethodCallExpression && upperBound.getArgumentList().getExpressions().isEmpty()) {
+                val methodExpr = upperBound.getMethodExpression()
+                if (methodExpr is PsiReferenceExpression && methodExpr.getReferenceName() == "size") {
+                    val qualifier = methodExpr.getQualifierExpression()
+                    if (qualifier is PsiReferenceExpression /* we don't convert to .indices if qualifier is method call or something because of possible side effects */) {
+                        val listType = PsiElementFactory.SERVICE.getInstance(converter.project).createTypeByFQClassName(CommonClassNames.JAVA_UTIL_LIST)
+                        val qualifierType = qualifier.getType()
+                        if (qualifierType != null && listType.isAssignableFrom(qualifierType)) {
+                            return QualifiedExpression(converter.convertExpression(qualifier), Identifier("indices", false).assignNoPrototype())
+                        }
+                    }
+                }
+            }
+
+            // check if it's iteration through array indices
+            if (upperBound is PsiReferenceExpression /* we don't convert to .indices if qualifier is method call or something because of possible side effects */
+                    && upperBound.getReferenceName() == "length") {
+                val qualifier = upperBound.getQualifierExpression()
+                if (qualifier is PsiReferenceExpression && qualifier.getType() is PsiArrayType) {
+                    return QualifiedExpression(converter.convertExpression(qualifier), Identifier("indices", false).assignNoPrototype())
+                }
+            }
+        }
+
+        val end = converter.convertExpression(upperBound)
+        val endExpression = if (comparisonTokenType == JavaTokenType.LT)
+            BinaryExpression(end, LiteralExpression("1").assignNoPrototype(), "-").assignNoPrototype()
+        else
+            end
+        return RangeExpression(converter.convertExpression(start), endExpression)
     }
 
     override fun visitForeachStatement(statement: PsiForeachStatement) {
         val iteratorExpr = converter.convertExpression(statement.getIteratedValue())
         val iterator = if (iteratorExpr.isNullable) BangBangExpression(iteratorExpr).assignNoPrototype() else iteratorExpr
-        result = ForeachStatement(converter.convertParameter(statement.getIterationParameter()),
+        result = ForeachStatement(statement.getIterationParameter().declarationIdentifier(),
                                   iterator,
                                   convertStatementOrBlock(statement.getBody()),
                                   statement.isInSingleLine())

j2k/tests/test/org/jetbrains/jet/j2k/test/JavaToKotlinConverterTestGenerated.java

@@ -1222,6 +1222,11 @@ public class JavaToKotlinConverterTestGenerated extends AbstractJavaToKotlinConv
             doTest("j2k/tests/testData/ast/for/commonCaseForTest.java");
         }
         
+        @TestMetadata("falseForRange.java")
+        public void testFalseForRange() throws Exception {
+            doTest("j2k/tests/testData/ast/for/falseForRange.java");
+        }
+        
         @TestMetadata("forRangeWithBlock.java")
         public void testForRangeWithBlock() throws Exception {
             doTest("j2k/tests/testData/ast/for/forRangeWithBlock.java");
@@ -1237,6 +1242,26 @@ public class JavaToKotlinConverterTestGenerated extends AbstractJavaToKotlinConv
             doTest("j2k/tests/testData/ast/for/forRangeWithLT.java");
         }
         
+        @TestMetadata("forThroughArrayIndices.java")
+        public void testForThroughArrayIndices() throws Exception {
+            doTest("j2k/tests/testData/ast/for/forThroughArrayIndices.java");
+        }
+        
+        @TestMetadata("forThroughListIndices.java")
+        public void testForThroughListIndices() throws Exception {
+            doTest("j2k/tests/testData/ast/for/forThroughListIndices.java");
+        }
+        
+        @TestMetadata("forThroughNonArrayIndices.java")
+        public void testForThroughNonArrayIndices() throws Exception {
+            doTest("j2k/tests/testData/ast/for/forThroughNonArrayIndices.java");
+        }
+        
+        @TestMetadata("forThroughNonCollectionIndices.java")
+        public void testForThroughNonCollectionIndices() throws Exception {
+            doTest("j2k/tests/testData/ast/for/forThroughNonCollectionIndices.java");
+        }
+        
         @TestMetadata("forWithBlock.java")
         public void testForWithBlock() throws Exception {
             doTest("j2k/tests/testData/ast/for/forWithBlock.java");

j2k/tests/testData/ast/for/falseForRange.java

@@ -0,0 +1,2 @@
+//statement
+for (int i = 0; i * 2 <= 10; i++) { foo(i); }

j2k/tests/testData/ast/for/falseForRange.kt

@@ -0,0 +1,7 @@
+run {
+    var i = 0
+    while (i * 2 <= 10) {
+        foo(i)
+        i++
+    }
+}

j2k/tests/testData/ast/for/forThroughArrayIndices.java

@@ -0,0 +1,3 @@
+//statement
+int[] array = new int[10];
+for (int i = 0; i < array.length; i++) {array[i] = i;}

j2k/tests/testData/ast/for/forThroughArrayIndices.kt

@@ -0,0 +1,4 @@
+val array = IntArray(10)
+for (i in array.indices) {
+    array[i] = i
+}

j2k/tests/testData/ast/for/forThroughListIndices.java

@@ -0,0 +1,17 @@
+//file
+import java.util.List;
+import java.util.ArrayList;
+
+class C{
+    void foo1(List<String> list) {
+        for (int i = 0; i < list.size(); i++) {
+            list.set(i, "a");
+        }
+    }
+
+    void foo2(ArrayList<String> list) {
+        for (int i = 0; i < list.size(); i++) {
+            list.set(i, "a");
+        }
+    }
+}

j2k/tests/testData/ast/for/forThroughListIndices.kt

@@ -0,0 +1,15 @@
+import java.util.ArrayList
+
+class C {
+    fun foo1(list: List<String>) {
+        for (i in list.indices) {
+            list.set(i, "a")
+        }
+    }
+
+    fun foo2(list: ArrayList<String>) {
+        for (i in list.indices) {
+            list.set(i, "a")
+        }
+    }
+}

j2k/tests/testData/ast/for/forThroughNonArrayIndices.java

@@ -0,0 +1,12 @@
+//file
+class X {
+    public int length = 5;
+}
+
+class C{
+    void foo(X x) {
+        for (int i = 0; i < x.length; i++) {
+            System.out.print(i);
+        }
+    }
+}

j2k/tests/testData/ast/for/forThroughNonArrayIndices.kt

@@ -0,0 +1,11 @@
+class X {
+    public var length: Int = 5
+}
+
+class C {
+    fun foo(x: X) {
+        for (i in 0..x.length - 1) {
+            System.out.print(i)
+        }
+    }
+}

j2k/tests/testData/ast/for/forThroughNonCollectionIndices.java

@@ -0,0 +1,12 @@
+//file
+class X {
+    public int size() { return 5; }
+}
+
+class C{
+    void foo(X x) {
+        for (int i = 0; i < x.size(); i++) {
+            System.out.print(i);
+        }
+    }
+}

j2k/tests/testData/ast/for/forThroughNonCollectionIndices.kt

@@ -0,0 +1,13 @@
+class X {
+    public fun size(): Int {
+        return 5
+    }
+}
+
+class C {
+    fun foo(x: X) {
+        for (i in 0..x.size() - 1) {
+            System.out.print(i)
+        }
+    }
+}