Merge pull request #12500 from github/redsun82/swift-dispatcher-rework

Swift: rework fetching and dispatching

Author: redsun82

swift/extractor/BUILD.bazel

@@ -12,7 +12,6 @@ swift_cc_binary(
         "//swift/extractor/config",
         "//swift/extractor/infra",
         "//swift/extractor/invocation",
-        "//swift/extractor/mangler",
         "//swift/extractor/remapping",
         "//swift/extractor/translators",
         "//swift/third_party/swift-llvm-support",

swift/extractor/SwiftExtractor.cpp

@@ -44,6 +44,26 @@ static void archiveFile(const SwiftExtractorConfiguration& config, swift::Source
   }
 }
 
+// TODO: This should be factored out/replaced with simplified version of custom mangling
+static std::string mangledDeclName(const swift::ValueDecl& decl) {
+  std::string_view moduleName = decl.getModuleContext()->getRealName().str();
+  // ASTMangler::mangleAnyDecl crashes when called on `ModuleDecl`
+  if (decl.getKind() == swift::DeclKind::Module) {
+    return std::string{moduleName};
+  }
+  swift::Mangle::ASTMangler mangler;
+  if (decl.getKind() == swift::DeclKind::TypeAlias) {
+    // In cases like this (when coming from PCM)
+    //  typealias CFXMLTree = CFTree
+    //  typealias CFXMLTreeRef = CFXMLTree
+    // mangleAnyDecl mangles both CFXMLTree and CFXMLTreeRef into 'So12CFXMLTreeRefa'
+    // which is not correct and causes inconsistencies. mangleEntity makes these two distinct
+    // prefix adds a couple of special symbols, we don't necessary need them
+    return mangler.mangleEntity(&decl);
+  }
+  return mangler.mangleAnyDecl(&decl, /* prefix = */ false);
+}
+
 static fs::path getFilename(swift::ModuleDecl& module,
                             swift::SourceFile* primaryFile,
                             const swift::Decl* lazyDeclaration) {
@@ -52,15 +72,15 @@ static fs::path getFilename(swift::ModuleDecl& module,
   }
   if (lazyDeclaration) {
     // this code will be thrown away in the near future
-    SwiftMangler mangler;
-    auto mangled = mangler.mangledName(*lazyDeclaration);
+    auto decl = llvm::dyn_cast<swift::ValueDecl>(lazyDeclaration);
+    assert(decl);
+    auto mangled = mangledDeclName(*decl);
     // mangled name can be too long to use as a file name, so we can't use it directly
     mangled = picosha2::hash256_hex_string(mangled);
     std::string ret;
     ret += module.getRealName().str();
     ret += '_';
-    // lazyDeclaration must be a ValueDecl, as already asserted in SwiftMangler::mangledName
-    ret += llvm::cast<swift::ValueDecl>(lazyDeclaration)->getBaseName().userFacingName();
+    ret += decl->getBaseName().userFacingName();
     ret += '_';
     // half a SHA2 is enough
     ret += std::string_view(mangled).substr(0, mangled.size() / 2);

swift/extractor/infra/SwiftDispatcher.h

@@ -6,12 +6,12 @@
 #include <swift/Basic/SourceManager.h>
 #include <swift/Parse/Token.h>
 
-#include "swift/extractor/trap/TrapLabelStore.h"
 #include "swift/extractor/trap/TrapDomain.h"
 #include "swift/extractor/infra/SwiftTagTraits.h"
 #include "swift/extractor/trap/generated/TrapClasses.h"
 #include "swift/extractor/infra/SwiftLocationExtractor.h"
 #include "swift/extractor/infra/SwiftBodyEmissionStrategy.h"
+#include "swift/extractor/infra/SwiftMangledName.h"
 #include "swift/extractor/config/SwiftExtractorState.h"
 #include "swift/extractor/infra/log/SwiftLogging.h"
 
@@ -24,24 +24,31 @@ namespace codeql {
 // 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*,
-                               const swift::CapturedValue*,
-                               const swift::PoundAvailableInfo*,
-                               const swift::AvailabilitySpec*>;
+  using Handle = std::variant<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*,
+                              const swift::CapturedValue*,
+                              const swift::PoundAvailableInfo*,
+                              const swift::AvailabilitySpec*>;
 
   template <typename E>
-  static constexpr bool IsStorable = std::is_constructible_v<Store::Handle, const E&>;
+  static constexpr bool IsFetchable = std::is_constructible_v<Handle, const E&>;
 
   template <typename E>
-  static constexpr bool IsLocatable = std::is_base_of_v<LocatableTag, TrapTagOf<E>>;
+  static constexpr bool IsLocatable =
+      std::is_base_of_v<LocatableTag, TrapTagOf<E>> && !std::is_base_of_v<TypeTag, TrapTagOf<E>>;
+
+  template <typename E>
+  static constexpr bool IsDeclPointer = std::is_convertible_v<E, const swift::Decl*>;
+
+  template <typename E>
+  static constexpr bool IsTypePointer = std::is_convertible_v<E, const swift::TypeBase*>;
 
  public:
   // all references and pointers passed as parameters to this constructor are supposed to outlive
@@ -112,7 +119,7 @@ class SwiftDispatcher {
   TrapLabel<UnspecifiedElementTag> emitUnspecified(std::optional<TrapLabel<ElementTag>>&& parent,
                                                    const char* property,
                                                    int index) {
-    UnspecifiedElement entry{trap.createLabel<UnspecifiedElementTag>()};
+    UnspecifiedElement entry{trap.createTypedLabel<UnspecifiedElementTag>()};
     entry.error = "element was unspecified by the extractor";
     entry.parent = std::move(parent);
     entry.property = property;
@@ -135,35 +142,22 @@ 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, typename... Args, std::enable_if_t<IsStorable<E>>* = nullptr>
-  TrapLabelOf<E> fetchLabel(const E& e, Args&&... args) {
+  template <typename E, std::enable_if_t<IsFetchable<E>>* = nullptr>
+  TrapLabelOf<E> fetchLabel(const E& e, swift::Type type = {}) {
     if constexpr (std::is_constructible_v<bool, const E&>) {
       if (!e) {
         // this will be treated on emission
         return undefined_label;
       }
     }
-    // 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(std::holds_alternative<std::monostate>(waitingForNewLabel) &&
-           "fetchLabel called before assignNewLabel");
-    if (auto l = store.get(e)) {
-      return *l;
+    auto& stored = store[e];
+    if (!stored.valid()) {
+      auto inserted = fetching.insert(e);
+      assert(inserted.second && "detected infinite fetchLabel recursion");
+      stored = createLabel(e, type);
+      fetching.erase(e);
     }
-    waitingForNewLabel = e;
-    // TODO: add tracing logs for visited stuff, maybe within the translators?
-    visit(e, std::forward<Args>(args)...);
-    Log::flush();
-    // TODO when everything is moved to structured C++ classes, this should be moved to createEntry
-    if (auto l = store.get(e)) {
-      if constexpr (IsLocatable<E>) {
-        locationExtractor.attachLocation(sourceManager, e, *l);
-      }
-      return *l;
-    }
-    assert(!"assignNewLabel not called during visit");
-    return {};
+    return TrapLabelOf<E>::unsafeCreateFromUntyped(stored);
   }
 
   // convenience `fetchLabel` overload for `swift::Type` (which is just a wrapper for
@@ -174,39 +168,28 @@ class SwiftDispatcher {
     return fetchLabelFromUnion<AstNodeTag>(node);
   }
 
-  template <typename E, std::enable_if_t<IsStorable<E*>>* = nullptr>
+  template <typename E, std::enable_if_t<IsFetchable<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, std::enable_if_t<IsStorable<E>>* = nullptr>
-  TrapLabel<ConcreteTrapTagOf<E>> assignNewLabel(const E& e, Args&&... args) {
-    assert(waitingForNewLabel == Store::Handle{e} && "assignNewLabel called on wrong entity");
-    auto label = trap.createLabel<ConcreteTrapTagOf<E>>(std::forward<Args>(args)...);
-    store.insert(e, label);
-    waitingForNewLabel = std::monostate{};
-    return label;
-  }
-
-  template <typename E, typename... Args, std::enable_if_t<IsStorable<E*>>* = nullptr>
-  TrapLabel<ConcreteTrapTagOf<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>
-  auto createEntry(const E& e, Args&&... args) {
-    return TrapClassOf<E>{assignNewLabel(e, std::forward<Args>(args)...)};
+  template <typename E>
+  auto createEntry(const E& e) {
+    auto found = store.find(&e);
+    assert(found != store.end() && "createEntry called on non-fetched label");
+    auto label = TrapLabel<ConcreteTrapTagOf<E>>::unsafeCreateFromUntyped(found->second);
+    if constexpr (IsLocatable<E>) {
+      locationExtractor.attachLocation(sourceManager, e, label);
+    }
+    return TrapClassOf<E>{label};
   }
 
   // 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>
-  auto createUncachedEntry(const E& e, Args&&... args) {
-    auto label = trap.createLabel<TrapTagOf<E>>(std::forward<Args>(args)...);
+  template <typename E>
+  auto createUncachedEntry(const E& e) {
+    auto label = trap.createTypedLabel<TrapTagOf<E>>();
     locationExtractor.attachLocation(sourceManager, &e, label);
     return TrapClassOf<E>{label};
   }
@@ -243,31 +226,65 @@ class SwiftDispatcher {
     trap.debug(args...);
   }
 
-  bool shouldEmitDeclBody(const swift::Decl& decl) {
-    encounteredModules.insert(decl.getModuleContext());
-    return bodyEmissionStrategy.shouldEmitDeclBody(decl);
-  }
-
   void emitComment(swift::Token& comment) {
-    CommentsTrap entry{trap.createLabel<CommentTag>(), comment.getRawText().str()};
+    CommentsTrap entry{trap.createTypedLabel<CommentTag>(), comment.getRawText().str()};
     trap.emit(entry);
     locationExtractor.attachLocation(sourceManager, comment, entry.id);
   }
 
-  void extractedDeclaration(const swift::Decl& decl) {
+ protected:
+  void visitPending() {
+    while (!toBeVisited.empty()) {
+      auto [next, type] = toBeVisited.back();
+      toBeVisited.pop_back();
+      // TODO: add tracing logs for visited stuff, maybe within the translators?
+      std::visit([this, type = type](const auto* e) { visit(e, type); }, next);
+    }
+  }
+
+ private:
+  template <typename E>
+  UntypedTrapLabel createLabel(const E& e, swift::Type type) {
+    if constexpr (IsDeclPointer<E> || IsTypePointer<E>) {
+      if (auto mangledName = name(e)) {
+        if (shouldVisit(e)) {
+          toBeVisited.emplace_back(e, type);
+        }
+        return trap.createLabel(mangledName);
+      }
+    }
+    // we always need to visit unnamed things
+    toBeVisited.emplace_back(e, type);
+    return trap.createLabel();
+  }
+
+  template <typename E>
+  bool shouldVisit(const E& e) {
+    if constexpr (IsDeclPointer<E>) {
+      encounteredModules.insert(e->getModuleContext());
+      if (bodyEmissionStrategy.shouldEmitDeclBody(*e)) {
+        extractedDeclaration(e);
+        return true;
+      }
+      skippedDeclaration(e);
+      return false;
+    }
+    return true;
+  }
+
+  void extractedDeclaration(const swift::Decl* decl) {
     if (isLazyDeclaration(decl)) {
-      state.emittedDeclarations.insert(&decl);
+      state.emittedDeclarations.insert(decl);
     }
   }
-  void skippedDeclaration(const swift::Decl& decl) {
+  void skippedDeclaration(const swift::Decl* decl) {
     if (isLazyDeclaration(decl)) {
-      state.pendingDeclarations.insert(&decl);
+      state.pendingDeclarations.insert(decl);
     }
   }
 
- private:
-  bool isLazyDeclaration(const swift::Decl& decl) {
-    swift::ModuleDecl* module = decl.getModuleContext();
+  static bool isLazyDeclaration(const swift::Decl* decl) {
+    swift::ModuleDecl* module = decl->getModuleContext();
     return module->isBuiltinModule() || module->getName().str() == "__ObjC" ||
            module->isNonSwiftModule();
   }
@@ -311,6 +328,8 @@ class SwiftDispatcher {
     return false;
   }
 
+  virtual SwiftMangledName name(const swift::Decl* decl) = 0;
+  virtual SwiftMangledName name(const swift::TypeBase* type) = 0;
   virtual void visit(const swift::Decl* decl) = 0;
   virtual void visit(const swift::Stmt* stmt) = 0;
   virtual void visit(const swift::StmtCondition* cond) = 0;
@@ -324,13 +343,19 @@ class SwiftDispatcher {
   virtual void visit(const swift::TypeBase* type) = 0;
   virtual void visit(const swift::CapturedValue* capture) = 0;
 
+  template <typename T, std::enable_if<!std::is_base_of_v<swift::TypeRepr, T>>* = nullptr>
+  void visit(const T* e, swift::Type) {
+    visit(e);
+  }
+
   const swift::SourceManager& sourceManager;
   SwiftExtractorState& state;
   TrapDomain& trap;
-  Store store;
+  std::unordered_map<Handle, UntypedTrapLabel> store;
+  std::unordered_set<Handle> fetching;
+  std::vector<std::pair<Handle, swift::Type>> toBeVisited;
   SwiftLocationExtractor& locationExtractor;
   SwiftBodyEmissionStrategy& bodyEmissionStrategy;
-  Store::Handle waitingForNewLabel{std::monostate{}};
   std::unordered_set<swift::ModuleDecl*> encounteredModules;
   Logger logger{"dispatcher"};
 };