Swift: address review comments

swift/extractor/SwiftExtractor.cpp

@@ -78,20 +78,22 @@ static fs::path getFilename(swift::ModuleDecl& module,
   return resolvePath(filename);
 }
 
-static llvm::SmallVector<swift::Decl*> getTopLevelDecls(swift::ModuleDecl& module,
-                                                        swift::SourceFile* primaryFile,
-                                                        const swift::Decl* lazyDeclaration) {
-  llvm::SmallVector<swift::Decl*> ret;
+static llvm::SmallVector<const swift::Decl*> getTopLevelDecls(swift::ModuleDecl& module,
+                                                              swift::SourceFile* primaryFile,
+                                                              const swift::Decl* lazyDeclaration) {
+  llvm::SmallVector<const swift::Decl*> ret;
   if (lazyDeclaration) {
-    ret.push_back(const_cast<swift::Decl*>(lazyDeclaration));
+    ret.push_back(lazyDeclaration);
     return ret;
   }
   ret.push_back(&module);
+  llvm::SmallVector<swift::Decl*> topLevelDecls;
   if (primaryFile) {
-    primaryFile->getTopLevelDecls(ret);
+    primaryFile->getTopLevelDecls(topLevelDecls);
   } else {
-    module.getTopLevelDecls(ret);
+    module.getTopLevelDecls(topLevelDecls);
   }
+  ret.insert(ret.end(), topLevelDecls.data(), topLevelDecls.data() + topLevelDecls.size());
   return ret;
 }
 
@@ -100,7 +102,7 @@ static TrapType getTrapType(swift::SourceFile* primaryFile, const swift::Decl* l
     return TrapType::source;
   }
   if (lazyDeclaration) {
-    return TrapType::lazy_declarations;
+    return TrapType::lazy_declaration;
   }
   return TrapType::module;
 }
@@ -199,10 +201,12 @@ void codeql::extractSwiftFiles(SwiftExtractorState& state, swift::CompilerInstan
         continue;
       }
       archiveFile(state.configuration, *sourceFile);
-      encounteredModules = extractDeclarations(state, compiler, *module, sourceFile, nullptr);
+      encounteredModules =
+          extractDeclarations(state, compiler, *module, sourceFile, /*lazy declaration*/ nullptr);
     }
     if (!isFromSourceFile) {
-      encounteredModules = extractDeclarations(state, compiler, *module, nullptr, nullptr);
+      encounteredModules = extractDeclarations(state, compiler, *module, /*source file*/ nullptr,
+                                               /*lazy declaration*/ nullptr);
     }
     for (auto encountered : encounteredModules) {
       if (state.encounteredModules.count(encountered) == 0) {
@@ -214,23 +218,22 @@ void codeql::extractSwiftFiles(SwiftExtractorState& state, swift::CompilerInstan
 }
 
 static void cleanupPendingDeclarations(SwiftExtractorState& state) {
-  std::vector<const swift::Decl*> worklist;
-  std::copy(std::begin(state.pendingDeclarations), std::end(state.pendingDeclarations),
-            std::back_inserter(worklist));
-
+  std::vector<const swift::Decl*> worklist(std::begin(state.pendingDeclarations),
+                                           std::end(state.pendingDeclarations));
   for (auto decl : worklist) {
     if (state.emittedDeclarations.count(decl)) {
       state.pendingDeclarations.erase(decl);
     }
   }
 }
+
 static void extractLazy(SwiftExtractorState& state, swift::CompilerInstance& compiler) {
   cleanupPendingDeclarations(state);
-  std::vector<const swift::Decl*> worklist;
-  std::copy(std::begin(state.pendingDeclarations), std::end(state.pendingDeclarations),
-            std::back_inserter(worklist));
+  std::vector<const swift::Decl*> worklist(std::begin(state.pendingDeclarations),
+                                           std::end(state.pendingDeclarations));
   for (auto pending : worklist) {
-    extractDeclarations(state, compiler, *pending->getModuleContext(), nullptr, pending);
+    extractDeclarations(state, compiler, *pending->getModuleContext(), /*source file*/ nullptr,
+                        pending);
   }
 }
 

swift/extractor/config/SwiftExtractorState.h

@@ -25,7 +25,10 @@ struct SwiftExtractorState {
   // The path for the modules outputted by the underlying frontend run, ignoring path redirection
   std::vector<std::filesystem::path> originalOutputModules;
 
+  // All lazy named declarations that were already emitted
   std::unordered_set<const swift::Decl*> emittedDeclarations;
+
+  // Lazy named declarations that were not yet emitted and will be emitted each one separately
   std::unordered_set<const swift::Decl*> pendingDeclarations;
 };
 

swift/extractor/infra/SwiftDispatcher.h

@@ -60,19 +60,6 @@ class SwiftDispatcher {
     return std::move(encounteredModules);
   }
 
-  void extractedDeclaration(const swift::Decl* decl) {
-    swift::ModuleDecl* module = decl->getModuleContext();
-    if (module->isBuiltinModule() || module->getName().str() == "__ObjC") {
-      state.emittedDeclarations.insert(decl);
-    }
-  }
-  void skippedDeclaration(const swift::Decl* decl) {
-    swift::ModuleDecl* module = decl->getModuleContext();
-    if (module->isBuiltinModule() || module->getName().str() == "__ObjC") {
-      state.pendingDeclarations.insert(decl);
-    }
-  }
-
   template <typename Entry>
   void emit(Entry&& entry) {
     bool valid = true;
@@ -264,7 +251,23 @@ class SwiftDispatcher {
     locationExtractor.attachLocation(sourceManager, comment, entry.id);
   }
 
+  void extractedDeclaration(const swift::Decl& decl) {
+    if (isLazyDeclaration(decl)) {
+      state.emittedDeclarations.insert(&decl);
+    }
+  }
+  void skippedDeclaration(const swift::Decl& decl) {
+    if (isLazyDeclaration(decl)) {
+      state.pendingDeclarations.insert(&decl);
+    }
+  }
+
  private:
+  bool isLazyDeclaration(const swift::Decl& decl) {
+    swift::ModuleDecl* module = decl.getModuleContext();
+    return module->isBuiltinModule() || module->getName().str() == "__ObjC";
+  }
+
   template <typename T, typename = void>
   struct HasSize : std::false_type {};
 

swift/extractor/infra/TargetDomains.cpp

@@ -13,7 +13,7 @@ static const char* typeToStr(TrapType type) {
       return "invocations";
     case TrapType::linkage:
       return "linkage";
-    case TrapType::lazy_declarations:
+    case TrapType::lazy_declaration:
       return "lazy_decls";
     default:
       return "";

swift/extractor/infra/TargetDomains.h

@@ -12,7 +12,7 @@ enum class TrapType {
   module,
   invocation,
   linkage,
-  lazy_declarations,
+  lazy_declaration,
 };
 
 std::filesystem::path getTrapPath(const SwiftExtractorState& state,

swift/extractor/translators/DeclTranslator.h

@@ -71,11 +71,11 @@ class DeclTranslator : public AstTranslatorBase<DeclTranslator> {
     std::optional<TrapClassOf<D>> entry;
     auto id = dispatcher.assignNewLabel(decl, mangledName(decl));
     if (dispatcher.shouldEmitDeclBody(decl)) {
-      dispatcher.extractedDeclaration(&decl);
+      dispatcher.extractedDeclaration(decl);
       entry.emplace(id);
       fillDecl(decl, *entry);
     } else {
-      dispatcher.skippedDeclaration(&decl);
+      dispatcher.skippedDeclaration(decl);
     }
     return entry;
   }