Package state in GlobalState (#4631)

* Move PackageDB to GlobalState * Avoid re-processing all package files when package specs have not changed * remove TODOs * don't use unique_ptr in PackageDB::deepCopy * Apply suggestions from code review Co-authored-by: Jake Zimmerman <zimmerman.jake@gmail.com> Co-authored-by: Jake Zimmerman <zimmerman.jake@gmail.com>
sorbet · Oct 12, 2021 · 0f72b1c · 0f72b1c
1 parent a33dd2e
commit 0f72b1c
Show file tree

Hide file tree

Showing 10 changed files with 399 additions and 183 deletions.
diff --git a/core/BUILD b/core/BUILD
@@ -6,6 +6,7 @@ cc_library(
             "*.h",
             "types/*.cc",
             "lsp/*.cc",
+            "packages/*.cc",
         ],
         exclude = [
             # workaround https://github.com/flycheck/flycheck/issues/248 in emacs
@@ -18,7 +19,11 @@ cc_library(
     hdrs = [
         "GlobalSubstitution.h",
         "core.h",
-    ] + glob(["errors/*.h"]) + glob(["lsp/*.h"]),
+    ] + glob([
+        "errors/*.h",
+        "lsp/*.h",
+        "packages/*.h",
+    ]),
     linkstatic = select({
         "//tools/config:linkshared": 0,
         "//conditions:default": 1,

diff --git a/core/GlobalState.cc b/core/GlobalState.cc
@@ -1767,6 +1767,7 @@ unique_ptr<GlobalState> GlobalState::deepCopy(bool keepId) const {
     for (auto &semanticExtension : this->semanticExtensions) {
         result->semanticExtensions.emplace_back(semanticExtension->deepCopy(*this, *result));
     }
+    result->packageDB_ = packageDB_.deepCopy();
     result->sanityCheck();
     {
         Timer timeit2(tracer(), "GlobalState::deepCopyOut");
@@ -1903,6 +1904,14 @@ FileRef GlobalState::findFileByPath(string_view path) const {
     return FileRef();
 }
 
+const packages::PackageDB &GlobalState::packageDB() const {
+    return packageDB_;
+}
+
+packages::UnfreezePackages GlobalState::unfreezePackages() {
+    return packageDB_.unfreeze();
+}
+
 unique_ptr<GlobalState> GlobalState::markFileAsTombStone(unique_ptr<GlobalState> what, FileRef fref) {
     ENFORCE(fref.id() < what->filesUsed());
     what->files[fref.id()]->sourceType = File::Type::TombStone;

diff --git a/core/GlobalState.h b/core/GlobalState.h
@@ -8,6 +8,8 @@
 #include "core/Names.h"
 #include "core/Symbols.h"
 #include "core/lsp/Query.h"
+#include "core/packages/PackageDB.h"
+#include "core/packages/PackageInfo.h"
 #include "main/pipeline/semantic_extension/SemanticExtension.h"
 #include <memory>
 
@@ -150,6 +152,9 @@ class GlobalState final {
     static std::unique_ptr<GlobalState> markFileAsTombStone(std::unique_ptr<GlobalState>, FileRef fref);
     FileRef findFileByPath(std::string_view path) const;
 
+    const packages::PackageDB &packageDB() const;
+    packages::UnfreezePackages unfreezePackages();
+
     void mangleRenameSymbol(SymbolRef what, NameRef origName);
     spdlog::logger &tracer() const;
     unsigned int namesUsedTotal() const;
@@ -297,6 +302,8 @@ class GlobalState final {
     UnorderedSet<int> onlyErrorClasses;
     bool wasModified_ = false;
 
+    core::packages::PackageDB packageDB_;
+
     bool freezeSymbolTable();
     bool freezeNameTable();
     bool freezeFileTable();

diff --git a/core/packages/PackageDB.cc b/core/packages/PackageDB.cc
@@ -0,0 +1,127 @@
+#include "core/packages/PackageDB.h"
+#include "core/Loc.h"
+
+using namespace std;
+
+namespace sorbet::core::packages {
+
+namespace {
+class NonePackage final : public PackageInfo {
+public:
+    core::NameRef mangledName() const {
+        return NameRef::noName();
+    }
+
+    const vector<string> &pathPrefixes() const {
+        ENFORCE(false);
+        return prefixes;
+    }
+
+    unique_ptr<PackageInfo> deepCopy() const {
+        ENFORCE(false);
+        return make_unique<NonePackage>();
+    }
+
+    Loc definitionLoc() const {
+        ENFORCE(false);
+        return Loc::none();
+    }
+
+    ~NonePackage() {}
+
+private:
+    const vector<string> prefixes;
+};
+static const NonePackage NONE_PKG;
+} // namespace
+
+UnfreezePackages::UnfreezePackages(PackageDB &db) : db(db) {
+    ENFORCE(db.frozen);
+    db.writerThread = this_thread::get_id();
+    db.frozen = false;
+}
+
+UnfreezePackages::~UnfreezePackages() {
+    ENFORCE(!db.frozen);
+    db.writerThread = std::thread::id();
+    db.frozen = true;
+    // Note in the future we may want to change the data structures involved in a way that requires
+    // a finalization step. Controlling freeze/unfreeze with RAII gives a good hook to do that.
+}
+
+NameRef PackageDB::enterPackage(unique_ptr<PackageInfo> pkg) {
+    ENFORCE(!frozen);
+    ENFORCE(writerThread == this_thread::get_id(), "PackageDB writes are not thread safe");
+    auto nr = pkg->mangledName();
+    auto prev = packages.find(nr);
+    if (prev == packages.end()) {
+        for (const auto &prefix : pkg->pathPrefixes()) {
+            packagesByPathPrefix[prefix] = nr;
+        }
+    } else {
+        // Package files do not have full featured content hashing. If the contents of one changes
+        // we always run slow-path and fully rebuild the set of packages. In some cases, the LSP
+        // fast-path may re-run on an unchanged package file. Sanity check to ensure the loc and
+        // prefixes are the same.
+        ENFORCE(prev->second->definitionLoc() == pkg->definitionLoc());
+        ENFORCE(prev->second->pathPrefixes() == pkg->pathPrefixes());
+    }
+    packages[nr] = move(pkg);
+    return nr;
+}
+
+NameRef PackageDB::lookupPackage(NameRef pkgMangledName) const {
+    ENFORCE(pkgMangledName.exists());
+    auto it = packages.find(pkgMangledName);
+    if (it == packages.end()) {
+        return NameRef::noName();
+    }
+    return it->first;
+}
+
+const PackageInfo &PackageDB::getPackageForFile(const core::GlobalState &gs, core::FileRef file) const {
+    ENFORCE(frozen);
+    auto &fileData = file.data(gs);
+    string_view path = fileData.path();
+    int curPrefixPos = path.find_last_of('/');
+    while (curPrefixPos != string::npos) {
+        const auto &it = packagesByPathPrefix.find(path.substr(0, curPrefixPos + 1));
+        if (it != packagesByPathPrefix.end()) {
+            const auto &pkg = getPackageInfo(it->second);
+            ENFORCE(pkg.exists());
+            return pkg;
+        }
+
+        if (fileData.sourceType == core::File::Type::Package) {
+            // When looking up a `__package.rb` file do not search parent directories
+            break;
+        }
+        curPrefixPos = path.find_last_of('/', curPrefixPos - 1);
+    }
+    return NONE_PKG;
+}
+
+const PackageInfo &PackageDB::getPackageInfo(core::NameRef mangledName) const {
+    auto it = packages.find(mangledName);
+    if (it == packages.end()) {
+        return NONE_PKG;
+    }
+    return *it->second;
+}
+
+PackageDB PackageDB::deepCopy() const {
+    ENFORCE(frozen);
+    PackageDB result;
+    result.packages.reserve(this->packages.size());
+    for (auto const &[nr, pkgInfo] : this->packages) {
+        result.packages[nr] = pkgInfo->deepCopy();
+    }
+    result.packagesByPathPrefix = this->packagesByPathPrefix;
+    return result;
+}
+
+UnfreezePackages PackageDB::unfreeze() {
+    return UnfreezePackages(*this);
+}
+
+} // namespace sorbet::core::packages
diff --git a/core/packages/PackageDB.h b/core/packages/PackageDB.h
@@ -0,0 +1,49 @@
+#ifndef SORBET_CORE_PACKAGES_PACKAGEDB_H
+#define SORBET_CORE_PACKAGES_PACKAGEDB_H
+
+#include "common/common.h"
+#include "core/Files.h"
+#include "core/Names.h"
+#include "core/packages/PackageInfo.h"
+
+namespace sorbet::core::packages {
+
+class PackageDB;
+
+class UnfreezePackages final {
+public:
+    PackageDB &db;
+
+    UnfreezePackages(PackageDB &);
+    ~UnfreezePackages();
+};
+
+class PackageDB final {
+public:
+    NameRef enterPackage(std::unique_ptr<PackageInfo> pkg);
+    NameRef lookupPackage(NameRef pkgMangledName) const;
+
+    const PackageInfo &getPackageForFile(const core::GlobalState &gs, core::FileRef file) const;
+    const PackageInfo &getPackageInfo(core::NameRef mangledName) const;
+
+    PackageDB deepCopy() const;
+
+    UnfreezePackages unfreeze();
+
+    PackageDB() = default;
+    PackageDB(const PackageDB &) = delete;
+    PackageDB(PackageDB &&) = default;
+    PackageDB &operator=(const PackageDB &) = delete;
+    PackageDB &operator=(PackageDB &&) = default;
+
+private:
+    UnorderedMap<core::NameRef, std::unique_ptr<packages::PackageInfo>> packages;
+    UnorderedMap<std::string, core::NameRef> packagesByPathPrefix;
+    bool frozen = true;
+    std::thread::id writerThread;
+
+    friend class UnfreezePackages;
+};
+
+} // namespace sorbet::core::packages
+#endif // SORBET_CORE_PACKAGES_PACKAGEDB_H
diff --git a/core/packages/PackageInfo.cc b/core/packages/PackageInfo.cc
@@ -0,0 +1,15 @@
+#include "core/packages/PackageInfo.h"
+#include "core/Loc.h"
+#include "core/NameRef.h"
+
+using namespace std;
+
+namespace sorbet::core::packages {
+bool PackageInfo::exists() const {
+    return mangledName().exists();
+}
+
+PackageInfo::~PackageInfo() {
+    // see https://eli.thegreenplace.net/2010/11/13/pure-virtual-destructors-in-c
+}
+} // namespace sorbet::core::packages
diff --git a/core/packages/PackageInfo.h b/core/packages/PackageInfo.h
@@ -0,0 +1,28 @@
+#ifndef SORBET_CORE_PACKAGES_PACKAGEINFO_H
+#define SORBET_CORE_PACKAGES_PACKAGEINFO_H
+
+#include <vector>
+
+namespace sorbet::core {
+class NameRef;
+class Loc;
+} // namespace sorbet::core
+
+namespace sorbet::core::packages {
+class PackageInfo {
+public:
+    virtual core::NameRef mangledName() const = 0;
+    virtual const std::vector<std::string> &pathPrefixes() const = 0;
+    virtual std::unique_ptr<PackageInfo> deepCopy() const = 0;
+    virtual core::Loc definitionLoc() const = 0;
+    virtual bool exists() const final;
+
+    virtual ~PackageInfo() = 0;
+    PackageInfo() = default;
+    PackageInfo(PackageInfo &) = delete;
+    explicit PackageInfo(const PackageInfo &) = default;
+    PackageInfo &operator=(PackageInfo &&) = delete;
+    PackageInfo &operator=(const PackageInfo &) = delete;
+};
+} // namespace sorbet::core::packages
+#endif
diff --git a/main/lsp/LSPTypechecker.cc b/main/lsp/LSPTypechecker.cc
@@ -582,10 +582,12 @@ vector<ast::ParsedFile> LSPTypechecker::getResolved(const vector<core::FileRef>
     ENFORCE(this_thread::get_id() == typecheckerThreadId, "Typechecker can only be used from the typechecker thread.");
     vector<ast::ParsedFile> updatedIndexed;
 
+    bool addAllPackages = false;
     for (auto fref : frefs) {
         if (fref.data(*gs).sourceType == core::File::Type::Package) {
             // Will be added in second loop.
             ENFORCE(config->opts.stripePackages);
+            addAllPackages = true;
             continue;
         }
 
@@ -595,8 +597,9 @@ vector<ast::ParsedFile> LSPTypechecker::getResolved(const vector<core::FileRef>
         }
     }
 
-    if (config->opts.stripePackages) {
-        // We must include every package file to resolve these files properly.
+    if (addAllPackages) {
+        ENFORCE(config->opts.stripePackages);
+        // We must include every package file to rebuild the full PackageDB.
         for (u4 i = 1; i < gs->filesUsed(); i++) {
             core::FileRef fref(i);
             if (fref.data(*gs).sourceType == core::File::Type::Package) {