--stripe-packages: typecheck performance skip modules in package files (

sorbet#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

@@ -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));

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

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

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

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