--gen-packages mode for imports (sorbet#9080)

* --gen-packages

* track package references

* combined autocorrect

* docs

* stop after resolving

* test

* switch to multiplexJob instead of Parallel::iterate (which uses multiplexJobWait)

* one more early continue and de-indent

* only run package reference tracking when the flag is enabled

* don't allow both --gen-packages and --lsp

* Speed up untracking package references

On the slow path, the package references sets will start off empty, so
looping over them should be fast. However, because calls to untrack and
track were interleaved, the untrack calls would encounter the entries
added by the track calls and do extra work that's not needed (since the
track calls would a file A, but the untrack call would be for some other
file B). By doing all the untracking upfront, we can make it faster.

* test for invalid CLI options

* fix non determinism in test

* keep global state param const and pass in nonConstPackageDB instead

* store unique_ptr instead of optional in the queue

* Revert "store unique_ptr instead of optional in the queue"

This reverts commit 02d2c57.

Author: neilparikh

core/GlobalState.cc

@@ -2073,9 +2073,8 @@ unique_ptr<GlobalState> GlobalState::copyForIndexThread(
     const vector<string> &extraPackageFilesDirectorySlashDeprecatedPrefixes,
     const vector<string> &extraPackageFilesDirectorySlashPrefixes,
     const vector<string> &packageSkipRBIExportEnforcementDirs, const vector<string> &allowRelaxedPackagerChecksFor,
-    const vector<string> &packagerLayers, string errorHint) const {
+    const vector<string> &packagerLayers, string errorHint, bool genPackages) const {
     ENFORCE(fileTableFrozen);
-
     auto result = make_unique<GlobalState>(this->errorQueue, this->epochManager);
 
     result->initEmpty();
@@ -2094,7 +2093,7 @@ unique_ptr<GlobalState> GlobalState::copyForIndexThread(
         result->setPackagerOptions(extraPackageFilesDirectoryUnderscorePrefixes,
                                    extraPackageFilesDirectorySlashDeprecatedPrefixes,
                                    extraPackageFilesDirectorySlashPrefixes, packageSkipRBIExportEnforcementDirs,
-                                   allowRelaxedPackagerChecksFor, packagerLayers, errorHint);
+                                   allowRelaxedPackagerChecksFor, packagerLayers, errorHint, genPackages);
     }
 
     return result;
@@ -2105,7 +2104,7 @@ unique_ptr<GlobalState> GlobalState::copyForLSPTypechecker(
     const vector<string> &extraPackageFilesDirectorySlashDeprecatedPrefixes,
     const vector<string> &extraPackageFilesDirectorySlashPrefixes,
     const vector<string> &packageSkipRBIExportEnforcementDirs, const vector<string> &allowRelaxedPackagerChecksFor,
-    const vector<string> &packagerLayers, string errorHint) const {
+    const vector<string> &packagerLayers, string errorHint, bool genPackages) const {
     auto result = make_unique<GlobalState>(this->errorQueue, this->epochManager);
 
     result->initEmpty();
@@ -2121,7 +2120,7 @@ unique_ptr<GlobalState> GlobalState::copyForLSPTypechecker(
         result->setPackagerOptions(extraPackageFilesDirectoryUnderscorePrefixes,
                                    extraPackageFilesDirectorySlashDeprecatedPrefixes,
                                    extraPackageFilesDirectorySlashPrefixes, packageSkipRBIExportEnforcementDirs,
-                                   allowRelaxedPackagerChecksFor, packagerLayers, errorHint);
+                                   allowRelaxedPackagerChecksFor, packagerLayers, errorHint, genPackages);
     }
 
     return result;
@@ -2132,7 +2131,7 @@ GlobalState::copyForSlowPath(const vector<string> &extraPackageFilesDirectoryUnd
                              const vector<string> &extraPackageFilesDirectorySlashPrefixes,
                              const vector<string> &packageSkipRBIExportEnforcementDirs,
                              const vector<string> &allowRelaxedPackagerChecksFor, const vector<string> &packagerLayers,
-                             string errorHint) const {
+                             string errorHint, bool genPackages) const {
     auto result = make_unique<GlobalState>(this->errorQueue, this->epochManager);
 
     // We omit a call to `initEmpty` here, as the only intended use of this function is to have its symbol table
@@ -2164,7 +2163,7 @@ GlobalState::copyForSlowPath(const vector<string> &extraPackageFilesDirectoryUnd
         result->setPackagerOptions(extraPackageFilesDirectoryUnderscorePrefixes,
                                    extraPackageFilesDirectorySlashDeprecatedPrefixes,
                                    extraPackageFilesDirectorySlashPrefixes, packageSkipRBIExportEnforcementDirs,
-                                   allowRelaxedPackagerChecksFor, packagerLayers, errorHint);
+                                   allowRelaxedPackagerChecksFor, packagerLayers, errorHint, genPackages);
     }
 
     return result;
