Rust: Support break and continue in loops

Author: paldepind

rust/ql/lib/codeql/rust/controlflow/internal/Completion.qll

@@ -6,6 +6,8 @@ private import SuccessorType
 private newtype TCompletion =
   TSimpleCompletion() or
   TBooleanCompletion(Boolean b) or
+  TBreakCompletion() or
+  TContinueCompletion() or
   TReturnCompletion()
 
 /** A completion of a statement or an expression. */
@@ -72,6 +74,28 @@ class BooleanCompletion extends ConditionalCompletion, TBooleanCompletion {
   override string toString() { result = "boolean(" + value + ")" }
 }
 
+/**
+ * A completion that represents a break.
+ */
+class BreakCompletion extends TBreakCompletion, Completion {
+  override BreakSuccessor getAMatchingSuccessorType() { any() }
+
+  override predicate isValidForSpecific(AstNode e) { e instanceof BreakExpr }
+
+  override string toString() { result = "break" }
+}
+
+/**
+ * A completion that represents a continue.
+ */
+class ContinueCompletion extends TContinueCompletion, Completion {
+  override ContinueSuccessor getAMatchingSuccessorType() { any() }
+
+  override predicate isValidForSpecific(AstNode e) { e instanceof ContinueExpr }
+
+  override string toString() { result = "continue" }
+}
+
 /**
  * A completion that represents a return.
  */

rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll

@@ -141,17 +141,48 @@ class LoopExprTree extends PostOrderTree instanceof LoopExpr {
   override predicate first(AstNode node) { first(super.getBody(), node) }
 
   override predicate succ(AstNode pred, AstNode succ, Completion c) {
-    // Edge from the last node in the body to the loop itself
+    // Edge back to the start for final expression and continue expressions
     last(super.getBody(), pred, c) and
-    completionIsNormal(c) and
+    (completionIsNormal(c) or c instanceof ContinueCompletion) and
+    this.first(succ)
+    or
+    // Edge for exiting the loop with a break expressions
+    last(super.getBody(), pred, c) and
+    c instanceof BreakCompletion and
     succ = this
+  }
+}
+
+class ReturnExprTree extends PostOrderTree instanceof ReturnExpr {
+  override predicate propagatesAbnormal(AstNode child) { child = super.getExpr() }
+
+  override predicate first(AstNode node) {
+    first(super.getExpr(), node)
+    or
+    not super.hasExpr() and node = this
+  }
+
+  override predicate succ(AstNode pred, AstNode succ, Completion c) {
+    last(super.getExpr(), pred, c) and succ = this
+  }
+}
+
+class BreakExprTree extends PostOrderTree instanceof BreakExpr {
+  override predicate propagatesAbnormal(AstNode child) { child = super.getExpr() }
+
+  override predicate first(AstNode node) {
+    first(super.getExpr(), node)
     or
-    // Tie the knot with an edge from the loop back to the first node
-    pred = this and
-    first(super.getBody(), succ)
+    not super.hasExpr() and node = this
+  }
+
+  override predicate succ(AstNode pred, AstNode succ, Completion c) {
+    last(super.getExpr(), pred, c) and succ = this
   }
 }
 
+class ContinueExprTree extends LeafTree instanceof ContinueExpr { }
+
 class LiteralExprTree extends LeafTree instanceof LiteralExpr { }
 
 class PathExprTree extends LeafTree instanceof PathExpr { }

rust/ql/lib/codeql/rust/controlflow/internal/SuccessorType.qll

@@ -4,6 +4,8 @@ cached
 newtype TSuccessorType =
   TSuccessorSuccessor() or
   TBooleanSuccessor(Boolean b) or
+  TBreakSuccessor() or
+  TContinueSuccessor() or
   TReturnSuccessor()
 
 /** The type of a control flow successor. */
@@ -32,11 +34,21 @@ abstract private class ConditionalSuccessor extends SuccessorTypeImpl {
   override string toString() { result = this.getValue().toString() }
 }
 
-/** A Boolean control flow successor. */
+/** A boolean control flow successor for a boolean conditon. */
 final class BooleanSuccessor extends ConditionalSuccessor, TBooleanSuccessor {
   BooleanSuccessor() { this = TBooleanSuccessor(value) }
 }
 
+/** A `break` control flow successor. */
+final class BreakSuccessor extends SuccessorTypeImpl, TBreakSuccessor {
+  final override string toString() { result = "break" }
+}
+
+/** A `continue` control flow successor. */
+final class ContinueSuccessor extends SuccessorTypeImpl, TContinueSuccessor {
+  final override string toString() { result = "continue" }
+}
+
 /**
  * A `return` control flow successor.
  */

rust/ql/test/library-tests/controlflow/test.rs

@@ -6,16 +6,31 @@ fn main() -> i64 {
     }
 }
 
-fn spin(n: i64) {
-    let mut i = 0;
+fn next(n: i64) -> i64 {
+    if n % 2 == 0 {
+        n / 2
+    } else {
+        3 * n + 1
+    }
+}
+
+fn spin(n: i64) -> bool {
+    let mut i = n;
     loop {
-        i += 1;
+        i = next(i);
+        if i == 1 {
+            break;
+        }
+        if i % 2 != 0 {
+            continue;
+        }
+        i = i / 2
     }
+    return true;
 }
 
 fn decrement(n: i64) -> i64 {
-    12;
-    if n == 0 {
+    if n <= 0 {
         0
     } else {
         n - 1