Merge pull request #10175 from github/redsun82/swift-missing-extractions

Swift: fix missing extractions

Author: redsun82

swift/codegen/schema.yml

@@ -1121,6 +1121,8 @@ ModuleDecl:
   _extends: TypeDecl
   is_builtin_module: predicate
   is_system_module: predicate
+  imported_modules: ModuleDecl*
+  exported_modules: ModuleDecl*
 
 ConstructorRefCallExpr:
   _extends: SelfApplyExpr

swift/extractor/SwiftBuiltinSymbols.h

@@ -0,0 +1,71 @@
+#pragma once
+
+#include <array>
+
+namespace codeql {
+constexpr std::array swiftBuiltins = {
+    "zeroInitializer",
+    "BridgeObject",
+    "Word",
+    "NativeObject",
+    "RawPointer",
+    "Executor",
+    "Job",
+    "RawUnsafeContinuation",
+    "addressof",
+    "initialize",
+    "reinterpretCast",
+    "Int1",
+    "Int8",
+    "Int16",
+    "Int32",
+    "Int64",
+    "IntLiteral",
+    "FPIEEE16",
+    "FPIEEE32",
+    "FPIEEE64",
+    "FPIEEE80",
+    "Vec2xInt8",
+    "Vec4xInt8",
+    "Vec8xInt8",
+    "Vec16xInt8",
+    "Vec32xInt8",
+    "Vec64xInt8",
+    "Vec2xInt16",
+    "Vec4xInt16",
+    "Vec8xInt16",
+    "Vec16xInt16",
+    "Vec32xInt16",
+    "Vec64xInt16",
+    "Vec2xInt32",
+    "Vec4xInt32",
+    "Vec8xInt32",
+    "Vec16xInt32",
+    "Vec32xInt32",
+    "Vec64xInt32",
+    "Vec2xInt64",
+    "Vec4xInt64",
+    "Vec8xInt64",
+    "Vec16xInt64",
+    "Vec32xInt64",
+    "Vec64xInt64",
+    "Vec2xFPIEEE16",
+    "Vec4xFPIEEE16",
+    "Vec8xFPIEEE16",
+    "Vec16xFPIEEE16",
+    "Vec32xFPIEEE16",
+    "Vec64xFPIEEE16",
+    "Vec2xFPIEEE32",
+    "Vec4xFPIEEE32",
+    "Vec8xFPIEEE32",
+    "Vec16xFPIEEE32",
+    "Vec32xFPIEEE32",
+    "Vec64xFPIEEE32",
+    "Vec2xFPIEEE64",
+    "Vec4xFPIEEE64",
+    "Vec8xFPIEEE64",
+    "Vec16xFPIEEE64",
+    "Vec32xFPIEEE64",
+    "Vec64xFPIEEE64",
+};
+}

swift/extractor/SwiftExtractor.cpp

@@ -6,15 +6,12 @@
 #include <queue>
 
 #include <swift/AST/SourceFile.h>
-#include <swift/Basic/FileTypes.h>
-#include <llvm/ADT/SmallString.h>
-#include <llvm/Support/FileSystem.h>
-#include <llvm/Support/Path.h>
+#include <swift/AST/Builtins.h>
 
-#include "swift/extractor/trap/generated/TrapClasses.h"
 #include "swift/extractor/trap/TrapDomain.h"
 #include "swift/extractor/visitors/SwiftVisitor.h"
 #include "swift/extractor/TargetTrapFile.h"
+#include "swift/extractor/SwiftBuiltinSymbols.h"
 
 using namespace codeql;
 using namespace std::string_literals;
@@ -65,7 +62,40 @@ static std::string getFilename(swift::ModuleDecl& module, swift::SourceFile* pri
     filename += module.getName().str();
     return filename;
   }