@@ -2315,10 +2314,11 @@ void GlobalState::setPackagerOptions(const vector<string> &extraPackageFilesDire
                                      const vector<string> &extraPackageFilesDirectorySlashPrefixes,
                                      const vector<string> &packageSkipRBIExportEnforcementDirs,
                                      const vector<string> &allowRelaxedPackagerChecksFor,
-                                     const vector<string> &packagerLayers, string errorHint) {
+                                     const vector<string> &packagerLayers, string errorHint, bool genPackages) {
     ENFORCE_NO_TIMER(!packageDB_.frozen);
 
     packageDB_.enabled_ = true;
+    packageDB_.genPackages_ = genPackages;
     packageDB_.extraPackageFilesDirectoryUnderscorePrefixes_ = extraPackageFilesDirectoryUnderscorePrefixes;
     packageDB_.extraPackageFilesDirectorySlashDeprecatedPrefixes_ = extraPackageFilesDirectorySlashDeprecatedPrefixes;
     packageDB_.extraPackageFilesDirectorySlashPrefixes_ = extraPackageFilesDirectorySlashPrefixes;

core/GlobalState.h

@@ -362,7 +362,7 @@ class GlobalState final {
                             const std::vector<std::string> &extraPackageFilesDirectorySlashPrefixes,
                             const std::vector<std::string> &packageSkipRBIExportEnforcementDirs,
                             const std::vector<std::string> &skipImportVisibilityCheckFor,
-                            const std::vector<std::string> &packagerLayers, std::string errorHint);
+                            const std::vector<std::string> &packagerLayers, std::string errorHint, bool genPackages);
     packages::UnfreezePackages unfreezePackages();
 
     NameRef nextMangledName(ClassOrModuleRef owner, NameRef origName);
@@ -481,19 +481,18 @@ class GlobalState final {
                        const std::vector<std::string> &extraPackageFilesDirectorySlashPrefixes,
                        const std::vector<std::string> &packageSkipRBIExportEnforcementDirs,
                        const std::vector<std::string> &allowRelaxedPackagerChecksFor,
-                       const std::vector<std::string> &packagerLayers, std::string errorHint) const;
+                       const std::vector<std::string> &packagerLayers, std::string errorHint, bool genPackages) const;
 
     // Minimally copy the global state, including the file table, to initialize the LSPTypechecker.
     // NOTE: this very intentionally will not copy the symbol or name tables. The symbol tables aren't used or populated
     // during indexing, and the name tables will only be written to.
-    std::unique_ptr<GlobalState>
-    copyForLSPTypechecker(const bool packagerEnabled,
-                          const std::vector<std::string> &extraPackageFilesDirectoryUnderscorePrefixes,
-                          const std::vector<std::string> &extraPackageFilesDirectorySlashDeprecatedPrefixes,
-                          const std::vector<std::string> &extraPackageFilesDirectorySlashPrefixes,
-                          const std::vector<std::string> &packageSkipRBIExportEnforcementDirs,
-                          const std::vector<std::string> &allowRelaxedPackagerChecksFor,
-                          const std::vector<std::string> &packagerLayers, std::string errorHint) const;
+    std::unique_ptr<GlobalState> copyForLSPTypechecker(
+        const bool packagerEnabled, const std::vector<std::string> &extraPackageFilesDirectoryUnderscorePrefixes,
+        const std::vector<std::string> &extraPackageFilesDirectorySlashDeprecatedPrefixes,
+        const std::vector<std::string> &extraPackageFilesDirectorySlashPrefixes,
+        const std::vector<std::string> &packageSkipRBIExportEnforcementDirs,
+        const std::vector<std::string> &allowRelaxedPackagerChecksFor, const std::vector<std::string> &packagerLayers,
+        std::string errorHint, bool genPackages) const;
 
     // Copy the name table, file table and other parts of GlobalState that are required to start the slow path.
     // NOTE: this very intentionally will not copy the symbol table, and the expectation is that the symbol table will
