Swift: extract ImportDecl and ModuleDecl

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,10 +130,14 @@ 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);
   }
-  if (topLevelDecls.empty()) {
+  // TODO the following will be moved to the dispatcher when we start caching swift file objects
+  // for the moment, topLevelDecls always contains the current module, which does not have a file
+  // associated with it, so we need a special case when there are no top level declarations
+  if (topLevelDecls.size() == 1) {
     // In the case of empty files, the dispatcher is not called, but we still want to 'record' the
     // fact that the file was extracted
     llvm::SmallString<PATH_MAX> name(filename);
@@ -203,18 +218,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

@@ -112,6 +112,12 @@ class SwiftDispatcher {
     return assignNewLabel(&e, std::forward<Args>(args)...);
   }
 
+  // convenience methods for structured C++ creation
+  template <typename E, typename... Args, std::enable_if_t<!std::is_pointer_v<E>>* = nullptr>
+  auto createEntry(const E& e, Args&&... args) {
+    return TrapClassOf<E>{assignNewLabel(&e, std::forward<Args>(args)...)};
+  }
+
   template <typename Tag>
   TrapLabel<Tag> createLabel() {
     auto ret = arena.allocateLabel<Tag>();
@@ -173,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
@@ -186,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/extractor/visitors/TypeVisitor.cpp

@@ -122,13 +122,13 @@ void TypeVisitor::visitParenType(swift::ParenType* type) {
 }
 
 codeql::OptionalType TypeVisitor::translateOptionalType(const swift::OptionalType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   fillUnarySyntaxSugarType(type, entry);
   return entry;
 }
 
 codeql::ArraySliceType TypeVisitor::translateArraySliceType(const swift::ArraySliceType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   fillUnarySyntaxSugarType(type, entry);
   return entry;
 }
@@ -163,7 +163,7 @@ void TypeVisitor::visitLValueType(swift::LValueType* type) {
 
 codeql::PrimaryArchetypeType TypeVisitor::translatePrimaryArchetypeType(
     const swift::PrimaryArchetypeType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   fillArchetypeType(type, entry);
   return entry;
 }
@@ -229,7 +229,7 @@ void TypeVisitor::emitAnyGenericType(swift::AnyGenericType* type,
 
 codeql::NestedArchetypeType TypeVisitor::translateNestedArchetypeType(
     const swift::NestedArchetypeType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   entry.parent = dispatcher_.fetchLabel(type.getParent());
   entry.associated_type_declaration = dispatcher_.fetchLabel(type.getAssocType());
   fillArchetypeType(type, entry);
@@ -248,26 +248,26 @@ void TypeVisitor::fillArchetypeType(const swift::ArchetypeType& type, ArchetypeT
 }
 
 codeql::ExistentialType TypeVisitor::translateExistentialType(const swift::ExistentialType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   entry.constraint = dispatcher_.fetchLabel(type.getConstraintType());
   return entry;
 }
 
 codeql::DynamicSelfType TypeVisitor::translateDynamicSelfType(const swift::DynamicSelfType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   entry.static_self_type = dispatcher_.fetchLabel(type.getSelfType());
   return entry;
 }
 
 codeql::VariadicSequenceType TypeVisitor::translateVariadicSequenceType(
     const swift::VariadicSequenceType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   fillUnarySyntaxSugarType(type, entry);
   return entry;
 }
 
 codeql::InOutType TypeVisitor::translateInOutType(const swift::InOutType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   entry.object_type = dispatcher_.fetchLabel(type.getObjectType());
   return entry;
 }
@@ -300,19 +300,19 @@ void TypeVisitor::fillReferenceStorageType(const swift::ReferenceStorageType& ty
 
 codeql::ProtocolCompositionType TypeVisitor::translateProtocolCompositionType(
     const swift::ProtocolCompositionType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   entry.members = dispatcher_.fetchRepeatedLabels(type.getMembers());
   return entry;
 }
 
 codeql::BuiltinIntegerLiteralType TypeVisitor::translateBuiltinIntegerLiteralType(
     const swift::BuiltinIntegerLiteralType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinIntegerType TypeVisitor::translateBuiltinIntegerType(
     const swift::BuiltinIntegerType& type) {
-  auto entry = createEntry(type);
+  auto entry = createTypeEntry(type);
   if (type.isFixedWidth()) {
     entry.width = type.getFixedWidth();
   }
@@ -321,51 +321,51 @@ codeql::BuiltinIntegerType TypeVisitor::translateBuiltinIntegerType(
 
 codeql::BuiltinBridgeObjectType TypeVisitor::translateBuiltinBridgeObjectType(
     const swift::BuiltinBridgeObjectType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinDefaultActorStorageType TypeVisitor::translateBuiltinDefaultActorStorageType(
     const swift::BuiltinDefaultActorStorageType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinExecutorType TypeVisitor::translateBuiltinExecutorType(
     const swift::BuiltinExecutorType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinFloatType TypeVisitor::translateBuiltinFloatType(
     const swift::BuiltinFloatType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinJobType TypeVisitor::translateBuiltinJobType(const swift::BuiltinJobType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinNativeObjectType TypeVisitor::translateBuiltinNativeObjectType(
     const swift::BuiltinNativeObjectType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinRawPointerType TypeVisitor::translateBuiltinRawPointerType(
     const swift::BuiltinRawPointerType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinRawUnsafeContinuationType TypeVisitor::translateBuiltinRawUnsafeContinuationType(
     const swift::BuiltinRawUnsafeContinuationType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinUnsafeValueBufferType TypeVisitor::translateBuiltinUnsafeValueBufferType(
     const swift::BuiltinUnsafeValueBufferType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 codeql::BuiltinVectorType TypeVisitor::translateBuiltinVectorType(
     const swift::BuiltinVectorType& type) {
-  return createEntry(type);
+  return createTypeEntry(type);
 }
 
 }  // namespace codeql

swift/extractor/visitors/TypeVisitor.h

@@ -81,8 +81,8 @@ class TypeVisitor : public TypeVisitorBase<TypeVisitor> {
   void emitAnyGenericType(swift::AnyGenericType* type, TrapLabel<AnyGenericTypeTag> label);
 
   template <typename T>
-  auto createEntry(const T& type) {
-    TrapClassOf<T> entry{dispatcher_.assignNewLabel(type)};
+  auto createTypeEntry(const T& type) {
+    auto entry = dispatcher_.createEntry(type);
     fillType(type, entry);
     return entry;
   }

swift/integration-tests/posix-only/partial-modules/Modules.expected

@@ -0,0 +1,4 @@
+| file://:0:0:0:0 | A |
+| file://:0:0:0:0 | B |
+| file://:0:0:0:0 | main |
+| file://:0:0:0:0 | partial_modules |

swift/integration-tests/posix-only/partial-modules/Modules.ql

@@ -0,0 +1,5 @@
+import swift
+
+from ModuleDecl decl
+where not decl.isBuiltinModule() and not decl.isSystemModule()
+select decl