-  return module.getModuleFilename().str();
+  if (module.isBuiltinModule()) {
+    // The Builtin module has an empty filename, let's fix that
+    return "/__Builtin__";
+  }
+  auto filename = module.getModuleFilename().str();
+  // there is a special case of a module without an actual filename reporting `<imports>`: in this
+  // case we want to avoid the `<>` characters, in case a dirty DB is imported on Windows
+  if (filename == "<imports>") {
+    return "/__imports__";
+  }
+  return filename;
+}
+
+/* The builtin module is special, as it does not publish any top-level declaration
+ * It creates (and caches) declarations on demand when a lookup is carried out
+ * (see BuiltinUnit in swift/AST/FileUnit.h for the cache details, and getBuiltinValueDecl in
+ * swift/AST/Builtins.h for the creation details)
+ * As we want to create the Builtin trap file once and for all so that it works for other
+ * extraction runs, rather than collecting what we need we pre-populate the builtin trap with
+ * what we expect. This list might need thus to be expanded.
+ * Notice, that while swift/AST/Builtins.def has a list of builtin symbols, it does not contain
+ * all information required to instantiate builtin variants.
+ * Other possible approaches:
+ * * create one trap per builtin declaration when encountered
+ * * expand the list to all possible builtins (of which there are a lot)
+ */
+static void getBuiltinDecls(swift::ModuleDecl& builtinModule,
+                            llvm::SmallVector<swift::Decl*>& decls) {
+  llvm::SmallVector<swift::ValueDecl*> values;
+  for (auto symbol : swiftBuiltins) {
+    builtinModule.lookupValue(builtinModule.getASTContext().getIdentifier(symbol),
+                              swift::NLKind::QualifiedLookup, values);
+  }
+  decls.insert(decls.end(), values.begin(), values.end());
 }
 
 static llvm::SmallVector<swift::Decl*> getTopLevelDecls(swift::ModuleDecl& module,
@@ -74,16 +104,19 @@ static llvm::SmallVector<swift::Decl*> getTopLevelDecls(swift::ModuleDecl& modul
   ret.push_back(&module);
   if (primaryFile) {
     primaryFile->getTopLevelDecls(ret);
+  } else if (module.isBuiltinModule()) {
+    getBuiltinDecls(module, ret);
   } else {
     module.getTopLevelDecls(ret);
   }
   return ret;
 }
 
-static void extractDeclarations(const SwiftExtractorConfiguration& config,
-                                swift::CompilerInstance& compiler,
-                                swift::ModuleDecl& module,
-                                swift::SourceFile* primaryFile = nullptr) {
+static std::unordered_set<swift::ModuleDecl*> extractDeclarations(
+    const SwiftExtractorConfiguration& config,
+    swift::CompilerInstance& compiler,
+    swift::ModuleDecl& module,
+    swift::SourceFile* primaryFile = nullptr) {
   auto filename = getFilename(module, primaryFile);
 
   // The extractor can be called several times from different processes with
@@ -92,7 +125,7 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
   auto trapTarget = createTargetTrapFile(config, filename);
   if (!trapTarget) {
     // another process arrived first, nothing to do for us
-    return;
+    return {};
   }
   TrapDomain trap{*trapTarget};
 
@@ -116,6 +149,7 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
   for (auto& comment : comments) {
     visitor.extract(comment);
   }
+  return std::move(visitor).getEncounteredModules();
 }
 
 static std::unordered_set<std::string> collectInputFilenames(swift::CompilerInstance& compiler) {
@@ -132,40 +166,27 @@ static std::unordered_set<std::string> collectInputFilenames(swift::CompilerInst
   return sourceFiles;
 }
 
-static std::unordered_set<swift::ModuleDecl*> collectModules(swift::CompilerInstance& compiler) {
-  // getASTContext().getLoadedModules() does not provide all the modules available within the
-  // program.
-  // We need to iterate over all the imported modules (recursively) to see the whole "universe."
-  std::unordered_set<swift::ModuleDecl*> allModules;
-  std::queue<swift::ModuleDecl*> worklist;
-  for (auto& [_, module] : compiler.getASTContext().getLoadedModules()) {
-    worklist.push(module);
-    allModules.insert(module);
-  }
-
-  while (!worklist.empty()) {
-    auto module = worklist.front();
-    worklist.pop();
-    llvm::SmallVector<swift::ImportedModule> importedModules;
-    // TODO: we may need more than just Exported ones
-    module->getImportedModules(importedModules, swift::ModuleDecl::ImportFilterKind::Exported);
-    for (auto& imported : importedModules) {
-      if (allModules.count(imported.importedModule) == 0) {
-        worklist.push(imported.importedModule);
-        allModules.insert(imported.importedModule);
-      }
-    }
+static std::vector<swift::ModuleDecl*> collectLoadedModules(swift::CompilerInstance& compiler) {
+  std::vector<swift::ModuleDecl*> ret;
+  for (const auto& [id, module] : compiler.getASTContext().getLoadedModules()) {
+    std::ignore = id;
+    ret.push_back(module);
   }
-  return allModules;
+  return ret;
 }
 
 void codeql::extractSwiftFiles(const SwiftExtractorConfiguration& config,
                                swift::CompilerInstance& compiler) {
   auto inputFiles = collectInputFilenames(compiler);
-  auto modules = collectModules(compiler);
+  std::vector<swift::ModuleDecl*> todo = collectLoadedModules(compiler);
+  std::unordered_set<swift::ModuleDecl*> seen{todo.begin(), todo.end()};
 
-  for (auto& module : modules) {
+  while (!todo.empty()) {
+    auto module = todo.back();
+    todo.pop_back();
+    llvm::errs() << "processing module " << module->getName() << '\n';
     bool isFromSourceFile = false;
+    std::unordered_set<swift::ModuleDecl*> encounteredModules;
     for (auto file : module->getFiles()) {
       auto sourceFile = llvm::dyn_cast<swift::SourceFile>(file);
       if (!sourceFile) {
@@ -176,10 +197,16 @@ void codeql::extractSwiftFiles(const SwiftExtractorConfiguration& config,
         continue;
       }
       archiveFile(config, *sourceFile);
-      extractDeclarations(config, compiler, *module, sourceFile);
+      encounteredModules = extractDeclarations(config, compiler, *module, sourceFile);
     }
     if (!isFromSourceFile) {
-      extractDeclarations(config, compiler, *module);
+      encounteredModules = extractDeclarations(config, compiler, *module);
+    }
+    for (auto encountered : encounteredModules) {
+      if (seen.count(encountered) == 0) {
+        todo.push_back(encountered);
+        seen.insert(encountered);
+      }
     }
   }
 }

swift/extractor/infra/SwiftDispatcher.h

@@ -55,6 +55,10 @@ class SwiftDispatcher {
     }
   }
 
+  const std::unordered_set<swift::ModuleDecl*> getEncounteredModules() && {
+    return std::move(encounteredModules);
+  }
+
   template <typename Entry>
   void emit(const Entry& entry) {
     trap.emit(entry);
@@ -228,8 +232,16 @@ class SwiftDispatcher {
   //  - extracting a primary source file: in this mode, we extract several files belonging to the
   //    same module one by one. In this mode, we restrict emission only to the same file ignoring
   //    all the other files.
+  // This is also used to register the modules we encounter.
+  // TODO calls to this function should be taken away from `DeclVisitor` and moved around with a
+  // clearer separation between naming entities (some decls, all types), deciding whether to emit
+  // them and finally visiting emitting the contents of the entity (which should remain in the
+  // visitors). Then this double responsibility (carrying out the test and registering encountered
+  // modules) should also be cleared out
   bool shouldEmitDeclBody(const swift::Decl& decl) {
-    if (decl.getModuleContext() != &currentModule) {
+    auto module = decl.getModuleContext();
+    if (module != &currentModule) {
+      encounteredModules.insert(module);
       return false;
     }
     // ModuleDecl is a special case: if it passed the previous test, it is the current module
@@ -333,6 +345,7 @@ class SwiftDispatcher {
   Store::Handle waitingForNewLabel{std::monostate{}};
   swift::ModuleDecl& currentModule;
   swift::SourceFile* currentPrimarySourceFile;
+  std::unordered_set<swift::ModuleDecl*> encounteredModules;
 };
 
 }  // namespace codeql