Swift: Avoid eager indexing of imported extensions

swift/extractor/mangler/SwiftMangler.cpp

@@ -8,7 +8,10 @@
 #include <swift/AST/GenericEnvironment.h>
 #include <swift/AST/GenericParamList.h>
 #include <swift/AST/ClangModuleLoader.h>
+#include <swift/ClangImporter/ClangModule.h>
 #include <clang/Basic/Module.h>
+#include <clang/AST/DeclObjC.h>
+#include <clang/Index/USRGeneration.h>
 
 using namespace codeql;
 
@@ -110,6 +113,28 @@ SwiftMangledName SwiftMangler::visitExtensionDecl(const swift::ExtensionDecl* de
 
   auto parent = getParent(decl);
   auto target = decl->getExtendedType();
+  if (auto category =
+          llvm::dyn_cast_or_null<clang::ObjCCategoryDecl>(decl->getClangNode().getAsDecl())) {
+    // Clang's category identity avoids loading every declaration in the imported module.
+    llvm::SmallString<128> usr;
+    if (!clang::index::generateUSRForDecl(category, usr)) {
+      auto module = category->getOwningModule();
+      return initMangled(decl) << fetch(target) << "|clang|"
+                               << (module ? module->getFullModuleName() : "") << "|"
+                               << usr.str().str();
+    }
+    LOG_WARNING("Unable to generate an imported category USR; using declaration indexes");
+  } else if (!decl->getClangNode() && llvm::isa<swift::ClangModuleUnit>(decl->getDeclContext())) {
+    // The importer creates one global-member extension per nominal type and Clang submodule.
+    for (auto member : decl->getAllMembers()) {
+      if (auto origin = member->getClangNode().getAsDecl()) {
+        auto module = origin->getOwningModule();
+        return initMangled(decl) << fetch(target) << "|clang_globals|"
+                                 << (module ? module->getFullModuleName() : "");
+      }
+    }
+    LOG_WARNING("Unable to identify an imported global-member module; using declaration indexes");
+  }
   auto index = getExtensionOrFilePrivateValueIndex(decl, parent);
   return initMangled(decl) << fetch(target) << index.index
                            << (index.kind == ExtensionOrFilePrivateValueKind::clang ? "_clang"

swift/extractor/mangler/tests/BUILD.bazel

@@ -0,0 +1,28 @@
+load("@rules_python//python:defs.bzl", "py_test")
+load("//swift:rules.bzl", "swift_cc_binary")
+
+swift_cc_binary(
+    name = "mangle",
+    srcs = ["mangle.cpp"],
+    deps = [
+        "//swift/extractor/mangler",
+        "//swift/logging",
+        "//swift/third_party/resources:prebuilt",
+    ],
+)
+
+py_test(
+    name = "test",
+    size = "medium",
+    srcs = ["test.py"],
+    args = [
+        "$(rlocationpath :mangle)",
+        "$(rlocationpath //swift/third_party/resources:dir)",
+    ],
+    data = [
+        ":mangle",
+        "//swift/third_party/resources:dir",
+    ],
+    target_compatible_with = ["@platforms//os:macos"],
+    deps = ["@rules_python//python/runfiles"],
+)

swift/extractor/mangler/tests/mangle.cpp

@@ -0,0 +1,133 @@
+#include "swift/extractor/mangler/SwiftMangler.h"
+#include "swift/logging/SwiftLogging.h"
+
+#include <swift/AST/ASTWalker.h>
+#include <swift/AST/ClangModuleLoader.h>
+#include <swift/AST/Decl.h>
+#include <swift/AST/Expr.h>
+#include <swift/AST/SourceFile.h>
+#include <swift/Basic/InitializeSwiftModules.h>
+#include <swift/Basic/LLVMInitialize.h>
+#include <swift/Basic/Statistic.h>
+#include <swift/Frontend/Frontend.h>
+#include <swift/FrontendTool/FrontendTool.h>
+
+#include <iostream>
+#include <unordered_set>
+#include <vector>
+
+const std::string_view codeql::programName = "mangle-test";
+const std::string_view codeql::extractorName = "swift";
+
+class References : public swift::ASTWalker {
+ public:
+  std::vector<swift::ValueDecl*> declarations;
+
+  PreWalkResult<swift::Expr*> walkToExprPre(swift::Expr* expression) override {
+    swift::ValueDecl* declaration = nullptr;
+    if (auto reference = llvm::dyn_cast<swift::DeclRefExpr>(expression)) {
+      declaration = reference->getDecl();
+    } else if (auto reference = llvm::dyn_cast<swift::MemberRefExpr>(expression)) {
+      declaration = reference->getMember().getDecl();
+    }
+    if (declaration && declaration->getDeclContext()->getAsDecl() &&
+        llvm::isa<swift::ExtensionDecl>(declaration->getDeclContext()->getAsDecl()) &&
+        seen.insert(declaration).second) {
+      declarations.push_back(declaration);
+    }
+    return Action::Continue(expression);
+  }
+
+ private:
+  std::unordered_set<const swift::Decl*> seen;
+};
+
+class Observer : public swift::FrontendObserver {
+ public:
+  explicit Observer(llvm::StringRef lateMember) : lateMember(lateMember) {}
+
+  void parsedArgs(swift::CompilerInvocation& invocation) override {
+    invocation.getFrontendOptions().KeepASTContext = true;
+  }
+
+  void performedCompilation(swift::CompilerInstance& compiler) override {
+    References references;
+    std::vector<swift::ValueDecl*> privateDeclarations;
+    for (auto file : compiler.getMainModule()->getFiles()) {
+      if (auto source = llvm::dyn_cast<swift::SourceFile>(file)) {
+        source->walk(references);
+        llvm::SmallVector<swift::Decl*> declarations;
+        source->getTopLevelDecls(declarations);
+        for (auto declaration : declarations) {
+          if (auto value = llvm::dyn_cast<swift::ValueDecl>(declaration);
+              value && value->getFormalAccess() == swift::AccessLevel::FilePrivate) {
+            privateDeclarations.push_back(value);
+          }
+        }
+      }
+    }
+    auto& counters = compiler.getStatsReporter()->getFrontendCounters();
+    auto before = counters.NumASTBytesAllocated;
+    codeql::SwiftRecursiveMangler mangler;
+    std::vector<std::string> names;
+    for (auto declaration : references.declarations) {
+      names.push_back(emit(mangler, declaration));
+    }
+    for (auto declaration : privateDeclarations) {
+      std::cout << "fileprivate\t" << mangler.mangleDecl(*declaration).str() << '\n';
+    }
+    auto added = counters.NumASTBytesAllocated - before;
+    if (!lateMember.empty() && !references.declarations.empty()) {
+      auto& context = compiler.getASTContext();
+      auto path = swift::ImportPath::Module::Builder(context, "Fixture.Late", '.');
+      if (!context.getModule(path.get())) {
+        std::cerr << "Late module import failed\n";
+        std::exit(1);
+      }
+      auto extension = llvm::cast<swift::ExtensionDecl>(
+          references.declarations.front()->getDeclContext()->getAsDecl());
+      for (auto sibling : extension->getExtendedNominal()->getExtensions()) {
+        for (auto member : sibling->getAllMembers()) {
+          if (auto value = llvm::dyn_cast<swift::ValueDecl>(member);
+              value && value->getName().getBaseName().userFacingName() == lateMember) {
+            emit(mangler, value);
+          }
+        }
+      }
+      for (size_t i = 0; i < references.declarations.size(); ++i) {
+        std::cout << "late-import-stable\t"
+                  << (mangler.mangleDecl(*references.declarations[i]).str() == names[i]) << '\n';
+      }
+    }
+    std::cout << "additional-ast-bytes\t" << added << '\n';
+    std::cout << "references\t" << references.declarations.size() << '\n';
+  }
+
+ private:
+  llvm::StringRef lateMember;
+
+  static std::string emit(codeql::SwiftRecursiveMangler& mangler, swift::ValueDecl* declaration) {
+    auto name = mangler.mangleDecl(*declaration).str();
+    auto repeated = mangler.mangleDecl(*declaration).str();
+    auto extension = llvm::cast<swift::ExtensionDecl>(declaration->getDeclContext()->getAsDecl());
+    std::cout << declaration->getName().getBaseName().userFacingName().str() << '\t' << name << '\t'
+              << mangler.mangleDecl(*extension).str() << '\t' << (name == repeated) << '\n';
+    return name;
+  }
+};
+
+int main(int argc, char** argv) {
+  PROGRAM_START(argc, argv);
+  INITIALIZE_LLVM();
+  initializeSwiftModules();
+  llvm::StringRef lateMember;
+  if (argc > 1 && std::string_view(argv[1]) == "--late-import") {
+    lateMember = "late";
+  } else if (argc > 1 && std::string_view(argv[1]) == "--late-global-import") {
+    lateMember = "lateValue";
+  }
+  Observer observer(lateMember);
+  auto offset = lateMember.empty() ? 1 : 2;
+  return swift::performFrontend({argv + offset, static_cast<size_t>(argc - offset)}, "mangle-test",
+                                reinterpret_cast<void*>(main), &observer);
+}