Unified: Support loop guards.

Author: aschackmull

shared/controlflow/codeql/controlflow/ControlFlowGraph.qll

@@ -161,6 +161,9 @@ signature module AstSig<LocationSig Location> {
     Expr getCollection();
   }
 
+  /** Gets the loop guard of the given `loop`, if any. */
+  default AstNode getLoopGuard(ForEachStmt loop) { none() }
+
   /**
    * A `break` statement.
    *
@@ -656,6 +659,7 @@ module Make0<LocationSig Location, AstSig<Location> Ast> {
         any(DoStmt dostmt).getCondition() = n or
         any(UntilStmt untilstmt).getCondition() = n or
         any(ForStmt forstmt).getCondition() = n or
+        getLoopGuard(_) = n or
         any(ConditionalExpr condexpr).getCondition() = n or
         any(CatchClause catch).getCondition() = n or
         any(Case case).getGuard() = n
@@ -1701,8 +1705,19 @@ module Make0<LocationSig Location, AstSig<Location> Ast> {
           n2.isBefore(foreachstmt.getVariable())
           or
           n1.isAfter(foreachstmt.getVariable()) and
+          (
+            n2.isBefore(getLoopGuard(foreachstmt))
+            or
+            not exists(getLoopGuard(foreachstmt)) and
+            n2.isBefore(foreachstmt.getBody())
+          )
+          or
+          n1.isAfterTrue(getLoopGuard(foreachstmt)) and
           n2.isBefore(foreachstmt.getBody())
           or
+          n1.isAfterFalse(getLoopGuard(foreachstmt)) and
+          n2.isAdditional(foreachstmt, loopHeaderTag())
+          or
           n1.isAfter(foreachstmt.getBody()) and
           n2.isAdditional(foreachstmt, loopHeaderTag())
           or

unified/ql/lib/codeql/unified/internal/ControlFlowGraph.qll

@@ -125,13 +125,13 @@ module Ast implements AstSig<Location> {
   class ForEachStmt extends LoopStmt instanceof U::ForEachStmt {
     override Stmt getBody() { result = U::ForEachStmt.super.getBody() }
 
-    // TODO support foreach guard
-    //
     Expr getVariable() { result = super.getPattern() }
 
     Expr getCollection() { result = super.getIterable() }
   }
 
+  AstNode getLoopGuard(ForEachStmt foreachstmt) { result = foreachstmt.(U::ForEachStmt).getGuard() }
+
   class BreakStmt = U::BreakExpr;
 
   class ContinueStmt = U::ContinueExpr;

unified/ql/test/library-tests/controlflow/basicblock-slices.expected

@@ -118,7 +118,8 @@
 | 140 | cfg.swift:140:7:140:7 | _ | '_' |
 | 141 | cfg.swift:141:9:141:12 | Block | 'Block' |
 | 143 | cfg.swift:143:3:145:3 | ForEachStmt | 'ForEachStmt -V 0 -> 10 -^ ... ... ...' |
-| 143 | cfg.swift:143:7:143:7 | y | 'y -> Block' |
+| 143 | cfg.swift:143:7:143:7 | y | 'y -> y -> 5 -^ ... != ...' |
+| 143 | cfg.swift:143:32:145:3 | Block | 'Block' |
 | 144 | cfg.swift:144:5:144:9 | print | 'print -> Argument -V y -^ print(...)' |
 | 147 | cfg.swift:147:3:160:3 | SwitchExpr | 'SwitchExpr -V x' |
 | 148 | cfg.swift:148:5:150:17 | SwitchCase | 'SwitchCase -V 0 -> 1 -^ OrPattern' |

unified/ql/test/library-tests/controlflow/cfg.expected

@@ -20,6 +20,9 @@ bbStep
 | cfg.swift:141:9:141:12 | Block | 'Block : successor -> _(-1)' |
 | cfg.swift:143:12:143:17 | ... ... ... | '... ... ... : empty -> SwitchExpr(+4)' |
 | cfg.swift:143:12:143:17 | ... ... ... | '... ... ... : non-empty -> y(+0)' |
+| cfg.swift:143:25:143:30 | ... != ... | '... != ... : false -> SwitchExpr(+4)' |
+| cfg.swift:143:25:143:30 | ... != ... | '... != ... : false -> y(+0)' |
+| cfg.swift:143:25:143:30 | ... != ... | '... != ... : true -> Block(+0)' |
 | cfg.swift:144:5:144:12 | print(...) | 'print(...) : successor -> SwitchExpr(+3)' |
 | cfg.swift:144:5:144:12 | print(...) | 'print(...) : successor -> y(-1)' |
 | cfg.swift:148:10:148:13 | OrPattern | 'OrPattern : match -> Block(+1)' |

unified/ql/test/library-tests/controlflow/cfg.swift

@@ -140,7 +140,7 @@ func patterns(x : Int) -> Bool {
   for _ in 0...10 // $ bbStep='... ... ... : empty -> ForEachStmt(+3)' bbStep='... ... ... : non-empty -> _(+0)'
         {  } // $ bbStep='Block : successor -> _(-1)' bbStep='Block : successor -> ForEachStmt(+2)'
 
-  for y in 0...10 where y != 5 { // $ bbStep='... ... ... : empty -> SwitchExpr(+4)' bbStep='... ... ... : non-empty -> y(+0)'
+  for y in 0...10 where y != 5 { // $ bbStep='... ... ... : empty -> SwitchExpr(+4)' bbStep='... ... ... : non-empty -> y(+0)' bbStep='... != ... : false -> SwitchExpr(+4)' bbStep='... != ... : false -> y(+0)' bbStep='... != ... : true -> Block(+0)'
     print(y) // $ bbStep='print(...) : successor -> SwitchExpr(+3)' bbStep='print(...) : successor -> y(-1)'
   }