Kotlin: Resugar for loops

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

@@ -1499,7 +1499,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>()
@@ -1514,7 +1514,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) {
@@ -3174,6 +3174,85 @@ 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 ||
+            (e.statements[0] as? IrVariable)?.origin != IrDeclarationOrigin.FOR_LOOP_ITERATOR ||
+            (e.statements[1] as? IrWhileLoop)?.origin != IrStatementOrigin.FOR_LOOP_INNER_WHILE) {
+            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 block = body.statements[1] as? IrBlock
+        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, block, 2, parent, callable) { p, idx ->
+            val loopId = tw.getFreshIdLabel<DbEnhancedforstmt>()
+            tw.writeStmts_enhancedforstmt(loopId, p, idx, callable)
+            loopId
+        }
+
+        extractExpressionExpr(expr, callable, id, 1, id)
+        extractVariableExpr(loopVar, callable, id, 0, id, extractInitializer = false)
+
+        return true
+    }
+
+
     /**
      * This tried to extract a block as an array update.
      * It returns true if it succeeds, and false otherwise.
@@ -3405,7 +3484,9 @@ open class KotlinFileExtractor(
                     }
                 }
                 is IrContainerExpression -> {
-                    if(!tryExtractArrayUpdate(e, callable, parent)) {
+                    if (!tryExtractArrayUpdate(e, callable, parent) &&
+                        !tryExtractForLoop(e, callable, parent)) {
+
                         val stmtParent = parent.stmt(e, callable)
                         val id = tw.getFreshIdLabel<DbBlock>()
                         val locId = tw.getLocation(e)
@@ -3417,10 +3498,10 @@ open class KotlinFileExtractor(
                     }
                 }
                 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
@@ -3993,9 +4074,12 @@ open class KotlinFileExtractor(
 
     private fun extractLoop(
         loop: IrLoop,
+        body: IrExpression?,
+        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)
 
@@ -4016,7 +4100,23 @@ open class KotlinFileExtractor(
             parent = stmtParent.parent
         }
 
-        val id = if (loop is IrWhileLoop) {
+        val id = getId(parent, idx)
+        tw.writeHasLocation(id, locId)
+
+        if (body != null) {
+            extractExpressionStmt(body, callable, id, bodyIdx)
+        }
+
+        return id
+    }
+
+    private fun extractLoopWithCondition(
+        loop: IrLoop,
+        stmtExprParent: StmtExprParent,
+        callable: Label<out DbCallable>
+    ) {
+        val id = extractLoop(loop, loop.body, 1, stmtExprParent, callable) { parent, idx ->
+            if (loop is IrWhileLoop) {
                 val id = tw.getFreshIdLabel<DbWhilestmt>()
                 tw.writeStmts_whilestmt(id, parent, idx, callable)
                 id
@@ -4025,13 +4125,8 @@ open class KotlinFileExtractor(
                 tw.writeStmts_dostmt(id, parent, idx, callable)
                 id
             }
-
-        tw.writeHasLocation(id, locId)
-        extractExpressionExpr(loop.condition, callable, id, 0, id)
-        val body = loop.body
-        if (body != null) {
-            extractExpressionStmt(body, callable, id, 1)
         }
+        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(...) |

java/ql/test/kotlin/library-tests/stmts/PrintAst.expected

@@ -116,22 +116,13 @@ stmts.kt:
 #   28|             0: [VarAccess] x
 #   28|             1: [VarAccess] y
 #   29|           1: [ContinueStmt] continue
-#   31|         2: [BlockStmt] { ... }
+#   31|         2: [EnhancedForStmt] for (... : ...)
-#   31|           0: [LocalVariableDeclStmt] var ...;
+#-----|           0: (Single Local Variable Declaration)
-#   31|             1: [LocalVariableDeclExpr] tmp0_iterator
+#   31|             1: [LocalVariableDeclExpr] i
-#   31|               0: [MethodAccess] iterator(...)
+#   31|           1: [MethodAccess] rangeTo(...)
-#   31|                 -1: [MethodAccess] rangeTo(...)
 #   31|             -1: [VarAccess] x
 #   31|             0: [VarAccess] y
-#   31|           1: [WhileStmt] while (...)
+#   31|           2: [BlockStmt] { ... }
-#   31|             0: [MethodAccess] hasNext(...)
-#   31|               -1: [VarAccess] tmp0_iterator
-#   31|             1: [BlockStmt] { ... }
-#   31|               0: [LocalVariableDeclStmt] var ...;
-#   31|                 1: [LocalVariableDeclExpr] i
-#   31|                   0: [MethodAccess] next(...)
-#   31|                     -1: [VarAccess] tmp0_iterator
-#   31|               1: [BlockStmt] { ... }
 #   32|             0: [ExprStmt] <Expr>;
 #   32|               0: [WhenExpr] when ...
 #   32|                 0: [WhenBranch] ... -> ...
@@ -139,23 +130,14 @@ stmts.kt:
 #   32|                     0: [VarAccess] x
 #   32|                     1: [VarAccess] y
 #   32|                   1: [BreakStmt] break
-#   35|         3: [BlockStmt] { ... }
+#   35|         3: [LabeledStmt] <Label>: ...
-#   35|           0: [LocalVariableDeclStmt] var ...;
+#   35|           0: [EnhancedForStmt] for (... : ...)
-#   35|             1: [LocalVariableDeclExpr] tmp1_iterator
+#-----|             0: (Single Local Variable Declaration)
-#   35|               0: [MethodAccess] iterator(...)
+#   35|               1: [LocalVariableDeclExpr] i
-#   35|                 -1: [MethodAccess] rangeTo(...)
+#   35|             1: [MethodAccess] rangeTo(...)
 #   35|               -1: [VarAccess] x
 #   35|               0: [VarAccess] y
-#   35|           1: [LabeledStmt] <Label>: ...
+#   35|             2: [BlockStmt] { ... }
-#   35|             0: [WhileStmt] while (...)
-#   35|               0: [MethodAccess] hasNext(...)
-#   35|                 -1: [VarAccess] tmp1_iterator
-#   35|               1: [BlockStmt] { ... }
-#   35|                 0: [LocalVariableDeclStmt] var ...;
-#   35|                   1: [LocalVariableDeclExpr] i
-#   35|                     0: [MethodAccess] next(...)
-#   35|                       -1: [VarAccess] tmp1_iterator
-#   35|                 1: [BlockStmt] { ... }
 #   36|               0: [BlockStmt] { ... }
 #   36|                 0: [DoStmt] do ... while (...)
 #   38|                   0: [GTExpr] ... > ...

java/ql/test/kotlin/library-tests/stmts/exprs.expected

@@ -64,13 +64,7 @@
 | stmts.kt:28:15:28:15 | y | VarAccess |
 | stmts.kt:31:10:31:10 | i | LocalVariableDeclExpr |
 | stmts.kt:31:15:31:15 | x | VarAccess |
-| stmts.kt:31:15:31:18 | hasNext(...) | MethodAccess |
-| stmts.kt:31:15:31:18 | iterator(...) | MethodAccess |
-| stmts.kt:31:15:31:18 | next(...) | MethodAccess |
 | stmts.kt:31:15:31:18 | rangeTo(...) | MethodAccess |
-| stmts.kt:31:15:31:18 | tmp0_iterator | LocalVariableDeclExpr |
-| stmts.kt:31:15:31:18 | tmp0_iterator | VarAccess |
-| stmts.kt:31:15:31:18 | tmp0_iterator | VarAccess |
 | stmts.kt:31:18:31:18 | y | VarAccess |
 | stmts.kt:32:9:32:24 | when ... | WhenExpr |
 | stmts.kt:32:13:32:13 | x | VarAccess |
@@ -78,13 +72,7 @@
 | stmts.kt:32:17:32:17 | y | VarAccess |
 | stmts.kt:35:18:35:18 | i | LocalVariableDeclExpr |
 | stmts.kt:35:23:35:23 | x | VarAccess |
-| stmts.kt:35:23:35:26 | hasNext(...) | MethodAccess |
-| stmts.kt:35:23:35:26 | iterator(...) | MethodAccess |
-| stmts.kt:35:23:35:26 | next(...) | MethodAccess |
 | stmts.kt:35:23:35:26 | rangeTo(...) | MethodAccess |
-| stmts.kt:35:23:35:26 | tmp1_iterator | LocalVariableDeclExpr |
-| stmts.kt:35:23:35:26 | tmp1_iterator | VarAccess |
-| stmts.kt:35:23:35:26 | tmp1_iterator | VarAccess |
 | stmts.kt:35:26:35:26 | y | VarAccess |
 | stmts.kt:37:13:37:36 | when ... | WhenExpr |
 | stmts.kt:37:17:37:17 | x | VarAccess |

java/ql/test/kotlin/library-tests/stmts/loops.expected

@@ -5,6 +5,6 @@ continueLabel
 jumpTarget
 | stmts.kt:25:24:25:33 | break | stmts.kt:23:11:27:5 | while (...) |
 | stmts.kt:29:9:29:16 | continue | stmts.kt:28:5:29:16 | while (...) |
-| stmts.kt:32:20:32:24 | break | stmts.kt:31:5:33:5 | while (...) |
+| stmts.kt:32:20:32:24 | break | stmts.kt:31:5:33:5 | for (... : ...) |
-| stmts.kt:37:24:37:36 | break | stmts.kt:35:13:39:5 | while (...) |
+| stmts.kt:37:24:37:36 | break | stmts.kt:35:13:39:5 | for (... : ...) |
 | stmts.kt:42:20:42:24 | break | stmts.kt:41:5:43:5 | while (...) |

java/ql/test/kotlin/library-tests/stmts/stmts.expected

@@ -40,22 +40,14 @@ enclosing
 | stmts.kt:25:24:25:33 | break | stmts.kt:25:13:25:33 | ... -> ... |
 | stmts.kt:28:5:29:16 | while (...) | stmts.kt:22:27:44:1 | { ... } |
 | stmts.kt:29:9:29:16 | continue | stmts.kt:28:5:29:16 | while (...) |
-| stmts.kt:31:5:33:5 | while (...) | stmts.kt:31:5:33:5 | { ... } |
+| stmts.kt:31:5:33:5 | for (... : ...) | stmts.kt:22:27:44:1 | { ... } |
-| stmts.kt:31:5:33:5 | { ... } | stmts.kt:22:27:44:1 | { ... } |
+| stmts.kt:31:21:33:5 | { ... } | stmts.kt:31:5:33:5 | for (... : ...) |
-| stmts.kt:31:5:33:5 | { ... } | stmts.kt:31:5:33:5 | while (...) |
-| stmts.kt:31:10:31:10 | var ...; | stmts.kt:31:5:33:5 | { ... } |
-| stmts.kt:31:15:31:18 | var ...; | stmts.kt:31:5:33:5 | { ... } |
-| stmts.kt:31:21:33:5 | { ... } | stmts.kt:31:5:33:5 | { ... } |
 | stmts.kt:32:9:32:24 | ... -> ... | stmts.kt:32:9:32:24 | <Expr>; |
 | stmts.kt:32:9:32:24 | <Expr>; | stmts.kt:31:21:33:5 | { ... } |
 | stmts.kt:32:20:32:24 | break | stmts.kt:32:9:32:24 | ... -> ... |
-| stmts.kt:35:13:39:5 | <Label>: ... | stmts.kt:35:13:39:5 | { ... } |
+| stmts.kt:35:13:39:5 | <Label>: ... | stmts.kt:22:27:44:1 | { ... } |
-| stmts.kt:35:13:39:5 | while (...) | stmts.kt:35:13:39:5 | <Label>: ... |
+| stmts.kt:35:13:39:5 | for (... : ...) | stmts.kt:35:13:39:5 | <Label>: ... |
-| stmts.kt:35:13:39:5 | { ... } | stmts.kt:22:27:44:1 | { ... } |
+| stmts.kt:35:29:39:5 | { ... } | stmts.kt:35:13:39:5 | for (... : ...) |
-| stmts.kt:35:13:39:5 | { ... } | stmts.kt:35:13:39:5 | while (...) |
-| stmts.kt:35:18:35:18 | var ...; | stmts.kt:35:13:39:5 | { ... } |
-| stmts.kt:35:23:35:26 | var ...; | stmts.kt:35:13:39:5 | { ... } |
-| stmts.kt:35:29:39:5 | { ... } | stmts.kt:35:13:39:5 | { ... } |
 | stmts.kt:36:9:38:25 | do ... while (...) | stmts.kt:36:9:38:25 | { ... } |
 | stmts.kt:36:9:38:25 | { ... } | stmts.kt:35:29:39:5 | { ... } |
 | stmts.kt:36:13:38:9 | { ... } | stmts.kt:36:9:38:25 | do ... while (...) |
@@ -125,21 +117,13 @@ enclosing
 | stmts.kt:25:24:25:33 | break | BreakStmt |
 | stmts.kt:28:5:29:16 | while (...) | WhileStmt |
 | stmts.kt:29:9:29:16 | continue | ContinueStmt |
-| stmts.kt:31:5:33:5 | while (...) | WhileStmt |
+| stmts.kt:31:5:33:5 | for (... : ...) | EnhancedForStmt |
-| stmts.kt:31:5:33:5 | { ... } | BlockStmt |
-| stmts.kt:31:5:33:5 | { ... } | BlockStmt |
-| stmts.kt:31:10:31:10 | var ...; | LocalVariableDeclStmt |
-| stmts.kt:31:15:31:18 | var ...; | LocalVariableDeclStmt |
 | stmts.kt:31:21:33:5 | { ... } | BlockStmt |
 | stmts.kt:32:9:32:24 | ... -> ... | WhenBranch |
 | stmts.kt:32:9:32:24 | <Expr>; | ExprStmt |
 | stmts.kt:32:20:32:24 | break | BreakStmt |
 | stmts.kt:35:13:39:5 | <Label>: ... | LabeledStmt |
-| stmts.kt:35:13:39:5 | while (...) | WhileStmt |
+| stmts.kt:35:13:39:5 | for (... : ...) | EnhancedForStmt |
-| stmts.kt:35:13:39:5 | { ... } | BlockStmt |
-| stmts.kt:35:13:39:5 | { ... } | BlockStmt |
-| stmts.kt:35:18:35:18 | var ...; | LocalVariableDeclStmt |
-| stmts.kt:35:23:35:26 | var ...; | LocalVariableDeclStmt |
 | stmts.kt:35:29:39:5 | { ... } | BlockStmt |
 | stmts.kt:36:9:38:25 | do ... while (...) | DoStmt |
 | stmts.kt:36:9:38:25 | { ... } | BlockStmt |