Swift: replace internal swift mangler with our own

Our mangler is split in two version:
* `SwiftTrapMangler`, with the same behaviour as the previous
  `SwiftMangler`, constructing mangled names with trap label references
* `SwiftRecursiveMangler` that replaces trap label references with
  recursive calls to its own `mangle` functions, effectively rolling out
  the entire chain of references

The latter is used to create lazy trap file names. Hashing is used to
avoid excessively long filenames.

Author: redsun82

swift/extractor/SwiftExtractor.cpp

@@ -14,7 +14,6 @@
 #include "swift/extractor/infra/file/Path.h"
 #include "swift/extractor/infra/SwiftLocationExtractor.h"
 #include "swift/extractor/infra/SwiftBodyEmissionStrategy.h"
-#include "swift/extractor/mangler/SwiftMangler.h"
 #include "swift/logging/SwiftAssert.h"
 
 using namespace codeql;
@@ -49,32 +48,20 @@ static void archiveFile(const SwiftExtractorConfiguration& config, swift::Source
   }
 }
 
-// TODO: This will be factored out/replaced with a 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);
-  }
-  if (decl.getKind() == swift::DeclKind::GenericTypeParam) {
-    // internal mangling does not distinguish generic type parameters with the same name and
-    // position of different functions. We prepend the context (that is, the function) to
-    // circumvent that
-    auto context = llvm::dyn_cast<swift::ValueDecl>(decl.getDeclContext()->getAsDecl());
-    assert(context);
-    return mangledDeclName(*context) + '_' + mangler.mangleAnyDecl(&decl, /* prefix = */ false);
+static std::string mangledDeclName(const swift::Decl& decl) {
+  // SwiftRecursiveMangler mangled name cannot be used directly as it can be too long for file names
+  // so we build some minimal human readable info for debuggability and append a hash
+  auto ret = decl.getModuleContext()->getRealName().str().str();
+  ret += '/';
+  ret += swift::Decl::getKindName(decl.getKind());
+  ret += '_';
+  if (auto valueDecl = llvm::dyn_cast<swift::ValueDecl>(&decl); valueDecl && valueDecl->hasName()) {
+    ret += valueDecl->getBaseName().userFacingName();
+    ret += '_';
   }
-  return mangler.mangleAnyDecl(&decl, /* prefix = */ false);
+  SwiftRecursiveMangler mangler;
+  ret += mangler.mangleDecl(decl).hash();
+  return ret;
 }
 
 static fs::path getFilename(swift::ModuleDecl& module,
@@ -84,20 +71,7 @@ static fs::path getFilename(swift::ModuleDecl& module,
     return resolvePath(primaryFile->getFilename());
   }
   if (lazyDeclaration) {
-    // this code will be thrown away in the near future
-    auto decl = llvm::dyn_cast<swift::ValueDecl>(lazyDeclaration);
-    CODEQL_ASSERT(decl, "not a ValueDecl");
-    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 += '_';
-    ret += decl->getBaseName().userFacingName();
-    ret += '_';
-    // half a SHA2 is enough
-    ret += std::string_view(mangled).substr(0, mangled.size() / 2);
-    return ret;
+    return mangledDeclName(*lazyDeclaration);
   }
   // PCM clang module
   if (module.isNonSwiftModule()) {

swift/extractor/infra/BUILD.bazel

@@ -11,5 +11,6 @@ swift_cc_library(
         "//swift/extractor/trap",
         "//swift/logging",
         "//swift/third_party/swift-llvm-support",
+        "@picosha2",
     ],
 )

swift/extractor/infra/SwiftMangledName.cpp

@@ -1,8 +1,17 @@
 #include "swift/extractor/infra/SwiftMangledName.h"
 #include "absl/strings/str_cat.h"
 
+#include <picosha2.h>
+
 namespace codeql {
 
+std::string SwiftMangledName::hash() const {
+  auto ret = picosha2::hash256_hex_string(value);
+  // half a hash is already enough for disambuiguation
+  ret.resize(ret.size() / 2);
+  return ret;
+}
+
 SwiftMangledName& SwiftMangledName::operator<<(UntypedTrapLabel label) & {
   absl::StrAppend(&value, "{", label.str(), "}");
   return *this;

swift/extractor/infra/SwiftMangledName.h

@@ -16,6 +16,9 @@ class SwiftMangledName {
   explicit operator bool() const { return !value.empty(); }
 
   const std::string& str() const { return value; }
+  operator std::string_view() const { return value; }
+
+  std::string hash() const;
 
   // let's avoid copying as long as we don't need it
   SwiftMangledName() = default;
@@ -29,7 +32,6 @@ class SwiftMangledName {
     (operator<<(std::forward<Args>(args)), ...);
   }
 
-  // streaming labels or ints into a SwiftMangledName just appends them
   SwiftMangledName& operator<<(UntypedTrapLabel label) &;
   SwiftMangledName& operator<<(unsigned i) &;
 
@@ -46,11 +48,6 @@ class SwiftMangledName {
     return *this;
   }
 
-  SwiftMangledName& operator<<(const SwiftMangledName& other) {
-    value += other.value;
-    return *this;
-  }
-
   template <typename T>
   SwiftMangledName&& operator<<(T&& arg) && {
     return std::move(operator<<(std::forward<T>(arg)));

swift/extractor/mangler/SwiftMangler.cpp

@@ -38,14 +38,6 @@ std::string_view getTypeKindStr(const swift::TypeBase* type) {
 
 }  // namespace
 
-template <typename E>
-UntypedTrapLabel SwiftMangler::fetch(E&& e) {
-  auto ret = dispatcher.fetchLabel(std::forward<E>(e));
-  // TODO use a generic logging handle for Swift entities here, once it's available
-  CODEQL_ASSERT(ret.valid(), "using an undefined label in mangling");
-  return ret;
-}
-
 SwiftMangledName SwiftMangler::initMangled(const swift::TypeBase* type) {
   return {getTypeKindStr(type), '_'};
 }
@@ -373,3 +365,29 @@ SwiftMangledName SwiftMangler::visitParametrizedProtocolType(
   ret << '>';
   return ret;
 }
+
+namespace {
+template <typename E>
+UntypedTrapLabel fetchLabel(SwiftDispatcher& dispatcher, const E* e) {
+  auto ret = dispatcher.fetchLabel(e);
+  // TODO use a generic logging handle for Swift entities here, once it's available
+  CODEQL_ASSERT(ret.valid(), "using an undefined label in mangling");
+  return ret;
+}
+}  // namespace
+
+SwiftMangledName SwiftTrapMangler::fetch(const swift::Decl* decl) {
+  return {fetchLabel(dispatcher, decl)};
+}
+
+SwiftMangledName SwiftTrapMangler::fetch(const swift::TypeBase* type) {
+  return {fetchLabel(dispatcher, type)};
+}
+
+SwiftMangledName SwiftRecursiveMangler::fetch(const swift::Decl* decl) {
+  return mangleDecl(*decl).hash();
+}
+
+SwiftMangledName SwiftRecursiveMangler::fetch(const swift::TypeBase* type) {
+  return mangleType(*type).hash();
+}

swift/extractor/mangler/SwiftMangler.h

@@ -36,7 +36,7 @@ class SwiftMangler : private swift::TypeVisitor<SwiftMangler, SwiftMangledName>,
   using DeclVisitor = swift::DeclVisitor<SwiftMangler, SwiftMangledName>;
 
  public:
-  explicit SwiftMangler(SwiftDispatcher& dispatcher) : dispatcher(dispatcher) {}
+  virtual ~SwiftMangler() = default;
 
   static SwiftMangledName mangleModuleName(std::string_view name);
 
@@ -104,18 +104,33 @@ class SwiftMangler : private swift::TypeVisitor<SwiftMangler, SwiftMangledName>,
   SwiftMangledName visitParametrizedProtocolType(const swift::ParameterizedProtocolType* type);
 
  private:
-  template <typename E>
-  UntypedTrapLabel fetch(E&& e);
+  std::unordered_map<const swift::Decl*, unsigned> preloadedExtensionIndexes;
+
+  virtual SwiftMangledName fetch(const swift::Decl* decl) = 0;
+  virtual SwiftMangledName fetch(const swift::TypeBase* type) = 0;
+  SwiftMangledName fetch(swift::Type type) { return fetch(type.getPointer()); }
 
+  void indexExtensions(llvm::ArrayRef<swift::Decl*> siblings);
+  unsigned int getExtensionIndex(const swift::ExtensionDecl* decl, const swift::Decl* parent);
   static SwiftMangledName initMangled(const swift::TypeBase* type);
   SwiftMangledName initMangled(const swift::Decl* decl);
   SwiftMangledName visitTypeDiscriminatedValueDecl(const swift::ValueDecl* decl);
+};
+
+class SwiftTrapMangler : public SwiftMangler {
+ public:
+  explicit SwiftTrapMangler(SwiftDispatcher& dispatcher) : dispatcher(dispatcher) {}
+
+ private:
+  SwiftMangledName fetch(const swift::Decl* decl) override;
+  SwiftMangledName fetch(const swift::TypeBase* type) override;
 
-  swift::Mangle::ASTMangler mangler;
   SwiftDispatcher& dispatcher;
-  std::unordered_map<const swift::Decl*, unsigned> preloadedExtensionIndexes;
-  void indexExtensions(llvm::ArrayRef<swift::Decl*> siblings);
-  unsigned int getExtensionIndex(const swift::ExtensionDecl* decl, const swift::Decl* parent);
+};
+
+class SwiftRecursiveMangler : public SwiftMangler {
+  SwiftMangledName fetch(const swift::Decl* decl) override;
+  SwiftMangledName fetch(const swift::TypeBase* type) override;
 };
 
 }  // namespace codeql

swift/extractor/translators/SwiftVisitor.h

@@ -58,7 +58,7 @@ class SwiftVisitor : private SwiftDispatcher {
   StmtTranslator stmtTranslator{*this};
   TypeTranslator typeTranslator{*this};
   PatternTranslator patternTranslator{*this};
-  SwiftMangler mangler{*this};
+  SwiftTrapMangler mangler{*this};
 };
 
 }  // namespace codeql