Merge pull request #9371 from MathiasVP/extract-key-path-application

Swift: Extract KeyPath applications

Author: MathiasVP

swift/codegen/schema.yml

@@ -390,6 +390,8 @@ InOutExpr:
 
 KeyPathApplicationExpr:
   _extends: Expr
+  base: Expr
+  key_path: Expr
 
 KeyPathDotExpr:
   _extends: Expr

swift/extractor/visitors/ExprVisitor.h

@@ -468,7 +468,7 @@ class ExprVisitor : public AstVisitorBase<ExprVisitor> {
         auto pathLabel = dispatcher_.fetchLabel(path);
         dispatcher_.emit(KeyPathExprParsedPathsTrap{label, pathLabel});
       }
-      if (auto root = expr->getParsedPath()) {
+      if (auto root = expr->getParsedRoot()) {
         auto rootLabel = dispatcher_.fetchLabel(root);
         dispatcher_.emit(KeyPathExprParsedRootsTrap{label, rootLabel});
       }
@@ -509,6 +509,22 @@ class ExprVisitor : public AstVisitorBase<ExprVisitor> {
     dispatcher_.emit(IfExprsTrap{label, condLabel, thenLabel, elseLabel});
   }
 
+  void visitKeyPathDotExpr(swift::KeyPathDotExpr* expr) {
+    auto label = dispatcher_.assignNewLabel(expr);
+    dispatcher_.emit(KeyPathDotExprsTrap{label});
+  }
+
+  void visitKeyPathApplicationExpr(swift::KeyPathApplicationExpr* expr) {
+    auto label = dispatcher_.assignNewLabel(expr);
+    assert(expr->getBase() && "KeyPathApplicationExpr has getBase()");
+    assert(expr->getKeyPath() && "KeyPathApplicationExpr has getKeyPath()");
+    
+    auto baseLabel = dispatcher_.fetchLabel(expr->getBase());
+    auto keyPathLabel = dispatcher_.fetchLabel(expr->getKeyPath());
+
+    dispatcher_.emit(KeyPathApplicationExprsTrap{label, baseLabel, keyPathLabel});
+  }
+
  private:
   TrapLabel<ArgumentTag> emitArgument(const swift::Argument& arg) {
     auto argLabel = dispatcher_.createLabel<ArgumentTag>();

swift/ql/lib/codeql/swift/generated/expr/KeyPathApplicationExpr.qll

@@ -3,4 +3,18 @@ import codeql.swift.elements.expr.Expr
 
 class KeyPathApplicationExprBase extends @key_path_application_expr, Expr {
   override string getAPrimaryQlClass() { result = "KeyPathApplicationExpr" }
+
+  Expr getBase() {
+    exists(Expr x |
+      key_path_application_exprs(this, x, _) and
+      result = x.resolve()
+    )
+  }
+
+  Expr getKeyPath() {
+    exists(Expr x |
+      key_path_application_exprs(this, _, x) and
+      result = x.resolve()
+    )
+  }
 }

swift/ql/lib/swift.dbscheme

@@ -888,7 +888,9 @@ in_out_exprs(
 );
 
 key_path_application_exprs(
-  unique int id: @key_path_application_expr
+  unique int id: @key_path_application_expr,
+  int base: @expr ref,
+  int key_path: @expr ref
 );
 
 key_path_dot_exprs(

swift/ql/test/extractor-tests/expressions/all.expected

@@ -125,6 +125,7 @@
 | expressions.swift:54:1:54:1 | _ |
 | expressions.swift:54:1:54:8 |  ... = ... |
 | expressions.swift:54:5:54:8 | #keyPath(...) |
+| expressions.swift:54:6:54:8 | TBD (UnresolvedDotExpr) |
 | expressions.swift:58:16:58:16 | 1234 |
 | expressions.swift:59:1:59:1 | unsafeFunction |
 | expressions.swift:59:1:59:34 | call to unsafeFunction |
@@ -232,3 +233,15 @@
 | expressions.swift:138:10:138:17 | ...[...] |
 | expressions.swift:138:13:138:13 | 3 |
 | expressions.swift:138:16:138:16 | 4 |
+| expressions.swift:152:26:152:26 | a |
+| expressions.swift:152:26:152:47 | \\...[...] |
+| expressions.swift:152:37:152:37 | keyPathInt |
+| expressions.swift:153:24:153:24 | a |
+| expressions.swift:153:24:153:43 | \\...[...] |
+| expressions.swift:153:35:153:35 | keyPathB |
+| expressions.swift:154:22:154:22 | a |
+| expressions.swift:154:22:154:41 | \\...[...] |
+| expressions.swift:154:22:154:56 | \\...[...] |
+| expressions.swift:154:33:154:33 | keyPathB |
+| expressions.swift:154:52:154:55 | #keyPath(...) |
+| expressions.swift:154:53:154:55 | TBD (UnresolvedDotExpr) |

swift/ql/test/extractor-tests/expressions/expressions.swift

@@ -136,4 +136,20 @@ func testProperties(hp : inout HasProperty) -> Int {
   var w = hp.normalField
   hp[1] = 2
   return hp[3, 4]
+}
+
+struct B {
+  var x : Int
+}
+
+struct A {
+  var b : B
+  var bs : [B]
+  var mayB : B?
+}
+
+func test(a : A, keyPathInt : WritableKeyPath<A, Int>, keyPathB : WritableKeyPath<A, B>) {
+  var apply_keyPathInt = a[keyPath: keyPathInt]
+  var apply_keyPathB = a[keyPath: keyPathB]
+  var nested_apply = a[keyPath: keyPathB][keyPath: \B.x]
 }

swift/ql/test/extractor-tests/expressions/semantics.expected

@@ -68,3 +68,9 @@
 | expressions.swift:137:3:137:7 | ...[...] | OrdinarySemantics |
 | expressions.swift:138:10:138:10 | hp | OrdinarySemantics |
 | expressions.swift:138:10:138:17 | ...[...] | OrdinarySemantics |
+| expressions.swift:152:26:152:26 | a | OrdinarySemantics |
+| expressions.swift:152:37:152:37 | keyPathInt | OrdinarySemantics |
+| expressions.swift:153:24:153:24 | a | OrdinarySemantics |
+| expressions.swift:153:35:153:35 | keyPathB | OrdinarySemantics |
+| expressions.swift:154:22:154:22 | a | OrdinarySemantics |
+| expressions.swift:154:33:154:33 | keyPathB | OrdinarySemantics |