Swift: extract ImportDecl and ModuleDecl

As `ASTMangler` crashes when called on `ModuleDecl`, we simply use
its name.

This might probably not work reliably in a scenario where multiple
modules are compiled with the same name (like `main`), but this is left
for future work. At the moment this cannot create DB inconsistencies.

Author: redsun82

swift/codegen/schema.yml

@@ -273,6 +273,9 @@ IfConfigDecl:
 
 ImportDecl:
   _extends: Decl
+  is_exported: predicate
+  module: ModuleDecl
+  declarations: ValueDecl*
 
 MissingMemberDecl:
   _extends: Decl
@@ -1067,6 +1070,8 @@ GenericTypeDecl:
 
 ModuleDecl:
   _extends: TypeDecl
+  is_builtin_module: predicate
+  is_system_module: predicate
 
 ConstructorRefCallExpr:
   _extends: SelfApplyExpr

swift/extractor/SwiftExtractor.cpp

@@ -64,8 +64,19 @@ static std::string getTrapFilename(swift::ModuleDecl& module, swift::SourceFile*
   return filename;
 }
 
+static llvm::SmallVector<swift::Decl*> getTopLevelDecls(swift::ModuleDecl& module,
+                                                        swift::SourceFile* primaryFile = nullptr) {
+  llvm::SmallVector<swift::Decl*> ret;
+  ret.push_back(&module);
+  if (primaryFile) {
+    primaryFile->getTopLevelDecls(ret);
+  } else {
+    module.getTopLevelDecls(ret);
+  }
+  return ret;
+}
+
 static void extractDeclarations(const SwiftExtractorConfiguration& config,
-                                llvm::ArrayRef<swift::Decl*> topLevelDecls,
                                 swift::CompilerInstance& compiler,
                                 swift::ModuleDecl& module,
                                 swift::SourceFile* primaryFile = nullptr) {
@@ -119,6 +130,7 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
   trap.emit(LocationsTrap{unknownLocationLabel, unknownFileLabel});
 
   SwiftVisitor visitor(compiler.getSourceMgr(), arena, trap, module, primaryFile);
+  auto topLevelDecls = getTopLevelDecls(module, primaryFile);
   for (auto decl : topLevelDecls) {
     visitor.extract(decl);
   }
@@ -203,18 +215,15 @@ void codeql::extractSwiftFiles(const SwiftExtractorConfiguration& config,
     // user code twice: once during the module build in a form of a source file, and then as
     // a pre-built module during building of the dependent source files.
     if (module->isSystemModule() || module->isBuiltinModule()) {
-      llvm::SmallVector<swift::Decl*> decls;
-      module->getTopLevelDecls(decls);
-      // TODO: pass ModuleDecl directly when we have module extraction in place?
-      extractDeclarations(config, decls, compiler, *module);
+      extractDeclarations(config, compiler, *module);
     } else {
       for (auto file : module->getFiles()) {
         auto sourceFile = llvm::dyn_cast<swift::SourceFile>(file);
         if (!sourceFile || inputFiles.count(sourceFile->getFilename().str()) == 0) {
           continue;
         }
         archiveFile(config, *sourceFile);
-        extractDeclarations(config, sourceFile->getTopLevelDecls(), compiler, *module, sourceFile);
+        extractDeclarations(config, compiler, *module, sourceFile);
       }
     }
   }

swift/extractor/infra/SwiftDispatcher.h

@@ -179,6 +179,11 @@ class SwiftDispatcher {
     return ret;
   }
 
+  template <typename... Args>
+  void emitDebugInfo(const Args&... args) {
+    trap.debug(std::forward<Args>(args)...);
+  }
+
   // In order to not emit duplicated entries for declarations, we restrict emission to only
   // Decls declared within the current "scope".
   // Depending on the whether we are extracting a primary source file or not the scope is defined as
@@ -192,7 +197,9 @@ class SwiftDispatcher {
     if (decl.getModuleContext() != &currentModule) {
       return false;
     }
-    if (!currentPrimarySourceFile) {
+    // ModuleDecl is a special case: if it passed the previous test, it is the current module
+    // but it never has a source file, so we short circuit to emit it in any case
+    if (!currentPrimarySourceFile || decl.getKind() == swift::DeclKind::Module) {
       return true;
     }
     if (auto context = decl.getDeclContext()) {

swift/extractor/trap/TrapOutput.h

@@ -41,8 +41,8 @@ class TrapOutput {
 
   template <typename... Args>
   void debug(const Args&... args) {
-    out_ << "// DEBUG: ";
-    (out_ << ... << args) << '\n';
+    out_ << "/* DEBUG:\n";
+    (out_ << ... << args) << "\n*/\n";
   }
 
  private:

swift/extractor/visitors/DeclVisitor.cpp

@@ -241,16 +241,15 @@ std::variant<codeql::AccessorDecl, codeql::AccessorDeclsTrap> DeclVisitor::trans
 
 std::optional<codeql::SubscriptDecl> DeclVisitor::translateSubscriptDecl(
     const swift::SubscriptDecl& decl) {
-  auto id = dispatcher_.assignNewLabel(decl, mangledName(decl));
-  if (!dispatcher_.shouldEmitDeclBody(decl)) {
+  auto entry = createNamedEntry(decl);
+  if (!entry) {
     return std::nullopt;
   }
-  SubscriptDecl entry{id};
-  entry.element_type = dispatcher_.fetchLabel(decl.getElementInterfaceType());
+  entry->element_type = dispatcher_.fetchLabel(decl.getElementInterfaceType());
   if (auto indices = decl.getIndices()) {
-    entry.params = dispatcher_.fetchRepeatedLabels(*indices);
+    entry->params = dispatcher_.fetchRepeatedLabels(*indices);
   }
-  fillAbstractStorageDecl(decl, entry);
+  fillAbstractStorageDecl(decl, *entry);
   return entry;
 }
 
@@ -262,7 +261,32 @@ codeql::ExtensionDecl DeclVisitor::translateExtensionDecl(const swift::Extension
   return entry;
 }
 
+codeql::ImportDecl DeclVisitor::translateImportDecl(const swift::ImportDecl& decl) {
+  auto entry = dispatcher_.createEntry(decl);
+  entry.is_exported = decl.isExported();
+  entry.module = dispatcher_.fetchLabel(decl.getModule());
+  entry.declarations = dispatcher_.fetchRepeatedLabels(decl.getDecls());
+  return entry;
+}
+
+std::optional<codeql::ModuleDecl> DeclVisitor::translateModuleDecl(const swift::ModuleDecl& decl) {
+  auto entry = createNamedEntry(decl);
+  if (!entry) {
+    return std::nullopt;
+  }
+  entry->is_builtin_module = decl.isBuiltinModule();
+  entry->is_system_module = decl.isSystemModule();
+  fillTypeDecl(decl, *entry);
+  return entry;
+}
+
 std::string DeclVisitor::mangledName(const swift::ValueDecl& decl) {
+  // ASTMangler::mangleAnyDecl crashes when called on `ModuleDecl`
+  // TODO find a more unique string working also when different modules are compiled with the same
+  // name
+  if (decl.getKind() == swift::DeclKind::Module) {
+    return static_cast<const swift::ModuleDecl&>(decl).getRealName().str().str();
+  }
   // prefix adds a couple of special symbols, we don't necessary need them
   return mangler.mangleAnyDecl(&decl, /* prefix = */ false);
 }

swift/extractor/visitors/DeclVisitor.h

@@ -50,6 +50,8 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
       const swift::AccessorDecl& decl);
   std::optional<codeql::SubscriptDecl> translateSubscriptDecl(const swift::SubscriptDecl& decl);
   codeql::ExtensionDecl translateExtensionDecl(const swift::ExtensionDecl& decl);
+  codeql::ImportDecl translateImportDecl(const swift::ImportDecl& decl);
+  std::optional<codeql::ModuleDecl> translateModuleDecl(const swift::ModuleDecl& decl);
 
  private:
   std::string mangledName(const swift::ValueDecl& decl);
@@ -66,6 +68,15 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
   void fillAbstractStorageDecl(const swift::AbstractStorageDecl& decl,
                                codeql::AbstractStorageDecl& entry);
 
+  template <typename D>
+  std::optional<TrapClassOf<D>> createNamedEntry(const D& decl) {
+    auto id = dispatcher_.assignNewLabel(decl, mangledName(decl));
+    if (dispatcher_.shouldEmitDeclBody(decl)) {
+      return TrapClassOf<D>{id};
+    }
+    return std::nullopt;
+  }
+
  private:
   swift::Mangle::ASTMangler mangler;
 };

swift/ql/lib/codeql/swift/generated/decl/ImportDecl.qll

@@ -1,6 +1,28 @@
 // generated by codegen/codegen.py
 import codeql.swift.elements.decl.Decl
+import codeql.swift.elements.decl.ModuleDecl
+import codeql.swift.elements.decl.ValueDecl
 
 class ImportDeclBase extends @import_decl, Decl {
   override string getAPrimaryQlClass() { result = "ImportDecl" }
+
+  predicate isExported() { import_decl_is_exported(this) }
+
+  ModuleDecl getModule() {
+    exists(ModuleDecl x |
+      import_decls(this, x) and
+      result = x.resolve()
+    )
+  }
+
+  ValueDecl getDeclaration(int index) {
+    exists(ValueDecl x |
+      import_decl_declarations(this, index, x) and
+      result = x.resolve()
+    )
+  }
+
+  ValueDecl getADeclaration() { result = getDeclaration(_) }
+
+  int getNumberOfDeclarations() { result = count(getADeclaration()) }
 }

swift/ql/lib/codeql/swift/generated/decl/ModuleDecl.qll

@@ -3,4 +3,8 @@ import codeql.swift.elements.decl.TypeDecl
 
 class ModuleDeclBase extends @module_decl, TypeDecl {
   override string getAPrimaryQlClass() { result = "ModuleDecl" }
+
+  predicate isBuiltinModule() { module_decl_is_builtin_module(this) }
+
+  predicate isSystemModule() { module_decl_is_system_module(this) }
 }

swift/ql/lib/swift.dbscheme

@@ -644,7 +644,20 @@ if_config_decls(  //dir=decl
 );
 
 import_decls(  //dir=decl
-  unique int id: @import_decl
+  unique int id: @import_decl,
+  int module: @module_decl ref
+);
+
+#keyset[id]
+import_decl_is_exported(  //dir=decl
+  int id: @import_decl ref
+);
+
+#keyset[id, index]
+import_decl_declarations(  //dir=decl
+  int id: @import_decl ref,
+  int index: int ref,
+  int declaration: @value_decl ref
 );
 
 missing_member_decls(  //dir=decl
@@ -2018,6 +2031,16 @@ module_decls(  //dir=decl
   unique int id: @module_decl
 );
 
+#keyset[id]
+module_decl_is_builtin_module(  //dir=decl
+  int id: @module_decl ref
+);
+
+#keyset[id]
+module_decl_is_system_module(  //dir=decl
+  int id: @module_decl ref
+);
+
 constructor_ref_call_exprs(  //dir=expr
   unique int id: @constructor_ref_call_expr
 );

swift/ql/test/TestUtils.qll

@@ -4,6 +4,8 @@ cached
 predicate toBeTested(Element e) {
   e instanceof File
   or
+  exists(ModuleDecl m | m = e and not m.isBuiltinModule() and not m.isSystemModule())
+  or
   exists(Locatable loc |
     loc.getLocation().getFile().getName().matches("%swift/ql/test%") and
     (