Unified: Support defer in CFG.

Author: aschackmull

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
   }
@@ -74,7 +79,25 @@ module Ast implements AstSig<Location> {
 
   class Expr = U::Expr;
 
-  class BlockStmt = U::Block;
+  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() }
@@ -150,12 +173,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 {

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

@@ -147,9 +147,11 @@
 | 171 | cfg.swift:171:12:171:16 | false | 'false -^ ReturnExpr' |
 | 175 | cfg.swift:175:1:191:1 | func testDefer | 'func testDefer' |
 | 175 | cfg.swift:175:16:175:16 | x | 'x -> Block' |
-| 177 | cfg.swift:177:3:179:3 | DeferStmt | 'DeferStmt -V Block' |
+| 177 | cfg.swift:177:3:179:3 | DeferStmt | 'DeferStmt' |
+| 177 | cfg.swift:177:9:179:3 | Block | 'Block' |
 | 178 | cfg.swift:178:5:178:9 | print | 'print -> Argument -V "4" -^ print(...)' |
-| 181 | cfg.swift:181:3:183:3 | DeferStmt | 'DeferStmt -V Block' |
+| 181 | cfg.swift:181:3:183:3 | DeferStmt | 'DeferStmt' |
+| 181 | cfg.swift:181:9:183:3 | Block | 'Block' |
 | 182 | cfg.swift:182:5:182:9 | print | 'print -> Argument -V "3" -^ print(...)' |
 | 185 | cfg.swift:185:3:190:3 | DeferStmt | 'DeferStmt -V Block' |
 | 186 | cfg.swift:186:5:186:9 | print | 'print -> Argument -V "1" -^ print(...)' |

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

@@ -1,4 +1,8 @@
 bbContinues
+| cfg.swift:177:3:177:3 | DeferStmt | 'DeferStmt goto DeferStmt(+4)' |
+| cfg.swift:181:3:181:3 | DeferStmt | 'DeferStmt goto DeferStmt(+4)' |
+| cfg.swift:182:5:182:14 | print(...) | 'print(...) goto Block(-5)' |
+| cfg.swift:188:7:188:16 | print(...) | 'print(...) goto Block(-7)' |
 bbStep
 | cfg.swift:20:9:20:14 | ... >= ... | '... >= ... : false -> Block(+0)' |
 | cfg.swift:20:9:20:14 | ... >= ... | '... >= ... : true -> GuardIfStmt(+3)' |

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

@@ -174,18 +174,18 @@ func patterns(x : Int) -> Bool {
 
 func testDefer(x : inout Int) {
   // Will print 1, 2, 3, 4
-  defer {
+  defer { // $ bbContinues='DeferStmt goto DeferStmt(+4)'
     print("4")
   }
 
-  defer {
-    print("3")
+  defer { // $ bbContinues='DeferStmt goto DeferStmt(+4)'
+    print("3") // $ bbContinues='print(...) goto Block(-5)'
   }
 
   defer {
     print("1")
      defer {
-      print("2")
+      print("2") // $ bbContinues='print(...) goto Block(-7)'
     }
   }
 }