Merge pull request #9373 from MathiasVP/cfg-for-key-paths

Author: MathiasVP

swift/ql/lib/codeql/swift/controlflow/BasicBlocks.qll

@@ -207,6 +207,10 @@ private module JoinBlockPredecessors {
     isPropertySetterElement(n, _, result)
     or
     isPropertyObserverElement(n, _, result)
+    or
+    result = n.(KeyPathElement).getAst()
+    or
+    result = n.(FuncDeclElement).getAst()
   }
 
   int getId(JoinBlockPredecessor jbp) {

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

@@ -10,7 +10,8 @@ newtype TControlFlowElement =
   } or
   TPropertyObserverElement(AccessorDecl observer, AssignExpr assign) {
     isPropertyObserverElement(observer, assign)
-  }
+  } or
+  TKeyPathElement(KeyPathExpr expr)
 
 predicate isLValue(Expr e) { any(AssignExpr assign).getDest() = e }
 
@@ -171,4 +172,18 @@ class FuncDeclElement extends ControlFlowElement, TFuncDeclElement {
   override string toString() { result = func.toString() }
 
   override Location getLocation() { result = func.getLocation() }
+
+  AbstractFunctionDecl getAst() { result = func }
+}
+
+class KeyPathElement extends ControlFlowElement, TKeyPathElement {
+  KeyPathExpr expr;
+
+  KeyPathElement() { this = TKeyPathElement(expr) }
+
+  override Location getLocation() { result = expr.getLocation() }
+
+  KeyPathExpr getAst() { result = expr }
+
+  override string toString() { result = expr.toString() }
 }

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

@@ -47,19 +47,23 @@ module CfgScope {
   }
 
   private class BodyStmtCallableScope extends Range_ instanceof AbstractFunctionDecl {
-    final override predicate entry(ControlFlowElement first) {
-      exists(Decls::FuncDeclTree tree |
-        tree.getAst() = this and
-        first = tree
-      )
-    }
+    Decls::FuncDeclTree tree;
 
-    final override predicate exit(ControlFlowElement last, Completion c) {
-      exists(Decls::FuncDeclTree tree |
-        tree.getAst() = this and
-        tree.last(last, c)
-      )
-    }
+    BodyStmtCallableScope() { tree.getAst() = this }
+
+    final override predicate entry(ControlFlowElement first) { first(tree, first) }
+
+    final override predicate exit(ControlFlowElement last, Completion c) { last(tree, last, c) }
+  }
+
+  private class KeyPathScope extends Range_ instanceof KeyPathExpr {
+    AstControlFlowTree tree;
+
+    KeyPathScope() { tree.getAst() = this.getParsedRoot().getFullyConverted() }
+
+    final override predicate entry(ControlFlowElement first) { first(tree, first) }
+
+    final override predicate exit(ControlFlowElement last, Completion c) { last(tree, last, c) }
   }
 }
 
@@ -1112,6 +1116,20 @@ module Exprs {
     }
   }
 
+  class KeyPathTree extends AstLeafTree {
+    override KeyPathExpr ast;
+  }
+
+  class KeyPathApplicationTree extends AstStandardPostOrderTree {
+    override KeyPathApplicationExpr ast;
+
+    final override ControlFlowElement getChildElement(int i) {
+      i = 0 and result.asAstNode() = ast.getBase().getFullyConverted()
+      or
+      i = 1 and result.asAstNode() = ast.getKeyPath().getFullyConverted()
+    }
+  }
+
   private class InOutTree extends AstStandardPostOrderTree {
     override InOutExpr ast;
 
@@ -1648,7 +1666,9 @@ private module Cached {
     result = scopeOfAst(n.asAstNode()) or
     result = scopeOfAst(n.(PropertyGetterElement).getRef()) or
     result = scopeOfAst(n.(PropertySetterElement).getAssignExpr()) or
-    result = scopeOfAst(n.(PropertyObserverElement).getAssignExpr())
+    result = scopeOfAst(n.(PropertyObserverElement).getAssignExpr()) or
+    result = n.(FuncDeclElement).getAst() or
+    result = n.(KeyPathElement).getAst()
   }
 
   cached

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

@@ -2,7 +2,7 @@ private import swift
 private import codeql.swift.generated.GetImmediateParent
 
 module CallableBase {
-  class TypeRange = @abstract_function_decl;
+  class TypeRange = @abstract_function_decl or @key_path_expr;
 
   class Range extends Scope::Range, TypeRange { }
 }

swift/ql/lib/codeql/swift/dataflow/internal/DataFlowDispatch.qll

@@ -1,6 +1,7 @@
 private import swift
 private import DataFlowPrivate
 private import DataFlowPublic
+private import codeql.swift.controlflow.ControlFlowGraph
 
 newtype TReturnKind = TNormalReturnKind()
 
@@ -53,7 +54,7 @@ class DataFlowCall extends ExprNode {
 cached
 private module Cached {
   cached
-  newtype TDataFlowCallable = TDataFlowFunc(FuncDecl func)
+  newtype TDataFlowCallable = TDataFlowFunc(CfgScope scope)
 
   /** Gets a viable run-time target for the call `call`. */
   cached