diff --git a/swift/extractor/SwiftExtractor.cpp b/swift/extractor/SwiftExtractor.cpp index 5d920373fbb..cdee01f446a 100644 --- a/swift/extractor/SwiftExtractor.cpp +++ b/swift/extractor/SwiftExtractor.cpp @@ -170,7 +170,7 @@ static std::unordered_set extractDeclarations( bodyEmissionStrategy); auto topLevelDecls = getTopLevelDecls(module, primaryFile, lazyDeclaration); for (auto decl : topLevelDecls) { - if (decl->isUnavailable()) { + if (decl->isUnavailable() && !llvm::isa(decl)) { continue; } visitor.extract(decl); diff --git a/swift/extractor/infra/SwiftBodyEmissionStrategy.cpp b/swift/extractor/infra/SwiftBodyEmissionStrategy.cpp index 7608379ac5f..71eebcbf6b0 100644 --- a/swift/extractor/infra/SwiftBodyEmissionStrategy.cpp +++ b/swift/extractor/infra/SwiftBodyEmissionStrategy.cpp @@ -12,8 +12,7 @@ using namespace codeql; // same module one by one. In this mode, we restrict emission only to the same file ignoring // all the other files. bool SwiftBodyEmissionStrategy::shouldEmitDeclBody(const swift::Decl& decl) { - auto module = decl.getModuleContext(); - if (module != ¤tModule) { + if (!currentTopLevelDecls.contains(&decl) && decl.getModuleContext() != ¤tModule) { return false; } if (currentLazyDeclaration && currentLazyDeclaration != &decl) { diff --git a/swift/extractor/infra/SwiftBodyEmissionStrategy.h b/swift/extractor/infra/SwiftBodyEmissionStrategy.h index 34eae944a55..b97600cc63c 100644 --- a/swift/extractor/infra/SwiftBodyEmissionStrategy.h +++ b/swift/extractor/infra/SwiftBodyEmissionStrategy.h @@ -12,13 +12,19 @@ class SwiftBodyEmissionStrategy { const swift::Decl* currentLazyDeclaration) : currentModule(currentModule), currentPrimarySourceFile(currentPrimarySourceFile), - currentLazyDeclaration(currentLazyDeclaration) {} + currentLazyDeclaration(currentLazyDeclaration) { + llvm::SmallVector decls; + currentModule.getTopLevelDecls(decls); + currentTopLevelDecls.insert(decls.begin(), decls.end()); + } + bool shouldEmitDeclBody(const swift::Decl& decl); private: swift::ModuleDecl& currentModule; swift::SourceFile* currentPrimarySourceFile; const swift::Decl* currentLazyDeclaration; + std::unordered_set currentTopLevelDecls; }; } // namespace codeql diff --git a/swift/extractor/mangler/SwiftMangler.cpp b/swift/extractor/mangler/SwiftMangler.cpp index 38f0683bd04..b44921853ba 100644 --- a/swift/extractor/mangler/SwiftMangler.cpp +++ b/swift/extractor/mangler/SwiftMangler.cpp @@ -8,6 +8,8 @@ #include #include #include +#include +#include using namespace codeql; @@ -38,6 +40,8 @@ std::string_view getTypeKindStr(const swift::TypeBase* type) { } // namespace +std::unordered_map SwiftMangler::preloadedExtensionIndexes; + SwiftMangledName SwiftMangler::initMangled(const swift::TypeBase* type) { return {getTypeKindStr(type), '_'}; } @@ -109,32 +113,58 @@ unsigned SwiftMangler::getExtensionIndex(const swift::ExtensionDecl* decl, // indexes once for each encountered parent into the `preloadedExtensionIndexes` mapping. // Because we mangle declarations only once in a given trap/dispatcher context, we can safely // discard preloaded indexes on use - if (auto found = preloadedExtensionIndexes.extract(decl)) { - return found.mapped(); + if (auto found = SwiftMangler::preloadedExtensionIndexes.find(decl); + found != SwiftMangler::preloadedExtensionIndexes.end()) { + return found->second; } if (auto parentModule = llvm::dyn_cast(parent)) { llvm::SmallVector siblings; parentModule->getTopLevelDecls(siblings); indexExtensions(siblings); + if (auto clangModule = parentModule->findUnderlyingClangModule()) { + indexClangExtensions(clangModule, decl->getASTContext().getClangModuleLoader()); + } } else if (auto iterableParent = llvm::dyn_cast(parent)) { indexExtensions(iterableParent->getAllMembers()); } else { // TODO use a generic logging handle for Swift entities here, once it's available CODEQL_ASSERT(false, "non-local context must be module or iterable decl context"); } - auto found = preloadedExtensionIndexes.extract(decl); + auto found = SwiftMangler::preloadedExtensionIndexes.find(decl); // TODO use a generic logging handle for Swift entities here, once it's available - CODEQL_ASSERT(found, "extension not found within parent"); - return found.mapped(); + CODEQL_ASSERT(found != SwiftMangler::preloadedExtensionIndexes.end(), + "extension not found within parent"); + return found->second; } void SwiftMangler::indexExtensions(llvm::ArrayRef siblings) { auto index = 0u; for (auto sibling : siblings) { if (sibling->getKind() == swift::DeclKind::Extension) { - preloadedExtensionIndexes.emplace(sibling, index); + SwiftMangler::preloadedExtensionIndexes.emplace(sibling, index); + index += 2; + } + } +} + +void SwiftMangler::indexClangExtensions(const clang::Module* clangModule, + swift::ClangModuleLoader* moduleLoader) { + if (!moduleLoader) { + return; + } + + auto index = 1u; + for (const auto& submodule : clangModule->submodules()) { + if (auto* swiftSubmodule = moduleLoader->getWrapperForModule(submodule)) { + llvm::SmallVector children; + swiftSubmodule->getTopLevelDecls(children); + for (const auto child : children) { + if (child->getKind() == swift::DeclKind::Extension) { + SwiftMangler::preloadedExtensionIndexes.emplace(child, index); + index += 2; + } + } } - ++index; } } diff --git a/swift/extractor/mangler/SwiftMangler.h b/swift/extractor/mangler/SwiftMangler.h index 7ed7b7c761b..da93f3473f6 100644 --- a/swift/extractor/mangler/SwiftMangler.h +++ b/swift/extractor/mangler/SwiftMangler.h @@ -106,13 +106,15 @@ class SwiftMangler : private swift::TypeVisitor, SwiftMangledName visitPackExpansionType(const swift::PackExpansionType* type); private: - std::unordered_map preloadedExtensionIndexes; + static std::unordered_map 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 siblings); + void indexClangExtensions(const clang::Module* clangModule, + swift::ClangModuleLoader* moduleLoader); unsigned int getExtensionIndex(const swift::ExtensionDecl* decl, const swift::Decl* parent); static SwiftMangledName initMangled(const swift::TypeBase* type); SwiftMangledName initMangled(const swift::Decl* decl);