Swift: extract lazy declarations

Author: AlexDenisov

swift/extractor/SwiftBuiltinSymbols.h

@@ -10,10 +10,10 @@
 
 #include "swift/extractor/translators/SwiftVisitor.h"
 #include "swift/extractor/infra/TargetDomains.h"
-#include "swift/extractor/SwiftBuiltinSymbols.h"
 #include "swift/extractor/infra/file/Path.h"
 #include "swift/extractor/infra/SwiftLocationExtractor.h"
 #include "swift/extractor/infra/SwiftBodyEmissionStrategy.h"
+#include "swift/extractor/mangler/SwiftMangler.h"
 
 using namespace codeql;
 using namespace std::string_literals;
@@ -43,10 +43,16 @@ static void archiveFile(const SwiftExtractorConfiguration& config, swift::Source
   }
 }
 
-static fs::path getFilename(swift::ModuleDecl& module, swift::SourceFile* primaryFile) {
+static fs::path getFilename(swift::ModuleDecl& module,
+                            swift::SourceFile* primaryFile,
+                            const swift::Decl* lazyDeclaration) {
   if (primaryFile) {
     return resolvePath(primaryFile->getFilename());
   }
+  if (lazyDeclaration) {
+    static SwiftMangler mangler;
+    return mangler.mangledName(*lazyDeclaration);
+  }
   // PCM clang module
   if (module.isNonSwiftModule()) {
     // Several modules with different names might come from .pcm (clang module) files
@@ -72,57 +78,48 @@ static fs::path getFilename(swift::ModuleDecl& module, swift::SourceFile* primar
   return resolvePath(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,
-                                                        swift::SourceFile* primaryFile = nullptr) {
+                                                        swift::SourceFile* primaryFile,
+                                                        const swift::Decl* lazyDeclaration) {
   llvm::SmallVector<swift::Decl*> ret;
+  if (lazyDeclaration) {
+    ret.push_back(const_cast<swift::Decl*>(lazyDeclaration));
+    return ret;
+  }
   ret.push_back(&module);
   if (primaryFile) {
     primaryFile->getTopLevelDecls(ret);
-  } else if (module.isBuiltinModule()) {
-    getBuiltinDecls(module, ret);
   } else {
     module.getTopLevelDecls(ret);
   }
   return ret;
 }
 
+static TrapType getTrapType(swift::SourceFile* primaryFile, const swift::Decl* lazyDeclaration) {
+  if (primaryFile) {
+    return TrapType::source;
+  }
+  if (lazyDeclaration) {
+    return TrapType::lazy_declarations;
+  }
+  return TrapType::module;
+}
+
 static std::unordered_set<swift::ModuleDecl*> extractDeclarations(
     SwiftExtractorState& state,
     swift::CompilerInstance& compiler,
     swift::ModuleDecl& module,
-    swift::SourceFile* primaryFile = nullptr) {
-  auto filename = getFilename(module, primaryFile);
+    swift::SourceFile* primaryFile,
+    const swift::Decl* lazyDeclaration) {
+  auto filename = getFilename(module, primaryFile, lazyDeclaration);
   if (primaryFile) {
     state.sourceFiles.push_back(filename);
   }
 
   // The extractor can be called several times from different processes with
   // the same input file(s). Using `TargetFile` the first process will win, and the following
   // will just skip the work
-  const auto trapType = primaryFile ? TrapType::source : TrapType::module;
+  const auto trapType = getTrapType(primaryFile, lazyDeclaration);
   auto trap = createTargetTrapDomain(state, filename, trapType);
   if (!trap) {
     // another process arrived first, nothing to do for us
@@ -143,9 +140,10 @@ static std::unordered_set<swift::ModuleDecl*> extractDeclarations(
 
   SwiftLocationExtractor locationExtractor(*trap);
   locationExtractor.emitFile(primaryFile);
-  SwiftBodyEmissionStrategy bodyEmissionStrategy(module, primaryFile);
-  SwiftVisitor visitor(compiler.getSourceMgr(), *trap, locationExtractor, bodyEmissionStrategy);
-  auto topLevelDecls = getTopLevelDecls(module, primaryFile);
+  SwiftBodyEmissionStrategy bodyEmissionStrategy(module, primaryFile, lazyDeclaration);
+  SwiftVisitor visitor(compiler.getSourceMgr(), state, *trap, locationExtractor,
+                       bodyEmissionStrategy);
+  auto topLevelDecls = getTopLevelDecls(module, primaryFile, lazyDeclaration);
   for (auto decl : topLevelDecls) {
     visitor.extract(decl);
   }
@@ -198,10 +196,10 @@ void codeql::extractSwiftFiles(SwiftExtractorState& state, swift::CompilerInstan
         continue;
       }
       archiveFile(state.configuration, *sourceFile);
-      encounteredModules = extractDeclarations(state, compiler, *module, sourceFile);
+      encounteredModules = extractDeclarations(state, compiler, *module, sourceFile, nullptr);
     }
     if (!isFromSourceFile) {
-      encounteredModules = extractDeclarations(state, compiler, *module);
+      encounteredModules = extractDeclarations(state, compiler, *module, nullptr, nullptr);
     }
     for (auto encountered : encounteredModules) {
       if (state.encounteredModules.count(encountered) == 0) {
@@ -211,3 +209,40 @@ 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));
+
+  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));
+
+  for (auto pending : worklist) {
+    extractDeclarations(state, compiler, *pending->getModuleContext(), nullptr, pending);
+  }
+}
+
+void codeql::extractExtractLazyDeclarations(SwiftExtractorState& state,
+                                            swift::CompilerInstance& compiler) {
+  // Just in case
+  const int upperBound = 100;
+  int iteration = 0;
+  while (!state.pendingDeclarations.empty() && iteration++ < upperBound) {
+    extractLazy(state, compiler);
+  }
+  if (iteration >= upperBound) {
+    std::cerr << "Swift extractor reach upper bound while extracting lazy declarations\n";
+    abort();
+  }
+}

swift/extractor/SwiftExtractor.cpp

@@ -7,4 +7,5 @@
 
 namespace codeql {
 void extractSwiftFiles(SwiftExtractorState& state, swift::CompilerInstance& compiler);
+void extractExtractLazyDeclarations(SwiftExtractorState& state, swift::CompilerInstance& compiler);
 }  // namespace codeql

swift/extractor/SwiftExtractor.h

@@ -24,6 +24,9 @@ struct SwiftExtractorState {
 
   // The path for the modules outputted by the underlying frontend run, ignoring path redirection
   std::vector<std::filesystem::path> originalOutputModules;
+
+  std::unordered_set<const swift::Decl*> emittedDeclarations;
+  std::unordered_set<const swift::Decl*> pendingDeclarations;
 };
 
 }  // namespace codeql

swift/extractor/config/SwiftExtractorState.h

@@ -16,6 +16,9 @@ bool SwiftBodyEmissionStrategy::shouldEmitDeclBody(const swift::Decl& decl) {
   if (module != &currentModule) {
     return false;
   }
+  if (currentLazyDeclaration && currentLazyDeclaration != &decl) {
+    return false;
+  }
   // 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) {

swift/extractor/infra/SwiftBodyEmissionStrategy.cpp

@@ -8,13 +8,17 @@ namespace codeql {
 class SwiftBodyEmissionStrategy {
  public:
   SwiftBodyEmissionStrategy(swift::ModuleDecl& currentModule,
-                            swift::SourceFile* currentPrimarySourceFile)
-      : currentModule(currentModule), currentPrimarySourceFile(currentPrimarySourceFile) {}
+                            swift::SourceFile* currentPrimarySourceFile,
+                            const swift::Decl* currentLazyDeclaration)
+      : currentModule(currentModule),
+        currentPrimarySourceFile(currentPrimarySourceFile),
+        currentLazyDeclaration(currentLazyDeclaration) {}
   bool shouldEmitDeclBody(const swift::Decl& decl);
 
  private:
   swift::ModuleDecl& currentModule;
   swift::SourceFile* currentPrimarySourceFile;
+  const swift::Decl* currentLazyDeclaration;
 };
 
 }  // namespace codeql

swift/extractor/infra/SwiftBodyEmissionStrategy.h

@@ -12,6 +12,7 @@
 #include "swift/extractor/trap/generated/TrapClasses.h"
 #include "swift/extractor/infra/SwiftLocationExtractor.h"
 #include "swift/extractor/infra/SwiftBodyEmissionStrategy.h"
+#include "swift/extractor/config/SwiftExtractorState.h"
 
 namespace codeql {
 
@@ -45,10 +46,12 @@ class SwiftDispatcher {
   // all references and pointers passed as parameters to this constructor are supposed to outlive
   // the SwiftDispatcher
   SwiftDispatcher(const swift::SourceManager& sourceManager,
+                  SwiftExtractorState& state,
                   TrapDomain& trap,
                   SwiftLocationExtractor& locationExtractor,
                   SwiftBodyEmissionStrategy& bodyEmissionStrategy)
       : sourceManager{sourceManager},
+        state{state},
         trap{trap},
         locationExtractor{locationExtractor},
         bodyEmissionStrategy{bodyEmissionStrategy} {}
@@ -57,6 +60,9 @@ class SwiftDispatcher {
     return std::move(encounteredModules);
   }
 
+  void extractedDeclaration(const swift::Decl* decl) { state.emittedDeclarations.insert(decl); }
+  void skippedDeclaration(const swift::Decl* decl) { state.pendingDeclarations.insert(decl); }
+
   template <typename Entry>
   void emit(Entry&& entry) {
     bool valid = true;
@@ -302,6 +308,7 @@ class SwiftDispatcher {
   virtual void visit(const swift::CapturedValue* capture) = 0;
 
   const swift::SourceManager& sourceManager;
+  SwiftExtractorState& state;
   TrapDomain& trap;
   Store store;
   SwiftLocationExtractor& locationExtractor;

swift/extractor/infra/SwiftDispatcher.h

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

swift/extractor/infra/TargetDomains.cpp

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

swift/extractor/infra/TargetDomains.h

@@ -92,6 +92,7 @@ class Observer : public swift::FrontendObserver {
   void performedSemanticAnalysis(swift::CompilerInstance& compiler) override {
     codeql::extractSwiftFiles(state, compiler);
     codeql::extractSwiftInvocation(state, compiler, invocationTrap);
+    codeql::extractExtractLazyDeclarations(state, compiler);
   }
 
  private: