Swift: cache file paths

This required a bit of a generalization of `TrapLabelStore` to not
work only with pointers.

Author: redsun82

swift/extractor/SwiftExtractor.cpp

@@ -104,29 +104,19 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
   dumpArgs(*trapTarget, config);
   TrapDomain trap{*trapTarget};
 
-  // TODO: move default location emission elsewhere, possibly in a separate global trap file
+  // TODO: remove this and recreate it with IPA when we have that
   // the following cannot conflict with actual files as those have an absolute path starting with /
-  auto unknownFileLabel = trap.createLabel<FileTag>("unknown");
-  auto unknownLocationLabel = trap.createLabel<LocationTag>("unknown");
-  trap.emit(FilesTrap{unknownFileLabel});
-  trap.emit(LocationsTrap{unknownLocationLabel, unknownFileLabel});
+  File unknownFileEntry{trap.createLabel<FileTag>("unknown")};
+  Location unknownLocationEntry{trap.createLabel<LocationTag>("unknown")};
+  unknownLocationEntry.file = unknownFileEntry.id;
+  trap.emit(unknownFileEntry);
+  trap.emit(unknownLocationEntry);
 
   SwiftVisitor visitor(compiler.getSourceMgr(), trap, module, primaryFile);
   auto topLevelDecls = getTopLevelDecls(module, primaryFile);
   for (auto decl : topLevelDecls) {
     visitor.extract(decl);
   }
