Merge pull request github#11014 from tamasvajk/kotlin-for-loop-var

Kotlin: Resugar `for` loops

Author: tamasvajk

java/kotlin-extractor/src/main/kotlin/KotlinFileExtractor.kt

@@ -1494,7 +1494,7 @@ open class KotlinFileExtractor(
         }
     }
 
-    private fun extractVariableExpr(v: IrVariable, callable: Label<out DbCallable>, parent: Label<out DbExprparent>, idx: Int, enclosingStmt: Label<out DbStmt>) {
+    private fun extractVariableExpr(v: IrVariable, callable: Label<out DbCallable>, parent: Label<out DbExprparent>, idx: Int, enclosingStmt: Label<out DbStmt>, extractInitializer: Boolean = true) {
         with("variable expr", v) {
             val varId = useVariable(v)
             val exprId = tw.getFreshIdLabel<DbLocalvariabledeclexpr>()
@@ -1509,7 +1509,7 @@ open class KotlinFileExtractor(
             tw.writeCallableEnclosingExpr(exprId, callable)
             tw.writeStatementEnclosingExpr(exprId, enclosingStmt)
             val i = v.initializer
-            if (i != null) {
+            if (i != null && extractInitializer) {
                 extractExpressionExpr(i, callable, exprId, 0, enclosingStmt)
             }
             if (!v.isVar) {
@@ -3181,6 +3181,90 @@ open class KotlinFileExtractor(
         }
     }
 
+    /**
+     * This method tries to extract a block as an enhanced for loop.
+     * It returns true if it succeeds, and false otherwise.
+     */
+    private fun tryExtractForLoop(e: IrContainerExpression, callable: Label<out DbCallable>, parent: StmtExprParent): Boolean {
+        /*
+         * We're expecting the pattern
+         * {
+         *   val iterator = [expr].iterator()
+         *   while (iterator.hasNext()) {
+         *    val [loopVar] = iterator.next()
+         *    [block]
+         *   }
+         * }
+         */
+
+        if (e.origin != IrStatementOrigin.FOR_LOOP ||
+            e.statements.size != 2) {
+            return false
+        }
+
+        val iteratorVariable = e.statements[0] as? IrVariable
+        val innerWhile = e.statements[1] as? IrWhileLoop
+
+        if (iteratorVariable == null ||
+            iteratorVariable.origin != IrDeclarationOrigin.FOR_LOOP_ITERATOR ||
+            innerWhile == null ||
+            innerWhile.origin != IrStatementOrigin.FOR_LOOP_INNER_WHILE) {
+            return false
+        }
+
+        val initializer = iteratorVariable.initializer as? IrCall
+        if (initializer == null ||
+            initializer.origin != IrStatementOrigin.FOR_LOOP_ITERATOR ||
+            initializer.symbol.owner.name.asString() != "iterator") {
+            return false
+        }
+
+        val expr = initializer.dispatchReceiver
+        val cond = innerWhile.condition as? IrCall
+        val body = innerWhile.body as? IrBlock
+
+        if (expr == null ||
+            cond == null ||
+            cond.origin != IrStatementOrigin.FOR_LOOP_HAS_NEXT ||
+            (cond.dispatchReceiver as? IrGetValue)?.symbol?.owner != iteratorVariable ||
+            body == null ||
+            body.origin != IrStatementOrigin.FOR_LOOP_INNER_WHILE ||
+            body.statements.size < 2) {
+            return false
+        }
+
+        val loopVar = body.statements[0] as? IrVariable
+        val nextCall = loopVar?.initializer as? IrCall
+
+        if (loopVar == null ||
+            !(loopVar.origin == IrDeclarationOrigin.FOR_LOOP_VARIABLE || loopVar.origin == IrDeclarationOrigin.IR_TEMPORARY_VARIABLE) ||
+            nextCall == null ||
+            nextCall.origin != IrStatementOrigin.FOR_LOOP_NEXT ||
+            (nextCall.dispatchReceiver as? IrGetValue)?.symbol?.owner != iteratorVariable) {
+            return false
+        }
+
+        val id = extractLoop(innerWhile, null, parent, callable) { p, idx ->
+            tw.getFreshIdLabel<DbEnhancedforstmt>().also {
+                tw.writeStmts_enhancedforstmt(it, p, idx, callable)
+            }
+        }
+
+        extractVariableExpr(loopVar, callable, id, 0, id, extractInitializer = false)
+        extractExpressionExpr(expr, callable, id, 1, id)
+        val block = body.statements[1] as? IrBlock
+        if (body.statements.size == 2 && block != null) {
+            // Extract the body that was given to us by the compiler
+            extractExpressionStmt(block, callable, id, 2)
+        } else {
+            // Extract a block with all but the first (loop variable declaration) statement
+            extractBlock(body, body.statements.takeLast(body.statements.size - 1), id, 2, callable)
+        }
+
+        return true
+    }
+
+
     /**
      * This tried to extract a block as an array update.
      * It returns true if it succeeds, and false otherwise.
@@ -3412,22 +3496,17 @@ open class KotlinFileExtractor(
                     }
                 }
                 is IrContainerExpression -> {
-                    if(!tryExtractArrayUpdate(e, callable, parent)) {
-                        val stmtParent = parent.stmt(e, callable)
-                        val id = tw.getFreshIdLabel<DbBlock>()
-                        val locId = tw.getLocation(e)
-                        tw.writeStmts_block(id, stmtParent.parent, stmtParent.idx, callable)
-                        tw.writeHasLocation(id, locId)
-                        e.statements.forEachIndexed { i, s ->
-                            extractStatement(s, callable, id, i)
-                        }
+                    if (!tryExtractArrayUpdate(e, callable, parent) &&
+                        !tryExtractForLoop(e, callable, parent)) {
+
+                        extractBlock(e, e.statements, parent, callable)
                     }
                 }
                 is IrWhileLoop -> {
-                    extractLoop(e, parent, callable)
+                    extractLoopWithCondition(e, parent, callable)
                 }
                 is IrDoWhileLoop -> {
-                    extractLoop(e, parent, callable)
+                    extractLoopWithCondition(e, parent, callable)
                 }
                 is IrInstanceInitializerCall -> {
                     val irConstructor = declarationStack.peek().first as? IrConstructor
@@ -3892,6 +3971,32 @@ open class KotlinFileExtractor(
         }
     }
 
+    private fun extractBlock(
+        e: IrContainerExpression,
+        statements: List<IrStatement>,
+        parent: StmtExprParent,
+        callable: Label<out DbCallable>
+    ) {
+        val stmtParent = parent.stmt(e, callable)
+        extractBlock(e, statements, stmtParent.parent, stmtParent.idx, callable)
+    }
+
+    private fun extractBlock(
+        e: IrElement,
+        statements: List<IrStatement>,
+        parent: Label<out DbStmtparent>,
+        idx: Int,
+        callable: Label<out DbCallable>
+    ) {
+        val id = tw.getFreshIdLabel<DbBlock>()
+        val locId = tw.getLocation(e)
+        tw.writeStmts_block(id, parent, idx, callable)
+        tw.writeHasLocation(id, locId)
+        statements.forEachIndexed { i, s ->
+            extractStatement(s, callable, id, i)
+        }
+    }
+
     private inline fun <D: DeclarationDescriptor, reified B: IrSymbolOwner> getBoundSymbolOwner(symbol: IrBindableSymbol<D, B>, e: IrExpression): B? {
         if (symbol.isBound) {
             return symbol.owner
@@ -4000,9 +4105,11 @@ open class KotlinFileExtractor(
 
     private fun extractLoop(
         loop: IrLoop,
+        bodyIdx: Int?,
         stmtExprParent: StmtExprParent,
-        callable: Label<out DbCallable>
-    ) {
+        callable: Label<out DbCallable>,
+        getId: (Label<out DbStmtparent>, Int) -> Label<out DbStmt>
+    ) : Label<out DbStmt> {
         val stmtParent = stmtExprParent.stmt(loop, callable)
         val locId = tw.getLocation(loop)
 
@@ -4023,22 +4130,34 @@ open class KotlinFileExtractor(
             parent = stmtParent.parent
         }
 
-        val id = if (loop is IrWhileLoop) {
-            val id = tw.getFreshIdLabel<DbWhilestmt>()
-            tw.writeStmts_whilestmt(id, parent, idx, callable)
-            id
-        } else {
-            val id = tw.getFreshIdLabel<DbDostmt>()
-            tw.writeStmts_dostmt(id, parent, idx, callable)
-            id
-        }
-
+        val id = getId(parent, idx)
         tw.writeHasLocation(id, locId)
-        extractExpressionExpr(loop.condition, callable, id, 0, id)
+
         val body = loop.body
-        if (body != null) {
-            extractExpressionStmt(body, callable, id, 1)
+        if (body != null && bodyIdx != null) {
+            extractExpressionStmt(body, callable, id, bodyIdx)
         }
+
+        return id
+    }
+
+    private fun extractLoopWithCondition(
+        loop: IrLoop,
+        stmtExprParent: StmtExprParent,
+        callable: Label<out DbCallable>
+    ) {
+        val id = extractLoop(loop, 1, stmtExprParent, callable) { parent, idx ->
+            if (loop is IrWhileLoop) {
+                tw.getFreshIdLabel<DbWhilestmt>().also {
+                    tw.writeStmts_whilestmt(it, parent, idx, callable)
+                }
+            } else {
+                tw.getFreshIdLabel<DbDostmt>().also {
+                    tw.writeStmts_dostmt(it, parent, idx, callable)
+                }
+            }
+        }
+        extractExpressionExpr(loop.condition, callable, id, 0, id)
     }
 
     private fun IrValueParameter.isExtensionReceiver(): Boolean {

java/ql/test/kotlin/library-tests/arrays/arrayIterators.expected

@@ -1,6 +1,3 @@
-| arrayIterators.kt:6:15:6:15 | iterator(...) | iterator(java.lang.Object[]) | kotlin.jvm.internal.ArrayIteratorKt |
-| arrayIterators.kt:7:15:7:15 | iterator(...) | iterator(int[]) | kotlin.jvm.internal.ArrayIteratorsKt |
-| arrayIterators.kt:8:15:8:15 | iterator(...) | iterator(boolean[]) | kotlin.jvm.internal.ArrayIteratorsKt |
 | arrayIterators.kt:10:16:10:25 | iterator(...) | iterator(java.lang.Object[]) | kotlin.jvm.internal.ArrayIteratorKt |
 | arrayIterators.kt:11:16:11:25 | iterator(...) | iterator(int[]) | kotlin.jvm.internal.ArrayIteratorsKt |
 | arrayIterators.kt:12:16:12:25 | iterator(...) | iterator(boolean[]) | kotlin.jvm.internal.ArrayIteratorsKt |

java/ql/test/kotlin/library-tests/for-array-iterators/test.expected

@@ -1,9 +0,0 @@
-| test.kt:5:14:5:14 | hasNext(...) |
-| test.kt:5:14:5:14 | iterator(...) |
-| test.kt:5:14:5:14 | next(...) |
-| test.kt:6:14:6:14 | hasNext(...) |
-| test.kt:6:14:6:14 | iterator(...) |
-| test.kt:6:14:6:14 | next(...) |
-| test.kt:7:14:7:14 | hasNext(...) |
-| test.kt:7:14:7:14 | iterator(...) |
-| test.kt:7:14:7:14 | next(...) |