@@ -504,7 +503,7 @@ class GlobalState final {
                     const std::vector<std::string> &extraPackageFilesDirectorySlashPrefixes,
                     const std::vector<std::string> &packageSkipRBIExportEnforcementDirs,
                     const std::vector<std::string> &allowRelaxedPackagerChecksFor,
-                    const std::vector<std::string> &packagerLayers, std::string errorHint) const;
+                    const std::vector<std::string> &packagerLayers, std::string errorHint, bool genPackages) const;
 
     // Contains a path prefix that should be stripped from all printed paths.
     std::string pathPrefix;

core/errors/packager.h

@@ -38,5 +38,6 @@ constexpr ErrorClass InvalidMinTypedLevel{3729, StrictLevel::False};
 constexpr ErrorClass PreludePackageImport{3731, StrictLevel::False};
 // constexpr ErrorClass NoExplicitPreludeImport{3732, StrictLevel::False};
 // constexpr ErrorClass PreludeLowestLayer{3733, StrictLevel::False};
+constexpr ErrorClass IncorrectImportList{3734, StrictLevel::False};
 } // namespace sorbet::core::errors::Packager
 #endif

core/packages/PackageDB.cc

@@ -172,6 +172,7 @@ PackageDB PackageDB::deepCopy() const {
 
     // --- options ---
     result.enabled_ = this->enabled_;
+    result.genPackages_ = this->genPackages_;
     result.extraPackageFilesDirectoryUnderscorePrefixes_ = this->extraPackageFilesDirectoryUnderscorePrefixes_;
     result.extraPackageFilesDirectorySlashDeprecatedPrefixes_ =
         this->extraPackageFilesDirectorySlashDeprecatedPrefixes_;

core/packages/PackageDB.h

@@ -60,6 +60,11 @@ class PackageDB final {
         return this->enabled_;
     }
 
+    // Whether the --gen-packages mode is active.
+    bool genPackages() const {
+        return this->genPackages_;
+    }
+
     absl::Span<const std::string> extraPackageFilesDirectoryUnderscorePrefixes() const;
     absl::Span<const std::string> extraPackageFilesDirectorySlashDeprecatedPrefixes() const;
     absl::Span<const std::string> extraPackageFilesDirectorySlashPrefixes() const;
@@ -92,6 +97,7 @@ class PackageDB final {
 
 private:
     bool enabled_ = false;
+    bool genPackages_ = false;
     std::vector<std::string> extraPackageFilesDirectoryUnderscorePrefixes_;
     std::vector<std::string> extraPackageFilesDirectorySlashDeprecatedPrefixes_;
     std::vector<std::string> extraPackageFilesDirectorySlashPrefixes_;

core/packages/PackageInfo.cc

@@ -1,5 +1,6 @@
 #include "core/packages/PackageInfo.h"
 #include "absl/strings/str_join.h"
+#include "common/sort/sort.h"
 #include "core/GlobalState.h"
 #include "core/Loc.h"
 #include "core/NameRef.h"
@@ -466,6 +467,97 @@ optional<string> PackageInfo::pathTo(const core::GlobalState &gs, const MangledN
     return nullopt;
 }
 
+void PackageInfo::trackPackageReference(const core::FileRef file, const core::packages::MangledName package,
+                                        const PackageReferenceInfo packageReferenceInfo) {
+    auto r = referencedPackages.find(package);
+    if (r != referencedPackages.end()) {
+        r->second.first.insert(file);
+    } else {
+        referencedPackages[package] = {{file}, packageReferenceInfo};
+    }
+}
+
+void PackageInfo::untrackPackageReferencesFor(const core::FileRef file) {
+    for (auto &[_, v] : referencedPackages) {
+        auto files = v.first;
+        auto it = files.find(file);
+        if (it == files.end()) {
+            continue;
+        }
+        files.erase(it);
+    }
+    erase_if(referencedPackages, [](const auto &x) { return x.second.first.empty(); });
+}
+
+namespace {
+core::packages::ImportType fileToImportType(const core::GlobalState &gs, core::FileRef file) {
+    if (file.data(gs).isPackagedTestHelper()) {
+        return core::packages::ImportType::TestHelper;
+    } else if (file.data(gs).isPackagedTest()) {
+        return core::packages::ImportType::TestUnit;
+    } else {
+        return core::packages::ImportType::Normal;
+    }
+}
+
+core::packages::ImportType broadestImportType(const core::GlobalState &gs, const UnorderedSet<core::FileRef> &files) {
+    auto broadestImport = core::packages::ImportType::TestUnit;
+    for (auto f : files) {
+        auto importType = fileToImportType(gs, f);
+        if (importType < broadestImport) {
+            broadestImport = importType;
+        }
+    }
+    return broadestImport;
+}
+
+void mergeAdjacentEdits(std::vector<core::AutocorrectSuggestion::Edit> &edits) {
+    fast_sort(edits, [](const auto &lhs, const auto &rhs) {
+        if (lhs.loc == rhs.loc) {
+            return lhs.replacement < rhs.replacement;
+        }
+        return lhs.loc.beginPos() < rhs.loc.beginPos();
+    });
+    auto i = 0;
+    while (edits.size() > 0 && i < edits.size() - 1) {
+        if (edits[i].loc.beginPos() == edits[i + 1].loc.beginPos() && edits[i].loc.empty() &&
+            edits[i + 1].loc.empty()) {
+            // If we're inserting 2 imports at the same location, combine them into a single edit.
+            // TODO(neil): maybe this logic should be moved/added to AutocorrectSuggestion::apply, to handle other
+            // cases where 2 autocorrects add at the same loc?
+            edits[i].replacement += edits[i + 1].replacement;
+            edits.erase(edits.begin() + i + 1);
+        } else {
+            i++;
+        }
+    }
+}
+
+}; // namespace
+
+std::optional<core::AutocorrectSuggestion> PackageInfo::aggregateMissingImports(const core::GlobalState &gs) const {
+    std::vector<core::AutocorrectSuggestion::Edit> allEdits;
+    for (auto &[p, value] : referencedPackages) {
+        auto &pkgInfo = gs.packageDB().getPackageInfo(p);
+        auto files = value.first;
+        auto packageReferenceInfo = value.second;
+        if (!packageReferenceInfo.importNeeded || packageReferenceInfo.causesModularityError || !pkgInfo.exists()) {
+            continue;
+        }
+        auto importType = broadestImportType(gs, files);
+        auto autocorrect = this->addImport(gs, pkgInfo, importType);
+        if (autocorrect.has_value()) {
+            allEdits.insert(allEdits.end(), make_move_iterator(autocorrect.value().edits.begin()),
+                            make_move_iterator(autocorrect.value().edits.end()));
+        }
+    }
+    if (allEdits.size() == 0) {
+        return nullopt;
+    }
+    mergeAdjacentEdits(allEdits);
+    return core::AutocorrectSuggestion{"Add missing imports", std::move(allEdits)};
+}
+
 bool PackageInfo::operator==(const PackageInfo &rhs) const {
     return mangledName() == rhs.mangledName();
 }

core/packages/PackageInfo.h

@@ -83,6 +83,11 @@ struct VisibleTo {
         : mangledName(mangledName), type(type), loc(loc){};
 };
 
+struct PackageReferenceInfo {
+    bool importNeeded;
+    bool causesModularityError;
+};
+
 class PackageInfo {
 public:
     MangledName mangledName() const {
@@ -177,9 +182,14 @@ class PackageInfo {
 
     StrictDependenciesLevel strictDependenciesLevel = StrictDependenciesLevel::None;
 
+    // Map from pkg -> {list of files in this package that reference pkg, whether this package is missing an import
+    // for pkg and whether importing this package would cause a modularity error}
+    UnorderedMap<core::packages::MangledName, std::pair<UnorderedSet<core::FileRef>, PackageReferenceInfo>>
+        referencedPackages = {};
+
     // The id of the SCC that this package's normal imports belong to.
     //
-    // WARNING: Modifying the contents of the package DB after this operation will cause this id to go out of
+    // WARNING: Modifying the contents of the package DB after ComputePackageSCCs will cause this id to go out of
     // date.
     std::optional<int> sccID() const {
         ENFORCE(exists());
@@ -189,7 +199,7 @@ class PackageInfo {
     // The ID of the SCC that this package's tests belong to. This ID is only useful in the context of the package graph
     // condensation graph.
     //
-    // WARNING: Modifying the contents of the package DB after this operation will cause this id to go out of
+    // WARNING: Modifying the contents of the package DB after ComputePackageSCCs will cause this id to go out of
     // date.
     std::optional<int> testSccID() const {
         ENFORCE(exists());
@@ -249,8 +259,18 @@ class PackageInfo {
     bool isPreludePackage() const {
         return this->isPreludePackage_;
     }
+
+    // Track that this package references `package` in `file`
+    void trackPackageReference(const core::FileRef file, const core::packages::MangledName package,
+                               const PackageReferenceInfo packageReferenceInfo);
+
+    // Remove knowledge of what this package is in `file`.
+    // We do this so that when VisibilityChecker is re-run over `file`, we can delete stale information.
+    void untrackPackageReferencesFor(const core::FileRef file);
+
+    std::optional<core::AutocorrectSuggestion> aggregateMissingImports(const core::GlobalState &gs) const;
 };
-CheckSize(PackageInfo, 208, 8);
+CheckSize(PackageInfo, 240, 8);
 
 } // namespace sorbet::core::packages
 #endif

main/lsp/LSPIndexer.cc

@@ -271,18 +271,19 @@ TypecheckingPath LSPIndexer::getTypecheckingPath(const vector<shared_ptr<core::F
 }
 
 void LSPIndexer::transferInitializeState(InitializedTask &task) {
+    ENFORCE(!this->config->opts.genPackages);
     // Copying the global state here means that we snapshot before any files have been loaded. That means that the
     // indexer and typechecker's file tables will almost immediately diverge, but that's not an issue as we don't share
     // `core::FileRef` values between the two.
-    auto typecheckerGS = std::exchange(
-        this->gs,
-        this->gs->copyForLSPTypechecker(this->config->opts.cacheSensitiveOptions.sorbetPackages,
-                                        this->config->opts.extraPackageFilesDirectoryUnderscorePrefixes,
-                                        this->config->opts.extraPackageFilesDirectorySlashDeprecatedPrefixes,
-                                        this->config->opts.extraPackageFilesDirectorySlashPrefixes,
-                                        this->config->opts.packageSkipRBIExportEnforcementDirs,
-                                        this->config->opts.allowRelaxedPackagerChecksFor,
-                                        this->config->opts.packagerLayers, this->config->opts.sorbetPackagesHint));
+    auto typecheckerGS =
+        std::exchange(this->gs, this->gs->copyForLSPTypechecker(
+                                    this->config->opts.cacheSensitiveOptions.sorbetPackages,
+                                    this->config->opts.extraPackageFilesDirectoryUnderscorePrefixes,
+                                    this->config->opts.extraPackageFilesDirectorySlashDeprecatedPrefixes,
+                                    this->config->opts.extraPackageFilesDirectorySlashPrefixes,
+                                    this->config->opts.packageSkipRBIExportEnforcementDirs,
+                                    this->config->opts.allowRelaxedPackagerChecksFor, this->config->opts.packagerLayers,
+                                    this->config->opts.sorbetPackagesHint, false));
 
     task.setGlobalState(std::move(typecheckerGS));
     task.setKeyValueStore(std::move(this->kvstore));

main/options/options.cc

@@ -425,6 +425,7 @@ buildOptions(const vector<pipeline::semantic_extension::SemanticExtensionProvide
                                  "Remove the provided <prefix> from all printed paths. Defaults to the input "
                                  "directory passed to Sorbet, if any.",
                                  cxxopts::value<string>()->default_value(empty.pathPrefix), "<prefix>");
+    options.add_options(section)("gen-packages", "Generate package information", cxxopts::value<bool>());
     // }}}
 
     // ----- AUTOCORRECTS ------------------------------------------------- {{{
@@ -1273,6 +1274,15 @@ void readOptions(Options &opts,
             }
         }
 
+        opts.genPackages = raw["gen-packages"].as<bool>();
+        if (opts.genPackages && !opts.cacheSensitiveOptions.sorbetPackages) {
+            logger->error("--gen-packages can only be used when --stripe-packages is also enabled");
+            throw EarlyReturnWithCode(1);
+        }
+        if (opts.genPackages && opts.runLSP) {
+            logger->error("--gen-packages can not be used when --lsp is also enabled");
+            throw EarlyReturnWithCode(1);
+        }
         if (raw.count("allow-relaxed-packager-checks-for")) {
             if (!opts.cacheSensitiveOptions.sorbetPackages) {
                 logger->error("--allow-relaxed-packager-checks-for can only be specified in --sorbet-packages mode");