Merge pull request #22815 from aschackmull/unified/cfg-updates

Unified: Add CFG support for loop guards and defer

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

@@ -35,12 +35,17 @@ module Ast implements AstSig<Location> {
     not e.fromSource()
   }
 
+  private AstNode adjustedGetChild(AstNode n, int index) { deferScope(n, index, result) }
+
   AstNode getChild(AstNode n, int index) {
     result.getParent() = n and
     result.getParentIndex() = index and
     not n instanceof Callable and
     not skipControlFlow(n) and
-    not skipControlFlow(result)
+    not skipControlFlow(result) and
+    not adjustedGetChild(_, _) = result
+    or
+    adjustedGetChild(n, index) = result
     or
     n.(FunctionExpr).getCaptureDeclaration(index) = result
     or
@@ -76,7 +81,56 @@ module Ast implements AstSig<Location> {
 
   class Expr = U::Expr;
 
-  class BlockStmt = U::Block;
+  /**
+   * Holds if `s` is the `n`th statement in the constructed `try` block of
+   * `defer`.
+   *
+   * For example, the following `defer` statements
+   * ```
+   * {
+   *   s1
+   *   defer { s5 }
+   *   s2
+   *   defer { s4 }
+   *   s3
+   * }
+   * ```
+   * are transformed into these `try`-`finally` blocks:
+   * ```
+   * {
+   *   s1
+   *   try {
+   *     s2
+   *     try {
+   *       s3
+   *     } finally {
+   *       s4
+   *     }
+   *   } finally {
+   *     s5
+   *   }
+   * }
+   * ```
+   */
+  private predicate deferScope(U::DeferStmt defer, int n, Stmt s) {
+    exists(U::Block block, int deferIdx | block.getStmt(deferIdx) = defer |
+      block.getStmt(deferIdx + 1) = s and n = 0
+      or
+      exists(Stmt mid |
+        deferScope(defer, n - 1, mid) and
+        not mid instanceof U::DeferStmt and
+        block.getStmt(deferIdx + 1 + n) = s
+      )
+    )
+  }
+
+  class BlockStmt extends Stmt instanceof U::Block {
+    Stmt getStmt(int n) { result = super.getStmt(n) and not deferScope(_, _, result) }
+
+    Stmt getLastStmt() {
+      exists(int i | result = this.getStmt(i) and not exists(this.getStmt(i + 1)))
+    }
+  }
 
   class ExprStmt extends Stmt {
     ExprStmt() { none() }
@@ -127,13 +181,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;
@@ -152,12 +206,18 @@ module Ast implements AstSig<Location> {
     Expr getExpr() { result = super.getValue() }
   }
 
-  class TryStmt extends U::TryExpr {
-    AstNode getBody(int index) { index = 0 and result = super.getBody() }
+  class TryStmt extends Stmt {
+    TryStmt() { this instanceof U::TryExpr or this instanceof U::DeferStmt }
+
+    AstNode getBody(int index) {
+      index = 0 and result = this.(U::TryExpr).getBody()
+      or
+      deferScope(this, index, result)
+    }
 
-    CatchClause getCatch(int index) { result = super.getCatchClause(index) }
+    CatchClause getCatch(int index) { result = this.(U::TryExpr).getCatchClause(index) }
 
-    Stmt getFinally() { none() }
+    Stmt getFinally() { result = this.(U::DeferStmt).getBody() }
   }
 
   class CatchClause extends U::CatchClause {