Unify loop break/continue statement handling between java and kotlin

Author: tamasvajk

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

@@ -2153,8 +2153,6 @@ open class KotlinFileExtractor(
         }
     }
 
-    private val loopIdMap: MutableMap<IrLoop, Label<out DbKtloopstmt>> = mutableMapOf()
-
     // todo: calculating the enclosing ref type could be done through this, instead of walking up the declaration parent chain
     private val declarationStack: Stack<IrDeclaration> = Stack()
 
@@ -2402,32 +2400,10 @@ open class KotlinFileExtractor(
                     }
                 }
                 is IrWhileLoop -> {
-                    val stmtParent = parent.stmt(e, callable)
-                    val id = tw.getFreshIdLabel<DbWhilestmt>()
-                    loopIdMap[e] = id
-                    val locId = tw.getLocation(e)
-                    tw.writeStmts_whilestmt(id, stmtParent.parent, stmtParent.idx, callable)
-                    tw.writeHasLocation(id, locId)
-                    extractExpressionExpr(e.condition, callable, id, 0, id)
-                    val body = e.body
-                    if(body != null) {
-                        extractExpressionStmt(body, callable, id, 1)
-                    }
-                    loopIdMap.remove(e)
+                    extractLoop(e, parent, callable)
                 }
                 is IrDoWhileLoop -> {
-                    val stmtParent = parent.stmt(e, callable)
-                    val id = tw.getFreshIdLabel<DbDostmt>()
-                    loopIdMap[e] = id
-                    val locId = tw.getLocation(e)
-                    tw.writeStmts_dostmt(id, stmtParent.parent, stmtParent.idx, callable)
-                    tw.writeHasLocation(id, locId)
-                    extractExpressionExpr(e.condition, callable, id, 0, id)
-                    val body = e.body
-                    if(body != null) {
-                        extractExpressionStmt(body, callable, id, 1)
-                    }
-                    loopIdMap.remove(e)
+                    extractLoop(e, parent, callable)
                 }
                 is IrInstanceInitializerCall -> {
                     val stmtParent = parent.stmt(e, callable)
@@ -2928,6 +2904,49 @@ open class KotlinFileExtractor(
         }
     }
 
+    private fun extractLoop(
+        loop: IrLoop,
+        stmtExprParent: StmtExprParent,
+        callable: Label<out DbCallable>
+    ) {
+        val stmtParent = stmtExprParent.stmt(loop, callable)
+        val locId = tw.getLocation(loop)
+
+        val idx: Int
+        val parent: Label<out DbStmtparent>
+
+        val label = loop.label
+        if (label != null) {
+            val labeledStmt = tw.getFreshIdLabel<DbLabeledstmt>()
+            tw.writeStmts_labeledstmt(labeledStmt, stmtParent.parent, stmtParent.idx, callable)
+            tw.writeHasLocation(labeledStmt, locId)
+
+            tw.writeNamestrings(label, "", labeledStmt)
+            idx = 0
+            parent = labeledStmt
+        } else {
+            idx = stmtParent.idx
+            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
+        }
+
+        tw.writeHasLocation(id, locId)
+        extractExpressionExpr(loop.condition, callable, id, 0, id)
+        val body = loop.body
+        if (body != null) {
+            extractExpressionStmt(body, callable, id, 1)
+        }
+    }
+
     private fun IrValueParameter.isExtensionReceiver(): Boolean {
         val parentFun = parent as? IrFunction ?: return false
         return parentFun.extensionReceiverParameter == this
@@ -4201,7 +4220,7 @@ open class KotlinFileExtractor(
 
     private fun extractBreakContinue(
         e: IrBreakContinue,
-        id: Label<out DbBreakcontinuestmt>
+        id: Label<out DbNamedexprorstmt>
     ) {
         with("break/continue", e) {
             val locId = tw.getLocation(e)
@@ -4210,14 +4229,6 @@ open class KotlinFileExtractor(
             if (label != null) {
                 tw.writeNamestrings(label, "", id)
             }
-
-            val loopId = loopIdMap[e.loop]
-            if (loopId == null) {
-                logger.errorElement("Missing break/continue target", e)
-                return
-            }
-
-            tw.writeKtBreakContinueTargets(id, loopId)
         }
     }
 

java/ql/lib/config/semmlecode.dbscheme

@@ -1165,17 +1165,6 @@ ktCommentOwners(
   int owner: @top ref
 )
 
-@breakcontinuestmt = @breakstmt
-                   | @continuestmt;
-
-@ktloopstmt = @whilestmt
-            | @dostmt;
-
-ktBreakContinueTargets(
-  unique int id: @breakcontinuestmt ref,
-  int target: @ktloopstmt ref
-)
-
 ktExtensionFunctions(
   unique int id: @method ref,
   int typeid: @type ref,

java/ql/lib/semmle/code/java/Statement.qll

@@ -888,27 +888,3 @@ class SuperConstructorInvocationStmt extends Stmt, ConstructorCall, @superconstr
 
   override string getAPrimaryQlClass() { result = "SuperConstructorInvocationStmt" }
 }
-
-/** A Kotlin loop statement. */
-class KtLoopStmt extends Stmt, @ktloopstmt {
-  KtLoopStmt() {
-    this instanceof WhileStmt or
-    this instanceof DoStmt
-  }
-}
-
-/** A Kotlin `break` or `continue` statement. */
-abstract class KtBreakContinueStmt extends Stmt, @breakcontinuestmt {
-  KtLoopStmt loop;
-
-  KtBreakContinueStmt() { ktBreakContinueTargets(this, loop) }
-
-  /** Gets the target loop statement of this `break`. */
-  KtLoopStmt getLoopStmt() { result = loop }
-}
-
-/** A Kotlin `break` statement. */
-class KtBreakStmt extends BreakStmt, KtBreakContinueStmt { }
-
-/** A Kotlin `continue` statement. */
-class KtContinueStmt extends ContinueStmt, KtBreakContinueStmt { }

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

@@ -1,7 +1,6 @@
 breakLabel
 | stmts.kt:25:24:25:33 | break | loop |
 continueLabel
-breakTarget
+jumpTarget
 | stmts.kt:25:24:25:33 | break | stmts.kt:23:11:27:5 | while (...) |
-continueTarget
 | stmts.kt:29:9:29:16 | continue | stmts.kt:28:5:29:16 | while (...) |

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

@@ -4,6 +4,4 @@ query predicate breakLabel(BreakStmt s, string label) { s.getLabel() = label }
 
 query predicate continueLabel(ContinueStmt s, string label) { s.getLabel() = label }
 
-query predicate breakTarget(KtBreakStmt s, KtLoopStmt l) { s.getLoopStmt() = l }
-
-query predicate continueTarget(KtContinueStmt s, KtLoopStmt l) { s.getLoopStmt() = l }
+query predicate jumpTarget(JumpStmt s, StmtParent p) { s.getTarget() = p }

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

@@ -80,6 +80,7 @@ enclosing
 | stmts.kt:18:52:18:52 | <Expr>; | ExprStmt |
 | stmts.kt:19:5:19:16 | return ... | ReturnStmt |
 | stmts.kt:22:27:30:1 | { ... } | BlockStmt |
+| stmts.kt:23:11:27:5 | <Label>: ... | LabeledStmt |
 | stmts.kt:23:11:27:5 | while (...) | WhileStmt |
 | stmts.kt:23:27:27:5 | { ... } | BlockStmt |
 | stmts.kt:24:9:26:25 | do ... while (...) | DoStmt |