Swift: extract closure captures (hand-written part)

Author: d10c

swift/extractor/infra/SwiftDispatcher.h

@@ -29,7 +29,8 @@ class SwiftDispatcher {
                                const swift::Expr*,
                                const swift::Pattern*,
                                const swift::TypeRepr*,
-                               const swift::TypeBase*>;
+                               const swift::TypeBase*,
+                               const swift::CapturedValue*>;
 
   template <typename E>
   static constexpr bool IsStorable = std::is_constructible_v<Store::Handle, const E&>;
@@ -218,6 +219,10 @@ class SwiftDispatcher {
     attachLocation(locatable->getStartLoc(), locatable->getEndLoc(), locatableLabel);
   }
 
+  void attachLocation(const swift::CapturedValue* capture, TrapLabel<LocatableTag> locatableLabel) {
+    attachLocation(capture->getLoc(), locatableLabel);
+  }
+
   void attachLocation(const swift::IfConfigClause* clause, TrapLabel<LocatableTag> locatableLabel) {
     attachLocation(clause->Loc, clause->Loc, locatableLabel);
   }
@@ -364,6 +369,7 @@ class SwiftDispatcher {
   virtual void visit(const swift::Pattern* pattern) = 0;
   virtual void visit(const swift::TypeRepr* typeRepr, swift::Type type) = 0;
   virtual void visit(const swift::TypeBase* type) = 0;
+  virtual void visit(const swift::CapturedValue* capture) = 0;
 
   const swift::SourceManager& sourceManager;
   TrapDomain& trap;

swift/extractor/infra/SwiftTagTraits.h

@@ -32,6 +32,7 @@ MAP(std::filesystem::path, DbFileTag)
 MAP(swift::StmtCondition, StmtConditionTag)
 MAP(swift::StmtConditionElement, ConditionElementTag)
 MAP(swift::CaseLabelItem, CaseLabelItemTag)
+MAP(swift::CapturedValue, CapturedDeclTag)
 
 MAP(swift::Stmt, StmtTag)
   MAP(swift::BraceStmt, BraceStmtTag)

swift/extractor/translators/DeclTranslator.cpp

@@ -439,4 +439,14 @@ codeql::MissingMemberDecl DeclTranslator::translateMissingMemberDecl(
   entry.name = decl.getName().getBaseName().userFacingName().str();
   return entry;
 }
+
+codeql::CapturedDecl DeclTranslator::translateCapturedValue(const swift::CapturedValue& capture) {
+  codeql::CapturedDecl entry{dispatcher.template assignNewLabel(capture)};
+  auto decl = capture.getDecl();
+  entry.decl = dispatcher.fetchLabel(decl);
+  entry.module = dispatcher.fetchLabel(decl->getModuleContext());
+  entry.is_direct = capture.isDirect();
+  entry.is_escaping = !capture.isNoEscape();
+  return entry;
+}
 }  // namespace codeql

swift/extractor/translators/DeclTranslator.h

@@ -48,6 +48,7 @@ class DeclTranslator : public AstTranslatorBase<DeclTranslator> {
   std::optional<codeql::OpaqueTypeDecl> translateOpaqueTypeDecl(const swift::OpaqueTypeDecl& decl);
   codeql::PoundDiagnosticDecl translatePoundDiagnosticDecl(const swift::PoundDiagnosticDecl& decl);
   codeql::MissingMemberDecl translateMissingMemberDecl(const swift::MissingMemberDecl& decl);
+  codeql::CapturedDecl translateCapturedValue(const swift::CapturedValue& capture);
 
  private:
   std::string mangledName(const swift::ValueDecl& decl);

swift/extractor/translators/ExprTranslator.cpp

@@ -477,6 +477,7 @@ void ExprTranslator::fillAbstractClosureExpr(const swift::AbstractClosureExpr& e
   assert(expr.getParameters() && "AbstractClosureExpr has getParameters()");
   entry.params = dispatcher.fetchRepeatedLabels(*expr.getParameters());
   entry.body = dispatcher.fetchLabel(expr.getBody());
+  entry.captures = dispatcher.fetchRepeatedLabels(expr.getCaptureInfo().getCaptures());
 }
 
 TrapLabel<ArgumentTag> ExprTranslator::emitArgument(const swift::Argument& arg) {

swift/extractor/translators/SwiftVisitor.h

@@ -39,6 +39,10 @@ class SwiftVisitor : private SwiftDispatcher {
     typeTranslator.translateAndEmit(*typeRepr, type);
   }
 
+  void visit(const swift::CapturedValue* capture) override {
+    declTranslator.translateAndEmit(*capture);
+  }
+
   DeclTranslator declTranslator{*this};
   ExprTranslator exprTranslator{*this};
   StmtTranslator stmtTranslator{*this};

swift/extractor/translators/TranslatorBase.h

@@ -118,6 +118,10 @@ class AstTranslatorBase : private swift::ASTVisitor<CrtpSubclass>,
     swift::ASTVisitor<CrtpSubclass>::visit(const_cast<E*>(&entity));
   }
 
+  void translateAndEmit(const swift::CapturedValue& e) {
+    dispatcher.emit(static_cast<CrtpSubclass*>(this)->translateCapturedValue(e));
+  }
+
  private:
   friend class swift::ASTVisitor<CrtpSubclass>;
 

swift/ql/lib/codeql/swift/elements/decl/CapturedDecl.qll

@@ -0,0 +1,5 @@
+private import codeql.swift.generated.decl.CapturedDecl
+
+class CapturedDecl extends Generated::CapturedDecl {
+  override string toString() { result = this.getDecl().toString() }
+}

swift/schema.py

@@ -340,7 +340,7 @@ class Argument(Locatable):
     expr: Expr | child
 
 class AbstractClosureExpr(Expr, Callable):
-    pass
+    captures: list["CapturedDecl"] | child
 
 class AnyTryExpr(Expr):
     sub_expr: Expr | child
@@ -362,6 +362,11 @@ class AssignExpr(Expr):
 class BindOptionalExpr(Expr):
     sub_expr: Expr | child
 
+class CapturedDecl(Decl):
+    decl: ValueDecl | doc("the declaration captured by the parent closure")
+    is_direct: predicate
+    is_escaping: predicate
+
 class CaptureListExpr(Expr):
     binding_decls: list[PatternBindingDecl] | child
     closure_body: "ClosureExpr" | child