Swift: introduce dispatcher

swift/codegen/cppgen.py

@@ -66,7 +66,7 @@ def generate(opts, renderer):
     processor = Processor({cls.name: cls for cls in schema.load(opts.schema).classes}, opts.trap_affix)
     out = opts.cpp_output
     renderer.render(cpp.ClassList(processor.get_classes(), opts.cpp_namespace, opts.trap_affix,
-                                  opts.cpp_include_dir), out / f"{opts.trap_affix}Classes.h")
+                                  opts.cpp_include_dir, opts.schema), out / f"{opts.trap_affix}Classes.h")
 
 
 tags = ("cpp", "schema")

swift/codegen/lib/cpp.py

@@ -104,6 +104,7 @@ class TrapList:
     namespace: str
     trap_affix: str
     include_dir: str
+    source: str
 
 
 @dataclass
@@ -112,6 +113,7 @@ class TagList:
 
     tags: List[Tag]
     namespace: str
+    source: str
 
 
 @dataclass
@@ -150,3 +152,4 @@ class ClassList:
     namespace: str
     trap_affix: str
     include_dir: str
+    source: str

swift/codegen/templates/cpp_classes.mustache

@@ -1,4 +1,4 @@
-// generated by {{generator}}
+// generated by {{generator}} from {{source}}
 // clang-format off
 #pragma once
 

swift/codegen/templates/trap_tags.mustache

@@ -1,4 +1,4 @@
-// generated by {{generator}}
+// generated by {{generator}} from {{source}}
 // clang-format off
 #pragma once
 

swift/codegen/templates/trap_traps.mustache

@@ -1,4 +1,4 @@
-// generated by {{generator}}
+// generated by {{generator}} from {{source}}
 // clang-format off
 #pragma once
 
@@ -29,12 +29,5 @@ inline std::ostream &operator<<(std::ostream &out, const {{name}}{{trap_affix}}
       << {{#get_streamer}}e.{{field_name}}{{/get_streamer}}{{/fields}} << ")";
   return out;
 }
-{{#id}}
-
-template <>
-struct TagToBindingTrapFunctor<typename {{type}}::Tag> {
-  using type = {{name}}{{trap_affix}};
-};
-{{/id}}
 {{/traps}}
 }

swift/codegen/trapgen.py

@@ -68,7 +68,7 @@ def generate(opts, renderer):
             for d in e.rhs:
                 tag_graph.setdefault(d.type, set()).add(e.lhs)
 
-    renderer.render(cpp.TrapList(traps, opts.cpp_namespace, opts.trap_affix, opts.cpp_include_dir),
+    renderer.render(cpp.TrapList(traps, opts.cpp_namespace, opts.trap_affix, opts.cpp_include_dir, opts.dbscheme),
                     out / f"{opts.trap_affix}Entries.h")
 
     tags = []
@@ -79,7 +79,7 @@ def generate(opts, renderer):
             index=index,
             id=tag,
         ))
-    renderer.render(cpp.TagList(tags, opts.cpp_namespace), out / f"{opts.trap_affix}Tags.h")
+    renderer.render(cpp.TagList(tags, opts.cpp_namespace, opts.dbscheme), out / f"{opts.trap_affix}Tags.h")
 
 
 tags = ("cpp", "dbscheme")

swift/extractor/BUILD.bazel

@@ -6,6 +6,8 @@ swift_cc_binary(
         "SwiftExtractor.cpp",
         "SwiftExtractor.h",
         "SwiftExtractorConfiguration.h",
+        "SwiftDispatcher.h",
+        "SwiftTagTraits.h",
         "main.cpp",
     ],
     visibility = ["//swift:__pkg__"],

swift/extractor/SwiftDispatcher.h

@@ -0,0 +1,179 @@
+#pragma once
+
+#include "swift/extractor/trap/TrapArena.h"
+#include "swift/extractor/trap/TrapLabelStore.h"
+#include "swift/extractor/trap/generated/TrapClasses.h"
+#include "swift/extractor/SwiftTagTraits.h"
+#include <swift/AST/SourceFile.h>
+#include <swift/Basic/SourceManager.h>
+#include <llvm/Support/FileSystem.h>
+
+namespace codeql {
+
+namespace detail {
+
+// The following `getKindName`s are used within "TBD" TRAP entries to visually mark an AST node as
+// not properly emitted yet.
+// TODO: To be replaced with QL counterpart
+template <typename Parent, typename Kind>
+inline std::string getKindName(Kind kind) {
+  return Parent::getKindName(kind).str();
+}
+
+template <>
+inline std::string getKindName<swift::TypeBase, swift::TypeKind>(swift::TypeKind kind) {
+  switch (kind) {
+#define TYPE(CLASS, PARENT)    \
+  case swift::TypeKind::CLASS: \
+    return #CLASS;
+#include "swift/AST/TypeNodes.def"
+    default:
+      return "Unknown";
+  }
+}
+
+template <>
+std::string inline getKindName<swift::TypeRepr, swift::TypeReprKind>(swift::TypeReprKind kind) {
+  switch (kind) {
+#define TYPEREPR(CLASS, PARENT)    \
+  case swift::TypeReprKind::CLASS: \
+    return #CLASS;
+#include "swift/AST/TypeReprNodes.def"
+    default:
+      return "Unknown";
+  }
+}
+
+}  // namespace detail
+
+// The main reponsibilities of the SwiftDispatcher are as follows:
+// * redirect specific AST node emission to a corresponding visitor (statements, expressions, etc.)
+// * storing TRAP labels for emitted AST nodes (in the TrapLabelStore) to avoid re-emission
+// Since SwiftDispatcher sees all the AST nodes, it also attaches a location to every 'locatable'
+// node (AST nodes that are not types: declarations, statements, expressions, etc.).
+class SwiftDispatcher {
+ public:
+  // sourceManager, arena, and trap are supposed to outlive the SwiftDispatcher
+  SwiftDispatcher(const swift::SourceManager& sourceManager, TrapArena& arena, TrapOutput& trap)
+      : sourceManager{sourceManager}, arena{arena}, trap{trap} {}
+
+  template <typename T>
+  void extract(T* entity) {
+    fetchLabel(entity);
+  }
+
+ private:
+  // This method gives a TRAP label for already emitted AST node.
+  // If the AST node was not emitted yet, then the emission is dispatched to a corresponding
+  // visitor (see `visit(T *)` methods below).
+  template <typename E>
+  TrapLabel<ToTag<E>> fetchLabel(E* e) {
+    // this is required so we avoid any recursive loop: a `fetchLabel` during the visit of `e` might
+    // end up calling `fetchLabel` on `e` itself, so we want the visit of `e` to call `fetchLabel`
+    // only after having called `assignNewLabel` on `e`
+    assert(!waitingForNewLabel && "fetchLabel called before assignNewLabel");
+    if (auto l = store.get(e)) {
+      return *l;
+    }
+    waitingForNewLabel = getCanonicalPointer(e);
+    visit(e);
+    if (auto l = store.get(e)) {
+      if constexpr (!std::is_base_of_v<swift::TypeBase, E>) {
+        attachLocation(e, *l);
+      }
+      return *l;
+    }
+    assert(!"assignNewLabel not called during visit");
+    return {};
+  }
+
+  // Due to the lazy emission approach, we must assign a label to a corresponding AST node before
+  // it actually gets emitted to handle recursive cases such as recursive calls, or recursive type
+  // declarations
+  template <typename E>
+  TrapLabel<ToTag<E>> assignNewLabel(E* e) {
+    assert(waitingForNewLabel == getCanonicalPointer(e) && "assignNewLabel called on wrong entity");
+    auto label = getLabel<ToTag<E>>();
+    trap.assignStar(label);
+    store.insert(e, label);
+    waitingForNewLabel = nullptr;
+    return label;
+  }
+
+  template <typename Tag>
+  TrapLabel<Tag> getLabel() {
+    return arena.allocateLabel<Tag>();
+  }
+
+  // This is a helper method to emit TRAP entries for AST nodes that we don't fully support yet.
+  template <typename Parent, typename Child>
+  void TBD(Child* entity, const std::string& suffix) {
+    using namespace std::string_literals;
+    auto label = assignNewLabel(entity);
+    auto kind = detail::getKindName<Parent>(static_cast<const Parent*>(entity)->getKind());
+    auto name = "TBD ("s + kind + suffix + ")";
+    if constexpr (std::is_same_v<Parent, swift::TypeBase>) {
+      trap.emit(UnknownTypesTrap{label, name});
+    } else {
+      trap.emit(UnknownAstNodesTrap{label, name});
+    }
+  }
+
+  template <typename Locatable>
+  void attachLocation(Locatable locatable, TrapLabel<LocatableTag> locatableLabel) {
+    attachLocation(&locatable, locatableLabel);
+  }
+
+  // Emits a Location TRAP entry and attaches it to an AST node
+  template <typename Locatable>
+  void attachLocation(Locatable* locatable, TrapLabel<LocatableTag> locatableLabel) {
+    auto start = locatable->getStartLoc();
+    auto end = locatable->getEndLoc();
+    if (!start.isValid() || !end.isValid()) {
+      // invalid locations seem to come from entities synthesized by the compiler
+      return;
+    }
+    std::string filepath = getFilepath(start);
+    auto fileLabel = arena.allocateLabel<FileTag>();
+    trap.assignKey(fileLabel, filepath);
+    // TODO: do not emit duplicate trap entries for Files
+    trap.emit(FilesTrap{fileLabel, filepath});
+    auto [startLine, startColumn] = sourceManager.getLineAndColumnInBuffer(start);
+    auto [endLine, endColumn] = sourceManager.getLineAndColumnInBuffer(end);
+    auto locLabel = arena.allocateLabel<LocationTag>();
+    trap.assignKey(locLabel, '{', fileLabel, "}:", startLine, ':', startColumn, ':', endLine, ':',
+                   endColumn);
+    trap.emit(LocationsTrap{locLabel, fileLabel, startLine, startColumn, endLine, endColumn});
+    trap.emit(LocatablesTrap{locatableLabel, locLabel});
+  }
+
+  std::string getFilepath(swift::SourceLoc loc) {
+    // TODO: this needs more testing
+    // TODO: check canonicaliztion of names on a case insensitive filesystems
+    // TODO: make symlink resolution conditional on CODEQL_PRESERVE_SYMLINKS=true
+    auto displayName = sourceManager.getDisplayNameForLoc(loc);
+    llvm::SmallString<PATH_MAX> realPath;
+    if (std::error_code ec = llvm::sys::fs::real_path(displayName, realPath)) {
+      std::cerr << "Cannot get real path: '" << displayName.str() << "': " << ec.message() << "\n";
+      return {};
+    }
+    return realPath.str().str();
+  }
+
+  // TODO: The following methods are supposed to redirect TRAP emission to correpsonding visitors,
+  // which are to be introduced in follow-up PRs
+  void visit(swift::Decl* decl) { TBD<swift::Decl>(decl, "Decl"); }
+  void visit(swift::Stmt* stmt) { TBD<swift::Stmt>(stmt, "Stmt"); }
+  void visit(swift::Expr* expr) { TBD<swift::Expr>(expr, "Expr"); }
+  void visit(swift::Pattern* pattern) { TBD<swift::Pattern>(pattern, "Pattern"); }
+  void visit(swift::TypeRepr* type) { TBD<swift::TypeRepr>(type, "TypeRepr"); }
+  void visit(swift::TypeBase* type) { TBD<swift::TypeBase>(type, "Type"); }
+
+  const swift::SourceManager& sourceManager;
+  TrapArena& arena;
+  TrapOutput& trap;
+  TrapLabelStore store;
+  const void* waitingForNewLabel{nullptr};
+};
+
+}  // namespace codeql

swift/extractor/SwiftExtractor.cpp

@@ -12,13 +12,15 @@
 #include <llvm/Support/FileSystem.h>
 #include <llvm/Support/Path.h>
 
-#include "swift/extractor/trap/TrapClasses.h"
-#include "swift/extractor/trap/TrapArena.h"
+#include "swift/extractor/trap/generated/TrapClasses.h"
 #include "swift/extractor/trap/TrapOutput.h"
+#include "swift/extractor/SwiftDispatcher.h"
 
 using namespace codeql;
 
-static void extractFile(const SwiftExtractorConfiguration& config, swift::SourceFile& file) {
+static void extractFile(const SwiftExtractorConfiguration& config,
+                        swift::CompilerInstance& compiler,
+                        swift::SourceFile& file) {
   if (std::error_code ec = llvm::sys::fs::create_directories(config.trapDir)) {
     std::cerr << "Cannot create TRAP directory: " << ec.message() << "\n";
     return;
@@ -79,12 +81,17 @@ static void extractFile(const SwiftExtractorConfiguration& config, swift::Source
 
   TrapOutput trap{trapStream};
   TrapArena arena{};
-  auto label = arena.allocateLabel<FileTag>();
-  trap.assignStar(label);
-  File f{};
-  f.id = label;
-  f.name = srcFilePath.str().str();
-  trap.emit(f);
+
+  // In the case of emtpy files, the dispatcher is not called, but we still want to 'record' the
+  // fact that the file was extracted
+  auto fileLabel = arena.allocateLabel<FileTag>();
+  trap.assignKey(fileLabel, srcFilePath.str().str());
+  trap.emit(FilesTrap{fileLabel, srcFilePath.str().str()});
+
+  SwiftDispatcher dispatcher(compiler.getSourceMgr(), arena, trap);
+  for (swift::Decl* decl : file.getTopLevelDecls()) {
+    dispatcher.extract(decl);
+  }
 
   // TODO: Pick a better name to avoid collisions
   std::string trapName = file.getFilename().str() + ".trap";
@@ -108,11 +115,11 @@ void codeql::extractSwiftFiles(const SwiftExtractorConfiguration& config,
         module->getFiles().front()->getKind() == swift::FileUnitKind::Source) {
       // We can only call getMainSourceFile if the first file is of a Source kind
       swift::SourceFile& file = module->getMainSourceFile();
-      extractFile(config, file);
+      extractFile(config, compiler, file);
     }
   } else {
     for (auto s : compiler.getPrimarySourceFiles()) {
-      extractFile(config, *s);
+      extractFile(config, compiler, *s);
     }
   }
 }

swift/extractor/SwiftTagTraits.h

@@ -0,0 +1,71 @@
+#pragma once
+
+// This file implements the mapping needed by the API defined in the TrapTagTraits.h
+#include <swift/AST/ASTVisitor.h>
+#include "swift/extractor/trap/TrapTagTraits.h"
+#include "swift/extractor/trap/generated/TrapTags.h"
+
+namespace codeql {
+
+// codegen goes with QL acronym convention (Sil instead of SIL), we need to remap it to Swift's
+// convention
+using SILBlockStorageTypeTag = SilBlockStorageTypeTag;
+using SILBoxTypeTag = SilBoxTypeTag;
+using SILFunctionTypeTag = SilFunctionTypeTag;
+using SILTokenTypeTag = SilTokenTypeTag;
+
+#define MAP_TYPE_TO_TAG(TYPE, TAG)           \
+  template <>                                \
+  struct detail::ToTagFunctor<swift::TYPE> { \
+    using type = TAG;                        \
+  }
+#define MAP_TAG(TYPE) MAP_TYPE_TO_TAG(TYPE, TYPE##Tag)
+#define MAP_SUBTAG(TYPE, PARENT)                           \
+  MAP_TAG(TYPE);                                           \
+  static_assert(std::is_base_of_v<PARENT##Tag, TYPE##Tag>, \
+                #PARENT "Tag must be a base of " #TYPE "Tag");
+
+#define OVERRIDE_TAG(TYPE, TAG)               \
+  template <>                                 \
+  struct detail::ToTagOverride<swift::TYPE> { \
+    using type = TAG;                         \
+  };                                          \
+  static_assert(std::is_base_of_v<TYPE##Tag, TAG>, "override is not a subtag");
+
+MAP_TAG(Stmt);
+#define ABSTRACT_STMT(CLASS, PARENT) MAP_SUBTAG(CLASS##Stmt, PARENT)
+#define STMT(CLASS, PARENT) ABSTRACT_STMT(CLASS, PARENT)
+#include "swift/AST/StmtNodes.def"
+
+MAP_TAG(Expr);
+#define ABSTRACT_EXPR(CLASS, PARENT) MAP_SUBTAG(CLASS##Expr, PARENT)
+#define EXPR(CLASS, PARENT) ABSTRACT_EXPR(CLASS, PARENT)
+#include "swift/AST/ExprNodes.def"
+
+MAP_TAG(Decl);
+#define ABSTRACT_DECL(CLASS, PARENT) MAP_SUBTAG(CLASS##Decl, PARENT)
+#define DECL(CLASS, PARENT) ABSTRACT_DECL(CLASS, PARENT)
+#include "swift/AST/DeclNodes.def"
+
+MAP_TAG(Pattern);
+#define ABSTRACT_PATTERN(CLASS, PARENT) MAP_SUBTAG(CLASS##Pattern, PARENT)
+#define PATTERN(CLASS, PARENT) ABSTRACT_PATTERN(CLASS, PARENT)
+#include "swift/AST/PatternNodes.def"
+
+MAP_TAG(TypeRepr);
+MAP_TYPE_TO_TAG(TypeBase, TypeTag);
+#define ABSTRACT_TYPE(CLASS, PARENT) MAP_SUBTAG(CLASS##Type, PARENT)
+#define TYPE(CLASS, PARENT) ABSTRACT_TYPE(CLASS, PARENT)
+#include "swift/AST/TypeNodes.def"
+
+OVERRIDE_TAG(FuncDecl, ConcreteFuncDeclTag);
+OVERRIDE_TAG(VarDecl, ConcreteVarDeclTag);
+
+#undef MAP_TAG
+#undef MAP_SUBTAG
+#undef MAP_TYPE_TO_TAG
+#undef OVERRIDE_TAG
+
+// All the other macros defined here are undefined by the .def files
+
+}  // namespace codeql