-  // TODO the following will be moved to the dispatcher when we start caching swift file objects
-  // for the moment, topLevelDecls always contains the current module, which does not have a file
-  // associated with it, so we need a special case when there are no top level declarations
-  if (topLevelDecls.size() == 1) {
-    // In the case of empty files, the dispatcher is not called, but we still want to 'record' the
-    // fact that the file was extracted
-    llvm::SmallString<PATH_MAX> name(filename);
-    llvm::sys::fs::make_absolute(name);
-    auto fileLabel = trap.createLabel<FileTag>(name.str().str());
-    trap.emit(FilesTrap{fileLabel, name.str().str()});
-  }
 }
 
 static std::unordered_set<std::string> collectInputFilenames(swift::CompilerInstance& compiler) {

swift/extractor/infra/FilePath.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <string>
+namespace codeql {
+
+// wrapper around `std::string` mainly intended to unambiguously go into an `std::variant`
+// TODO probably not needed once we can use `std::filesystem::path`
+struct FilePath {
+  FilePath() = default;
+  FilePath(const std::string& path) : path{path} {}
+  FilePath(std::string&& path) : path{std::move(path)} {}
+
+  std::string path;
+
+  bool operator==(const FilePath& other) const { return path == other.path; }
+};
+}  // namespace codeql
+
+namespace std {
+template <>
+struct hash<codeql::FilePath> {
+  size_t operator()(const codeql::FilePath& value) { return hash<string>{}(value.path); }
+};
+}  // namespace std

swift/extractor/infra/SwiftDispatcher.h

@@ -8,6 +8,7 @@
 #include "swift/extractor/trap/TrapDomain.h"
 #include "swift/extractor/infra/SwiftTagTraits.h"
 #include "swift/extractor/trap/generated/TrapClasses.h"
+#include "swift/extractor/infra/FilePath.h"
 
 namespace codeql {
 
@@ -17,6 +18,24 @@ namespace codeql {
 // 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 {
+  // types to be supported by assignNewLabel/fetchLabel need to be listed here
+  using Store = TrapLabelStore<const swift::Decl*,
+                               const swift::Stmt*,
+                               const swift::StmtCondition*,
+                               const swift::StmtConditionElement*,
+                               const swift::CaseLabelItem*,
+                               const swift::Expr*,
+                               const swift::Pattern*,
+                               const swift::TypeRepr*,
+                               const swift::TypeBase*,
+                               FilePath>;
+
+  template <typename E>
+  static constexpr bool IsStorable = std::is_constructible_v<Store::Handle, const E&>;
+
+  template <typename E>
+  static constexpr bool IsLocatable = std::is_base_of_v<LocatableTag, TrapTagOf<E>>;
+
  public:
   // all references and pointers passed as parameters to this constructor are supposed to outlive
   // the SwiftDispatcher
@@ -27,7 +46,12 @@ class SwiftDispatcher {
       : sourceManager{sourceManager},
         trap{trap},
         currentModule{currentModule},
-        currentPrimarySourceFile{currentPrimarySourceFile} {}
+        currentPrimarySourceFile{currentPrimarySourceFile} {
+    if (currentPrimarySourceFile) {
+      // we make sure the file is in the trap output even if the source is empty
+      fetchLabel(getFilePath(currentPrimarySourceFile->getFilename()));
+    }
+  }
 
   template <typename Entry>
   void emit(const Entry& entry) {
@@ -61,9 +85,11 @@ class SwiftDispatcher {
   // 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>
-  TrapLabelOf<E> fetchLabel(E* e) {
-    assert(e && "trying to fetch a label on nullptr, maybe fetchOptionalLabel is to be used?");
+  template <typename E, std::enable_if_t<IsStorable<E>>* = nullptr>
+  TrapLabelOf<E> fetchLabel(const E& e) {
+    if constexpr (std::is_constructible_v<bool, const E&>) {
+      assert(e && "fetching a label on a null entity, maybe fetchOptionalLabel is to be used?");
+    }
     // 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`.
@@ -76,7 +102,7 @@ class SwiftDispatcher {
     visit(e);
     // TODO when everything is moved to structured C++ classes, this should be moved to createEntry
     if (auto l = store.get(e)) {
-      if constexpr (!std::is_base_of_v<swift::TypeBase, E>) {
+      if constexpr (IsLocatable<E>) {
         attachLocation(e, *l);
       }
       return *l;
@@ -93,36 +119,37 @@ class SwiftDispatcher {
     return fetchLabelFromUnion<AstNodeTag>(node);
   }
 
-  TrapLabel<ConditionElementTag> fetchLabel(const swift::StmtConditionElement& element) {
-    return fetchLabel(&element);
+  template <typename E, std::enable_if_t<IsStorable<E*>>* = nullptr>
+  TrapLabelOf<E> fetchLabel(const E& e) {
+    return fetchLabel(&e);
   }
 
   // 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, typename... Args>
-  TrapLabelOf<E> assignNewLabel(E* e, Args&&... args) {
+  template <typename E, typename... Args, std::enable_if_t<IsStorable<E>>* = nullptr>
+  TrapLabelOf<E> assignNewLabel(const E& e, Args&&... args) {
     assert(waitingForNewLabel == Store::Handle{e} && "assignNewLabel called on wrong entity");
     auto label = trap.createLabel<TrapTagOf<E>>(std::forward<Args>(args)...);
     store.insert(e, label);
     waitingForNewLabel = std::monostate{};
     return label;
   }
 
-  template <typename E, typename... Args, std::enable_if_t<!std::is_pointer_v<E>>* = nullptr>
+  template <typename E, typename... Args, std::enable_if_t<IsStorable<E*>>* = nullptr>
   TrapLabelOf<E> assignNewLabel(const E& e, Args&&... args) {
     return assignNewLabel(&e, std::forward<Args>(args)...);
   }
 
   // convenience methods for structured C++ creation
-  template <typename E, typename... Args, std::enable_if_t<!std::is_pointer_v<E>>* = nullptr>
+  template <typename E, typename... Args>
   auto createEntry(const E& e, Args&&... args) {
-    return TrapClassOf<E>{assignNewLabel(&e, std::forward<Args>(args)...)};
+    return TrapClassOf<E>{assignNewLabel(e, std::forward<Args>(args)...)};
   }
 
   // used to create a new entry for entities that should not be cached
   // an example is swift::Argument, that are created on the fly and thus have no stable pointer
-  template <typename E, typename... Args, std::enable_if_t<!std::is_pointer_v<E>>* = nullptr>
+  template <typename E, typename... Args>
   auto createUncachedEntry(const E& e, Args&&... args) {
     auto label = trap.createLabel<TrapTagOf<E>>(std::forward<Args>(args)...);
     attachLocation(&e, label);
@@ -206,34 +233,23 @@ class SwiftDispatcher {
   }
 
  private:
-  // types to be supported by assignNewLabel/fetchLabel need to be listed here
-  using Store = TrapLabelStore<swift::Decl,
-                               swift::Stmt,
-                               swift::StmtCondition,
-                               swift::StmtConditionElement,
-                               swift::CaseLabelItem,
-                               swift::Expr,
-                               swift::Pattern,
-                               swift::TypeRepr,
-                               swift::TypeBase>;
-
   void attachLocation(swift::SourceLoc start,
                       swift::SourceLoc end,
                       TrapLabel<LocatableTag> locatableLabel) {
     if (!start.isValid() || !end.isValid()) {
       // invalid locations seem to come from entities synthesized by the compiler
       return;
     }
-    std::string filepath = getFilepath(start);
-    auto fileLabel = trap.createLabel<FileTag>(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 = trap.createLabel<LocationTag>('{', fileLabel, "}:", startLine, ':', startColumn,
-                                                  ':', endLine, ':', endColumn);
-    trap.emit(LocationsTrap{locLabel, fileLabel, startLine, startColumn, endLine, endColumn});
-    trap.emit(LocatableLocationsTrap{locatableLabel, locLabel});
+    auto file = getFilePath(sourceManager.getDisplayNameForLoc(start));
+    Location entry{{}};
+    entry.file = fetchLabel(file);
+    std::tie(entry.start_line, entry.start_column) = sourceManager.getLineAndColumnInBuffer(start);
+    std::tie(entry.end_line, entry.end_column) = sourceManager.getLineAndColumnInBuffer(end);
+    entry.id = trap.createLabel<LocationTag>('{', entry.file, "}:", entry.start_line, ':',
+                                             entry.start_column, ':', entry.end_line, ':',
+                                             entry.end_column);
+    emit(entry);
+    emit(LocatableLocationsTrap{locatableLabel, entry.id});
   }
 
   template <typename Tag, typename... Ts>
@@ -256,21 +272,18 @@ class SwiftDispatcher {
     return false;
   }
 
-  std::string getFilepath(swift::SourceLoc loc) {
+  static FilePath getFilePath(llvm::StringRef path) {
     // 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";
+    if (std::error_code ec = llvm::sys::fs::real_path(path, realPath)) {
+      std::cerr << "Cannot get real path: '" << path.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
   virtual void visit(swift::Decl* decl) = 0;
   virtual void visit(swift::Stmt* stmt) = 0;
   virtual void visit(const swift::StmtCondition* cond) = 0;
@@ -281,6 +294,12 @@ class SwiftDispatcher {
   virtual void visit(swift::TypeRepr* type) = 0;
   virtual void visit(swift::TypeBase* type) = 0;
 
+  void visit(const FilePath& file) {
+    auto entry = createEntry(file);
+    entry.name = file.path;
+    emit(entry);
+  }
+
   const swift::SourceManager& sourceManager;
   TrapDomain& trap;
   Store store;

swift/extractor/infra/SwiftTagTraits.h

@@ -5,6 +5,7 @@
 #include <swift/AST/ASTVisitor.h>
 #include "swift/extractor/trap/TrapTagTraits.h"
 #include "swift/extractor/trap/generated/TrapTags.h"
+#include "swift/extractor/infra/FilePath.h"
 
 namespace codeql {
 
@@ -16,12 +17,12 @@ using SILFunctionTypeTag = SilFunctionTypeTag;
 using SILTokenTypeTag = SilTokenTypeTag;
 using SILBoxTypeReprTag = SilBoxTypeReprTag;
 
-#define MAP_TYPE_TO_TAG(TYPE, TAG)           \
-  template <>                                \
-  struct detail::ToTagFunctor<swift::TYPE> { \
-    using type = TAG;                        \
+#define MAP_TYPE_TO_TAG(TYPE, TAG)    \
+  template <>                         \
+  struct detail::ToTagFunctor<TYPE> { \
+    using type = TAG;                 \
   }
-#define MAP_TAG(TYPE) MAP_TYPE_TO_TAG(TYPE, TYPE##Tag)
+#define MAP_TAG(TYPE) MAP_TYPE_TO_TAG(swift::TYPE, TYPE##Tag)
 #define MAP_SUBTAG(TYPE, PARENT)                           \
   MAP_TAG(TYPE);                                           \
   static_assert(std::is_base_of_v<PARENT##Tag, TYPE##Tag>, \
@@ -36,7 +37,7 @@ using SILBoxTypeReprTag = SilBoxTypeReprTag;
 
 MAP_TAG(Stmt);
 MAP_TAG(StmtCondition);
-MAP_TYPE_TO_TAG(StmtConditionElement, ConditionElementTag);
+MAP_TYPE_TO_TAG(swift::StmtConditionElement, ConditionElementTag);
 MAP_TAG(CaseLabelItem);
 #define ABSTRACT_STMT(CLASS, PARENT) MAP_SUBTAG(CLASS##Stmt, PARENT)
 #define STMT(CLASS, PARENT) ABSTRACT_STMT(CLASS, PARENT)
@@ -63,14 +64,16 @@ MAP_TAG(TypeRepr);
 #define TYPEREPR(CLASS, PARENT) ABSTRACT_TYPEREPR(CLASS, PARENT)
 #include <swift/AST/TypeReprNodes.def>
 
-MAP_TYPE_TO_TAG(TypeBase, TypeTag);
+MAP_TYPE_TO_TAG(swift::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);
 
+MAP_TYPE_TO_TAG(FilePath, FileTag);
+
 #undef MAP_TAG
 #undef MAP_SUBTAG
 #undef MAP_TYPE_TO_TAG

swift/extractor/trap/TrapLabelStore.h

@@ -20,18 +20,18 @@ namespace codeql {
 template <typename... Ts>
 class TrapLabelStore {
  public:
-  using Handle = std::variant<std::monostate, const Ts*...>;
+  using Handle = std::variant<std::monostate, Ts...>;
 
   template <typename T>
-  std::optional<TrapLabelOf<T>> get(const T* e) {
+  std::optional<TrapLabelOf<T>> get(const T& e) {
     if (auto found = store_.find(e); found != store_.end()) {
       return TrapLabelOf<T>::unsafeCreateFromUntyped(found->second);
     }
     return std::nullopt;
   }
 
   template <typename T>
-  void insert(const T* e, TrapLabelOf<T> l) {
+  void insert(const T& e, TrapLabelOf<T> l) {
     auto [_, inserted] = store_.emplace(e, l);
     assert(inserted && "already inserted");
   }

swift/extractor/trap/TrapTagTraits.h

@@ -26,7 +26,8 @@ struct ToTrapClassFunctor;
 }  // namespace detail
 
 template <typename T>
-using TrapTagOf = typename detail::ToTagOverride<std::remove_const_t<T>>::type;
+using TrapTagOf =
+    typename detail::ToTagOverride<std::remove_const_t<std::remove_pointer_t<T>>>::type;
 
 template <typename T>
 using TrapLabelOf = TrapLabel<TrapTagOf<T>>;

swift/extractor/visitors/SwiftVisitor.h

@@ -15,7 +15,7 @@ class SwiftVisitor : private SwiftDispatcher {
   using SwiftDispatcher::SwiftDispatcher;
 
   template <typename T>
-  void extract(T* entity) {
+  void extract(const T& entity) {
     fetchLabel(entity);
   }