From fb29d1f55fbc2d38b15d17d8641b00a171aaf63f Mon Sep 17 00:00:00 2001
From: Trevor Elliott <awesomelyawesome@gmail.com>
Date: Tue, 15 Feb 2022 15:27:10 -0800
Subject: [PATCH] Split package db loading up from namespace generation (#5311)

---
 main/realmain.cc             | 4 ++--
 packager/rbi_gen.cc          | 6 +-----
 packager/rbi_gen.h           | 3 +--
 test/pipeline_test_runner.cc | 6 +++---
 4 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/main/realmain.cc b/main/realmain.cc
index 9efda73c33d..358ce367696 100644
--- a/main/realmain.cc
+++ b/main/realmain.cc
@@ -455,7 +455,6 @@ int realmain(int argc, char *argv[]) {
     gs->errorUrlBase = opts.errorUrlBase;
     gs->semanticExtensions = move(extensions);
     vector<ast::ParsedFile> indexed;
-    UnorderedSet<core::ClassOrModuleRef> packageNamespaces;
 
     gs->requiresAncestorEnabled = opts.requiresAncestorEnabled;
 
@@ -585,7 +584,7 @@ int realmain(int argc, char *argv[]) {
 
             auto packageFileRefs = pipeline::reserveFiles(gs, packageFiles);
             auto packages = pipeline::index(*gs, packageFileRefs, opts, *workers, nullptr);
-            packageNamespaces = packager::RBIGenerator::buildPackageNamespace(*gs, packages, *workers);
+            packages = packager::Packager::findPackages(*gs, *workers, move(packages));
 #endif
         }
 
@@ -737,6 +736,7 @@ int realmain(int argc, char *argv[]) {
             logger->warn("Package rbi generation is disabled in sorbet-orig for faster builds");
             return 1;
 #else
+            auto packageNamespaces = packager::RBIGenerator::buildPackageNamespace(*gs, *workers);
             packager::RBIGenerator::run(*gs, packageNamespaces, opts.packageRBIOutput, *workers);
 #endif
         }
diff --git a/packager/rbi_gen.cc b/packager/rbi_gen.cc
index aa6164248fa..51d2f150421 100644
--- a/packager/rbi_gen.cc
+++ b/packager/rbi_gen.cc
@@ -1215,11 +1215,7 @@ class RBIExporter final {
 };
 } // namespace
 
-UnorderedSet<core::ClassOrModuleRef>
-RBIGenerator::buildPackageNamespace(core::GlobalState &gs, vector<ast::ParsedFile> &packageFiles, WorkerPool &workers) {
-    // Populate package database.
-    packageFiles = Packager::findPackages(gs, workers, move(packageFiles));
-
+UnorderedSet<core::ClassOrModuleRef> RBIGenerator::buildPackageNamespace(core::GlobalState &gs, WorkerPool &workers) {
     const auto &packageDB = gs.packageDB();
 
     auto &packages = packageDB.packages();
diff --git a/packager/rbi_gen.h b/packager/rbi_gen.h
index b91716d2424..1839825bac5 100644
--- a/packager/rbi_gen.h
+++ b/packager/rbi_gen.h
@@ -19,8 +19,7 @@ class RBIGenerator final {
     };
 
     // Exposed for testing.
-    static UnorderedSet<core::ClassOrModuleRef>
-    buildPackageNamespace(core::GlobalState &gs, std::vector<ast::ParsedFile> &packageFiles, WorkerPool &workers);
+    static UnorderedSet<core::ClassOrModuleRef> buildPackageNamespace(core::GlobalState &gs, WorkerPool &workers);
     static RBIOutput runOnce(const core::GlobalState &gs, core::NameRef pkg,
                              const UnorderedSet<core::ClassOrModuleRef> &packageNamespaces);
 
diff --git a/test/pipeline_test_runner.cc b/test/pipeline_test_runner.cc
index 78ea87d3e5f..2dd684d3110 100644
--- a/test/pipeline_test_runner.cc
+++ b/test/pipeline_test_runner.cc
@@ -350,7 +350,7 @@ TEST_CASE("PerPhaseTest") { // NOLINT
                 tree = ast::ParsedFile{rewriter::Rewriter::run(ctx, move(tree.tree)), tree.file};
                 tree = testSerialize(*gs, local_vars::LocalVars::run(ctx, move(tree)));
 
-                if (FileOps::getFileName(tree.file.data(*rbiGenGs).path()) == packageFileName) {
+                if (tree.file.data(*rbiGenGs).isPackage()) {
                     packageTrees.emplace_back(move(tree));
                 } else {
                     trees.emplace_back(move(tree));
@@ -373,8 +373,8 @@ TEST_CASE("PerPhaseTest") { // NOLINT
 
             // RBI generation
             {
-                auto packageNamespaces =
-                    packager::RBIGenerator::buildPackageNamespace(*rbiGenGs, packageTrees, *workers);
+                packageTrees = packager::Packager::findPackages(*rbiGenGs, *workers, move(packageTrees));
+                auto packageNamespaces = packager::RBIGenerator::buildPackageNamespace(*rbiGenGs, *workers);
                 for (auto &package : rbiGenGs->packageDB().packages()) {
                     auto output = packager::RBIGenerator::runOnce(*rbiGenGs, package, packageNamespaces);
                     if (!output.rbi.empty()) {