Merge pull request #9112 from AlexDenisov/alexdenisov/introduce-dispatcher

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

swift/extractor/trap/BUILD.bazel

@@ -2,14 +2,14 @@ genrule(
     name = "trapgen",
     srcs = ["//swift:dbscheme"],
     outs = [
-        "TrapEntries.h",
-        "TrapTags.h",
+        "generated/TrapEntries.h",
+        "generated/TrapTags.h",
     ],
     cmd = " ".join([
         "$(location //swift/codegen:trapgen)",
         "--dbscheme $<",
         "--cpp-include-dir " + package_name(),
-        "--cpp-output $(RULEDIR)",
+        "--cpp-output $(RULEDIR)/generated",
     ]),
     exec_tools = ["//swift/codegen:trapgen"],
 )
@@ -21,13 +21,13 @@ genrule(
         "//swift/codegen:schema_includes",
     ],
     outs = [
-        "TrapClasses.h",
+        "generated/TrapClasses.h",
     ],
     cmd = " ".join([
         "$(location //swift/codegen:cppgen)",
         "--schema $(location //swift/codegen:schema)",
         "--cpp-include-dir " + package_name(),
-        "--cpp-output $(RULEDIR)",
+        "--cpp-output $(RULEDIR)/generated",
     ]),
     exec_tools = ["//swift/codegen:cppgen"],
 )

swift/extractor/trap/TrapLabel.h

@@ -5,7 +5,7 @@
 #include <string>
 
 #include "swift/extractor/trap/TrapTagTraits.h"
-#include "swift/extractor/trap/TrapTags.h"
+#include "swift/extractor/trap/generated/TrapTags.h"
 
 namespace codeql {
 
@@ -13,6 +13,8 @@ class UntypedTrapLabel {
   uint64_t id_;
 
   friend class std::hash<UntypedTrapLabel>;
+  template <typename Tag>
+  friend class TrapLabel;
 
  protected:
   UntypedTrapLabel() : id_{0xffffffffffffffff} {}
@@ -41,6 +43,7 @@ class TrapLabel : public UntypedTrapLabel {
 
   // The caller is responsible for ensuring ID uniqueness.
   static TrapLabel unsafeCreateFromExplicitId(uint64_t id) { return {id}; }
+  static TrapLabel unsafeCreateFromUntyped(UntypedTrapLabel label) { return {label.id_}; }
 
   template <typename OtherTag>
   TrapLabel(const TrapLabel<OtherTag>& other) : UntypedTrapLabel(other) {

swift/extractor/trap/TrapLabelStore.h

@@ -0,0 +1,61 @@
+#pragma once
+
+#include <cassert>
+#include <optional>
+#include <unordered_map>
+
+#include <swift/AST/ASTVisitor.h>
+
+#include "swift/extractor/trap/TrapLabel.h"
+#include "swift/extractor/trap/TrapTagTraits.h"
+#include "swift/extractor/trap/generated/TrapTags.h"
+
+namespace codeql {
+
+// The following is needed to avoid the problem of subclass pointers not necessarily coinciding
+// with superclass ones in case of multiple inheritance
+// The interesting part here is implicit conversion from a derived class pointer to the parameter
+inline const swift::Decl* getCanonicalPointer(const swift::Decl* e) {
+  return e;
+}
+inline const swift::Stmt* getCanonicalPointer(const swift::Stmt* e) {
+  return e;
+}
+inline const swift::Expr* getCanonicalPointer(const swift::Expr* e) {
+  return e;
+}
+inline const swift::Pattern* getCanonicalPointer(const swift::Pattern* e) {
+  return e;
+}
+inline const swift::TypeRepr* getCanonicalPointer(const swift::TypeRepr* e) {
+  return e;
+}
+inline const swift::TypeBase* getCanonicalPointer(const swift::TypeBase* e) {
+  return e;
+}
+
+// The extraction is done in a lazy/on-demand fashion:
+// Each emitted TRAP entry for an AST node gets a TRAP label assigned to it.
+// To avoid re-emission, we store the "AST node <> label" entry in the TrapLabelStore.
+class TrapLabelStore {
+ public:
+  template <typename T>
+  std::optional<TrapLabel<ToTag<T>>> get(const T* e) {
+    if (auto found = store_.find(getCanonicalPointer(e)); found != store_.end()) {
+      return TrapLabel<ToTag<T>>::unsafeCreateFromUntyped(found->second);
+    }
+    return std::nullopt;
+  }
+
+  template <typename T>
+  void insert(const T* e, TrapLabel<ToTag<T>> l) {
+    auto [_, inserted] = store_.emplace(getCanonicalPointer(e), l);
+    assert(inserted && "already inserted");
+  }
+
+ private:
+  // TODO: consider std::variant or llvm::PointerUnion instead of `void *`
+  std::unordered_map<const void*, UntypedTrapLabel> store_;
+};
+
+}  // namespace codeql

swift/extractor/trap/TrapTagTraits.h

@@ -1,21 +1,21 @@
 #pragma once
 
+// This file defines functors that can be specialized to define a mapping from arbitrary types to
+// label tags
+
 #include <type_traits>
 
 namespace codeql {
 
+namespace detail {
 template <typename T>
 struct ToTagFunctor;
 template <typename T>
 struct ToTagOverride : ToTagFunctor<T> {};
 
-template <typename T>
-using ToTag = typename ToTagOverride<std::remove_const_t<T>>::type;
+}  // namespace detail
 
 template <typename T>
-struct TagToBindingTrapFunctor;
-
-template <typename Tag>
-using TagToBindingTrap = typename TagToBindingTrapFunctor<Tag>::type;
+using ToTag = typename detail::ToTagOverride<std::remove_const_t<T>>::type;
 
 }  // namespace codeql

swift/ql/lib/codeql/swift/elements/Location.qll

@@ -1,4 +1,11 @@
-// generated by codegen/codegen.py, remove this comment if you wish to edit this file
 private import codeql.swift.generated.Location
 
-class Location extends LocationBase { }
+class Location extends LocationBase {
+  predicate hasLocationInfo(string path, int sl, int sc, int el, int ec) {
+    path = getFile().getFullName() and
+    sl = getStartLine() and
+    sc = getStartColumn() and
+    el = getEndLine() and
+    ec = getEndColumn()
+  }
+}

swift/ql/lib/codeql/swift/elements/UnknownAstNode.qll

@@ -1,4 +1,5 @@
-// generated by codegen/codegen.py, remove this comment if you wish to edit this file
 private import codeql.swift.generated.UnknownAstNode
 
-class UnknownAstNode extends UnknownAstNodeBase { }
+class UnknownAstNode extends UnknownAstNodeBase {
+  override string toString() { result = getName() }
+}

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

@@ -0,0 +1,2 @@
+| test.swift:1:1:1:9 | TBD (TopLevelCodeDecl) |
+| test.swift:1:5:1:5 | TBD (VarDecl) |

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

@@ -0,0 +1,4 @@
+import swift
+
+from Decl decl
+select decl

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

@@ -0,0 +1,2 @@
+let x = 42
+