Swift: add params to CFG

Author: Robert Marsh

swift/ql/lib/codeql/swift/controlflow/internal/ControlFlowElements.qll

@@ -3,6 +3,7 @@ private import swift
 cached
 newtype TControlFlowElement =
   TAstElement(AstNode n) or
+  TFuncDeclElement(AbstractFunctionDecl func) { func.hasBody() } or
   TPropertyGetterElement(Decl accessor, Expr ref) { isPropertyGetterElement(accessor, ref) } or
   TPropertySetterElement(AccessorDecl accessor, AssignExpr assign) {
     isPropertySetterElement(accessor, assign)
@@ -161,3 +162,13 @@ class PropertyObserverElement extends ControlFlowElement, TPropertyObserverEleme
 
   AssignExpr getAssignExpr() { result = assign }
 }
+
+
+class FuncDeclElement extends ControlFlowElement, TFuncDeclElement {
+  AbstractFunctionDecl func;
+  FuncDeclElement() { this = TFuncDeclElement(func) }
+
+  override string toString() { result = func.toString() }
+
+  override Location getLocation() { result = func.getLocation() }
+}

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

@@ -48,15 +48,15 @@ module CfgScope {
 
   private class BodyStmtCallableScope extends Range_ instanceof AbstractFunctionDecl {
     final override predicate entry(ControlFlowElement first) {
-      exists(Stmts::BraceStmtTree tree |
-        tree.getAst() = super.getBody() and
-        tree.firstInner(first)
+      exists(Decls::FuncDeclTree tree |
+        tree.getAst() = this and
+        first = tree
       )
     }
 
     final override predicate exit(ControlFlowElement last, Completion c) {
-      exists(Stmts::BraceStmtTree tree |
-        tree.getAst() = super.getBody() and
+      exists(Decls::FuncDeclTree tree |
+        tree.getAst() = this and
         tree.last(last, c)
       )
     }
@@ -884,6 +884,21 @@ module Decls {
       )
     }
   }
+
+  class FuncDeclTree extends StandardPreOrderTree, TFuncDeclElement {
+    AbstractFunctionDecl ast;
+
+    FuncDeclTree() { this = TFuncDeclElement(ast) }
+
+    AbstractFunctionDecl getAst() { result = ast }
+
+    final override ControlFlowElement getChildElement(int i) {
+      result.asAstNode() = ast.getParam(i)
+      or
+      result.asAstNode() = ast.getBody() and
+      i = ast.getNumberOfParams()
+    }
+  }
 }
 
 module Exprs {