Swift: Extract 'yield' statements.

Author: MathiasVP

swift/codegen/schema.yml

@@ -601,6 +601,7 @@ ThrowStmt:
 
 YieldStmt:
   _extends: Stmt
+  results: Expr*
 
 BoundGenericType:
   _extends: NominalOrBoundGenericNominalType

swift/extractor/visitors/StmtVisitor.h

@@ -206,6 +206,16 @@ class StmtVisitor : public AstVisitorBase<StmtVisitor> {
     dispatcher_.emit(FallthroughStmtsTrap{label, sourceLabel, destLabel});
   }
 
+  void visitYieldStmt(swift::YieldStmt* stmt) {
+    auto label = dispatcher_.assignNewLabel(stmt);
+    dispatcher_.emit(YieldStmtsTrap{label});
+    auto i = 0u;
+    for(auto* expr : stmt->getYields()) {
+      auto exprLabel = dispatcher_.fetchLabel(expr);
+      dispatcher_.emit(YieldStmtResultsTrap{label, i++, exprLabel});
+    }
+  }
+
  private:
   void emitLabeledStmt(const swift::LabeledStmt* stmt, TrapLabel<LabeledStmtTag> label) {
     if (stmt->getLabelInfo()) {

swift/ql/lib/codeql/swift/generated/stmt/YieldStmt.qll

@@ -1,6 +1,18 @@
 // generated by codegen/codegen.py
+import codeql.swift.elements.expr.Expr
 import codeql.swift.elements.stmt.Stmt
 
 class YieldStmtBase extends @yield_stmt, Stmt {
   override string getAPrimaryQlClass() { result = "YieldStmt" }
+
+  Expr getResult(int index) {
+    exists(Expr x |
+      yield_stmt_results(this, index, x) and
+      result = x.resolve()
+    )
+  }
+
+  Expr getAResult() { result = getResult(_) }
+
+  int getNumberOfResults() { result = count(getAResult()) }
 }

swift/ql/lib/swift.dbscheme

@@ -1273,6 +1273,13 @@ yield_stmts(
   unique int id: @yield_stmt
 );
 
+#keyset[id, index]
+yield_stmt_results(
+  int id: @yield_stmt ref,
+  int index: int ref,
+  int result: @expr ref
+);
+
 @bound_generic_type =
   @bound_generic_class_type
 | @bound_generic_enum_type

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

@@ -7,12 +7,14 @@
 | test.swift:2:7:2:7 | get |
 | test.swift:2:7:2:7 | self |
 | test.swift:2:7:2:7 | self |
+| test.swift:2:7:2:7 | self |
 | test.swift:2:7:2:7 | set |
 | test.swift:2:7:2:7 | value |
 | test.swift:2:7:2:7 | x |
 | test.swift:3:3:6:3 | var ... = ... |
 | test.swift:3:7:3:7 | (unnamed function decl) |
 | test.swift:3:7:3:7 | next |
+| test.swift:3:7:3:7 | self |
 | test.swift:4:5:4:5 | self |
 | test.swift:4:5:4:24 | get |
 | test.swift:5:5:5:5 | self |
@@ -26,6 +28,7 @@
 | test.swift:9:17:9:17 | get |
 | test.swift:9:17:9:17 | self |
 | test.swift:9:17:9:17 | self |
+| test.swift:9:17:9:17 | self |
 | test.swift:9:17:9:17 | set |
 | test.swift:9:17:9:17 | value |
 | test.swift:9:17:9:17 | x |
@@ -79,6 +82,7 @@
 | test.swift:41:7:41:7 | get |
 | test.swift:41:7:41:7 | self |
 | test.swift:41:7:41:7 | self |
+| test.swift:41:7:41:7 | self |
 | test.swift:41:7:41:7 | set |
 | test.swift:41:7:41:7 | value |
 | test.swift:42:3:42:3 | self |
@@ -111,6 +115,7 @@
 | test.swift:81:8:81:8 | normalField |
 | test.swift:82:3:87:3 | var ... = ... |
 | test.swift:82:7:82:7 | (unnamed function decl) |
+| test.swift:82:7:82:7 | self |
 | test.swift:82:7:82:7 | settableField |
 | test.swift:83:5:83:5 | newValue |
 | test.swift:83:5:83:11 | set |
@@ -127,9 +132,11 @@
 | test.swift:102:7:102:7 | normalField |
 | test.swift:102:7:102:7 | self |
 | test.swift:102:7:102:7 | self |
+| test.swift:102:7:102:7 | self |
 | test.swift:102:7:102:7 | set |
 | test.swift:102:7:102:7 | value |
 | test.swift:104:3:104:3 | (unnamed function decl) |
+| test.swift:104:3:104:3 | self |
 | test.swift:104:3:109:3 | subscript ... |
 | test.swift:104:13:104:13 | x |
 | test.swift:104:13:104:13 | x |
@@ -149,6 +156,7 @@
 | test.swift:115:7:115:7 | hasWillSet1 |
 | test.swift:115:7:115:7 | self |
 | test.swift:115:7:115:7 | self |
+| test.swift:115:7:115:7 | self |
 | test.swift:115:7:115:7 | set |
 | test.swift:115:7:115:7 | value |
 | test.swift:116:5:116:25 | willSet |
@@ -159,6 +167,7 @@
 | test.swift:119:7:119:7 | hasWillSet2 |
 | test.swift:119:7:119:7 | self |
 | test.swift:119:7:119:7 | self |
+| test.swift:119:7:119:7 | self |
 | test.swift:119:7:119:7 | set |
 | test.swift:119:7:119:7 | value |
 | test.swift:120:5:120:5 | newValue |
@@ -169,6 +178,7 @@
 | test.swift:123:7:123:7 | hasDidSet1 |
 | test.swift:123:7:123:7 | self |
 | test.swift:123:7:123:7 | self |
+| test.swift:123:7:123:7 | self |
 | test.swift:123:7:123:7 | set |
 | test.swift:123:7:123:7 | value |
 | test.swift:124:5:124:24 | didSet |
@@ -189,6 +199,7 @@
 | test.swift:131:7:131:7 | hasBoth |
 | test.swift:131:7:131:7 | self |
 | test.swift:131:7:131:7 | self |
+| test.swift:131:7:131:7 | self |
 | test.swift:131:7:131:7 | set |
 | test.swift:131:7:131:7 | value |
 | test.swift:132:5:132:5 | newValue |

swift/ql/test/extractor-tests/statements/YieldStmts.expected

@@ -0,0 +1,2 @@
+| main.swift:75:7:75:7 | yield ... | 0 | file://:0:0:0:0 | &... |
+| main.swift:78:7:78:14 | yield ... | 0 | main.swift:78:13:78:14 | &... |

swift/ql/test/extractor-tests/statements/YieldStmts.ql

@@ -0,0 +1,5 @@
+import swift
+
+from YieldStmt yield, int i
+where yield.getLocation().getFile().getName().matches("%swift/ql/test%")
+select yield, i, yield.getResult(i)

swift/ql/test/extractor-tests/statements/main.swift

@@ -70,3 +70,16 @@ if case .some(_) = x {
 let numbers = [1, 2, 3]
 for number in numbers where number % 2 == 0 {
 }
+
+struct HasModifyAccessorDecl {
+  var x : Int
+  var hasModify : Int {
+    _modify {
+      yield &x
+    }
+
+    get {
+      return 0
+    }
+  }
+}

swift/ql/test/library-tests/controlflow/graph/Cfg.expected

@@ -5129,8 +5129,6 @@ cfg.swift:
 #-----|  -> [...]
 
 #  394| (unnamed function decl)
-#-----|  -> yield ...
-#-----|  -> TBD (YieldStmt)
 
 #  394| enter (unnamed function decl)
 #-----|  -> (unnamed function decl)
@@ -5141,11 +5139,6 @@ cfg.swift:
 #  394| enter set
 #-----|  -> set
 
-#  394| exit (unnamed function decl)
-
-#  394| exit (unnamed function decl) (normal)
-#-----|  -> exit (unnamed function decl)
-
 #  394| exit get
 
 #  394| exit get (normal)
@@ -5163,12 +5156,6 @@ cfg.swift:
 
 #  394| value
 
-#  394| yield ...
-#-----|  -> exit (unnamed function decl) (normal)
-
-#  394| TBD (YieldStmt)
-#-----|  -> exit (unnamed function decl) (normal)
-
 #  395| deinit
 #-----|  -> self