From ba3f461aebdbcb47c35b330fc9bc58aae91e6f4f Mon Sep 17 00:00:00 2001
From: Nathaniel Roman <ngroman@gmail.com>
Date: Thu, 2 Dec 2021 16:08:32 -0800
Subject: [PATCH] --stripe-packages: typecheck performance skip modules in
 package files (#4947)

* skip type-checking of package modules

* Add test case for method calls in packages

* fix removal + stop putting nodes back

* switch to earlier shallow map

* clear rhs in pre

* fixup

* allow for ast leaks in CLI mode to reduce dealloc traffic

* leak only
---
 main/pipeline/pipeline.cc                     |  6 ++++
 packager/packager.cc                          | 35 +++++++++++++++++++
 packager/packager.h                           |  6 ++++
 .../packager/packagespec_methods/__package.rb | 10 ++++++
 .../packagespec_methods/packagespec.rbi       |  6 ++++
 5 files changed, 63 insertions(+)
 create mode 100644 test/testdata/packager/packagespec_methods/__package.rb
 create mode 100644 test/testdata/packager/packagespec_methods/packagespec.rbi

diff --git a/main/pipeline/pipeline.cc b/main/pipeline/pipeline.cc
index 4b85aef93c4..04c3c19c429 100644
--- a/main/pipeline/pipeline.cc
+++ b/main/pipeline/pipeline.cc
@@ -544,6 +544,12 @@ void typecheckOne(core::Context ctx, ast::ParsedFile resolved, const options::Op
         return;
     }
 
+#ifndef SORBET_REALMAIN_MIN
+    if (f.data(ctx).isPackage()) {
+        resolved = packager::Packager::removePackageModules(ctx, move(resolved), intentionallyLeakASTs);
+    }
+#endif
+
     resolved = definition_validator::runOne(ctx, std::move(resolved));
 
     resolved = class_flatten::runOne(ctx, move(resolved));
diff --git a/packager/packager.cc b/packager/packager.cc
index e89daccabeb..2f4e0ccc28a 100644
--- a/packager/packager.cc
+++ b/packager/packager.cc
@@ -23,6 +23,34 @@ namespace {
 constexpr string_view PACKAGE_FILE_NAME = "__package.rb"sv;
 constexpr core::NameRef TEST_NAME = core::Names::Constants::Test();
 
+bool isPackageModule(const core::GlobalState &gs, core::SymbolRef modName) {
+    while (modName.exists() && modName != core::Symbols::root()) {
+        if (modName == core::Symbols::PackageRegistry() || modName == core::Symbols::PackageTests()) {
+            return true;
+        }
+        modName = modName.data(gs)->owner;
+    }
+    return false;
+}
+
+class PrunePackageModules final {
+    const bool intentionallyLeakASTs;
+
+public:
+    PrunePackageModules(bool intentionallyLeakASTs) : intentionallyLeakASTs(intentionallyLeakASTs) {}
+
+    ast::ExpressionPtr postTransformClassDef(core::Context ctx, ast::ExpressionPtr tree) {
+        auto &klass = ast::cast_tree_nonnull<ast::ClassDef>(tree);
+        if (isPackageModule(ctx, klass.symbol)) {
+            if (intentionallyLeakASTs) {
+                intentionallyLeakMemory(tree.release());
+            }
+            return ast::MK::EmptyTree();
+        }
+        return tree;
+    }
+};
+
 bool isPrimaryTestNamespace(const core::NameRef ns) {
     return ns == TEST_NAME;
 }
@@ -1475,4 +1503,11 @@ vector<ast::ParsedFile> Packager::runIncremental(core::GlobalState &gs, vector<a
     return files;
 }
 
+ast::ParsedFile Packager::removePackageModules(core::Context ctx, ast::ParsedFile pf, bool intentionallyLeakASTs) {
+    ENFORCE(pf.file.data(ctx).isPackage());
+    PrunePackageModules prune(intentionallyLeakASTs);
+    pf.tree = ast::ShallowMap::apply(ctx, prune, move(pf.tree));
+    return pf;
+}
+
 } // namespace sorbet::packager
diff --git a/packager/packager.h b/packager/packager.h
index 96a12d11fac..41bb88cb76b 100644
--- a/packager/packager.h
+++ b/packager/packager.h
@@ -82,6 +82,12 @@ class Packager final {
     // Run packager incrementally. Note: `files` must contain all packages files. Does not support package changes.
     static std::vector<ast::ParsedFile> runIncremental(core::GlobalState &gs, std::vector<ast::ParsedFile> files);
 
+    // The structures created for `__package.rb` files from their imports are large and deep. This causes
+    // performance problems with typechecking. Use to remove these modules while retaining the PackageSpec
+    // class itself during typecheck.
+    static ast::ParsedFile removePackageModules(core::Context ctx, ast::ParsedFile pf,
+                                                bool intentionallyLeakASTs = false);
+
     Packager() = delete;
 };
 } // namespace sorbet::packager
diff --git a/test/testdata/packager/packagespec_methods/__package.rb b/test/testdata/packager/packagespec_methods/__package.rb
new file mode 100644
index 00000000000..047916d8bb5
--- /dev/null
+++ b/test/testdata/packager/packagespec_methods/__package.rb
@@ -0,0 +1,10 @@
+# typed: strict
+
+class MyPkg < PackageSpec
+  custom_method 'abc'
+  custom_method 'abc', 'too_many_args'
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: Too many arguments provided for method `PackageSpec.custom_method`. Expected: `1`, got: `2`
+
+  bad_method 'def'
+# ^^^^^^^^^^^^^^^^ error: Method `bad_method` does not exist on `T.class_of(MyPkg)`
+end
diff --git a/test/testdata/packager/packagespec_methods/packagespec.rbi b/test/testdata/packager/packagespec_methods/packagespec.rbi
new file mode 100644
index 00000000000..2c5a2e00faf
--- /dev/null
+++ b/test/testdata/packager/packagespec_methods/packagespec.rbi
@@ -0,0 +1,6 @@
+# typed: strict
+
+class ::PackageSpec
+  sig {params(x: String).void}
+  def self.custom_method(x); end
+end