Add package support to T::NonForcingConstants (sorbet#3470)

* Add helper method to mangle a package name in the same way as packager pass

* Update RBI to know about third optional arg to non_forcing_is_a

* Implement support for packaged non_forcing_is_a in resolver

* Add test case

* Update exps

* Add stubbed-out runtime implementation

* Handle `package` being a keyword arg

* Update exp files

* Also return if the package literal is not a string

* format_cxx

* Fix runtime behavior... for now

* Remove expectation test for now

* Change to lookupMangledPackageName

* Add test to exercise non-existent package and fix error message printing

* Just else

* Add a comment too

* Share the comment

* Add todo comment

core/NameRef.h

@@ -94,6 +94,8 @@ class NameRef final : private DebugOnlyCheck<NameRefDebugCheck> {
 
     NameRef prepend(GlobalState &gs, std::string_view s) const;
 
+    NameRef lookupMangledPackageName(const GlobalState &gs) const;
+
     std::string showRaw(const GlobalState &gs) const;
     std::string toString(const GlobalState &gs) const;
     std::string show(const GlobalState &gs) const;

core/Names.cc

@@ -7,6 +7,8 @@
 #include <numeric> // accumulate
 
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/str_split.h"
 
 using namespace std;
 
@@ -291,6 +293,14 @@ NameRef NameRef::prepend(GlobalState &gs, string_view s) const {
     return gs.enterNameUTF8(nameEq);
 }
 
+NameRef NameRef::lookupMangledPackageName(const GlobalState &gs) const {
+    auto name = this->data(gs);
+    ENFORCE(name->kind == NameKind::UTF8, "manglePackageName over non-utf8 name");
+    auto parts = absl::StrSplit(name->raw.utf8, "::");
+    string nameEq = absl::StrCat(absl::StrJoin(parts, "_"), "_Package");
+    return gs.lookupNameConstant(nameEq);
+}
+
 Name Name::deepCopy(const GlobalState &to) const {
     Name out;
     out.kind = this->kind;

gems/sorbet-runtime/lib/types/non_forcing_constants.rb

@@ -4,8 +4,11 @@
 module T::NonForcingConstants
   # NOTE: This method is documented on the RBI in Sorbet's payload, so that it
   # shows up in the hover/completion documentation via LSP.
-  T::Sig::WithoutRuntime.sig {params(val: BasicObject, klass: String).returns(T::Boolean)}
-  def self.non_forcing_is_a?(val, klass)
+  T::Sig::WithoutRuntime.sig {params(val: BasicObject, klass: String, package: T.nilable(String)).returns(T::Boolean)}
+  def self.non_forcing_is_a?(val, klass, package: nil)
+    # TODO(gdritter): once we have a runtime implementation of
+    # packages, we'll need to actually handle the `package` argument
+    # here.
     method_name = "T::NonForcingConstants.non_forcing_is_a?"
     if klass.empty?
       raise ArgumentError.new("The string given to `#{method_name}` must not be empty")
@@ -18,7 +21,10 @@ def self.non_forcing_is_a?(val, klass)
     parts.each do |part|
       if current_klass.nil?
         # First iteration
-        if part != ""
+        if part != "" && package.nil?
+          # if we've supplied a package, we're probably running in
+          # package mode, which means absolute references are
+          # meaningless
           raise ArgumentError.new("The string given to `#{method_name}` must be an absolute constant reference that starts with `::`")
         end
 

rbi/sorbet/t.rbi

@@ -186,6 +186,6 @@ end
 
 module T::NonForcingConstants
   # See <https://sorbet.org/docs/non-forcing-constants> for full docs.
-  sig {params(val: BasicObject, klass: String).returns(T::Boolean)}
-  def self.non_forcing_is_a?(val, klass); end
+  sig {params(val: BasicObject, klass: String, package: T.nilable(String)).returns(T::Boolean)}
+  def self.non_forcing_is_a?(val, klass, package: nil); end
 end

resolver/resolver.cc

@@ -1886,7 +1886,7 @@ class ResolveSignaturesWalk {
     void validateNonForcingIsA(core::Context ctx, const ast::Send &send) {
         constexpr string_view method = "T::NonForcingConstants.non_forcing_is_a?";
 
-        if (send.args.size() != 2) {
+        if (send.args.size() != 2 && send.args.size() != 3) {
             return;
         }
 
@@ -1913,6 +1913,38 @@ class ResolveSignaturesWalk {
             return;
         }
 
+        core::LiteralType *package = nullptr;
+        optional<core::LocOffsets> packageLoc;
+        if (send.args.size() == 3) {
+            // this means we got the third package arg
+            auto *kwargs = ast::cast_tree_const<ast::Hash>(send.args[2]);
+            if (!kwargs || kwargs->keys.size() != 1) {
+                // Infer will report an error
+                return;
+            }
+            auto *key = ast::cast_tree_const<ast::Literal>(kwargs->keys.front());
+            if (!key || !key->isSymbol(ctx) || key->asSymbol(ctx) != ctx.state.lookupNameUTF8("package")) {
+                return;
+            }
+
+            auto *packageNode = ast::cast_tree_const<ast::Literal>(kwargs->values.front());
+            packageLoc = std::optional<core::LocOffsets>{send.args[2]->loc};
+            if (packageNode == nullptr) {
+                if (auto e = ctx.beginError(send.args[2]->loc, core::errors::Resolver::LazyResolve)) {
+                    e.setHeader("`{}` only accepts string literals", method);
+                }
+                return;
+            }
+
+            package = core::cast_type<core::LiteralType>(packageNode->value.get());
+            if (package == nullptr || package->literalKind != core::LiteralType::LiteralTypeKind::String) {
+                // Infer will report a type error
+                return;
+            }
+        }
+        // if we got two args, then package should be null, and if we got three args, then package should be non-null
+        ENFORCE((send.args.size() == 2 && !package) || (send.args.size() == 3 && package));
+
         auto name = core::NameRef(ctx.state, literal->value);
         auto shortName = name.data(ctx)->shortName(ctx);
         if (shortName.empty()) {
@@ -1922,45 +1954,74 @@ class ResolveSignaturesWalk {
             return;
         }
 
+        // If this string _begins_ with `::`, then the first fragment will be an empty string; in multiple places below,
+        // we'll check to find out whether the first part is `""` or not, which means we're testing whether the string
+        // did or did not begin with `::`.
         auto parts = absl::StrSplit(shortName, "::");
         core::SymbolRef current;
         for (auto part : parts) {
             if (!current.exists()) {
                 // First iteration
-                if (part != "") {
-                    if (auto e = ctx.beginError(stringLoc, core::errors::Resolver::LazyResolve)) {
-                        e.setHeader(
-                            "The string given to `{}` must be an absolute constant reference that starts with `{}`",
-                            method, "::");
+                if (!package) {
+                    if (part != "") {
+                        if (auto e = ctx.beginError(stringLoc, core::errors::Resolver::LazyResolve)) {
+                            e.setHeader(
+                                "The string given to `{}` must be an absolute constant reference that starts with `{}`",
+                                method, "::");
+                        }
+                        return;
+                    }
+                    current = core::Symbols::root();
+                    continue;
+                } else {
+                    if (part == "") {
+                        if (auto e = ctx.beginError(stringLoc, core::errors::Resolver::LazyResolve)) {
+                            e.setHeader("The string given to `{}` should not be an absolute constant reference if a "
+                                        "package name is also provided",
+                                        method);
+                        }
+                        return;
                     }
-                    return;
-                }
 
-                current = core::Symbols::root();
-            } else {
-                auto member = ctx.state.lookupNameConstant(part);
-                if (!member.exists()) {
-                    if (auto e = ctx.beginError(stringLoc, core::errors::Resolver::LazyResolve)) {
-                        auto prettyCurrent =
-                            current == core::Symbols::root() ? "" : "::" + current.data(ctx)->show(ctx);
-                        auto pretty = fmt::format("{}::{}", prettyCurrent, part);
-                        e.setHeader("Unable to resolve constant `{}`", pretty);
+                    auto packageName = core::NameRef(ctx.state, package->value);
+                    auto mangledName = packageName.lookupMangledPackageName(ctx.state);
+                    // if the mangled name doesn't exist, then this means probably there's no package named this
+                    if (!mangledName.exists()) {
+                        if (auto e = ctx.beginError(*packageLoc, core::errors::Resolver::LazyResolve)) {
+                            e.setHeader("Unable to find package: `{}`", packageName.toString(ctx));
+                        }
+                        return;
+                    }
+                    current = core::Symbols::PackageRegistry().data(ctx)->findMember(ctx, mangledName);
+                    if (!current.exists()) {
+                        if (auto e = ctx.beginError(*packageLoc, core::errors::Resolver::LazyResolve)) {
+                            e.setHeader("Unable to find package `{}`", packageName.toString(ctx));
+                        }
+                        return;
                     }
-                    return;
                 }
+            }
 
-                auto newCurrent = current.data(ctx)->findMember(ctx, member);
-                if (!newCurrent.exists()) {
-                    if (auto e = ctx.beginError(stringLoc, core::errors::Resolver::LazyResolve)) {
-                        auto prettyCurrent =
-                            current == core::Symbols::root() ? "" : "::" + current.data(ctx)->show(ctx);
-                        auto pretty = fmt::format("{}::{}", prettyCurrent, part);
-                        e.setHeader("Unable to resolve constant `{}`", pretty);
-                    }
-                    return;
+            auto member = ctx.state.lookupNameConstant(part);
+            if (!member.exists()) {
+                if (auto e = ctx.beginError(stringLoc, core::errors::Resolver::LazyResolve)) {
+                    auto prettyCurrent = current == core::Symbols::root() ? "" : "::" + current.data(ctx)->show(ctx);
+                    auto pretty = fmt::format("{}::{}", prettyCurrent, part);
+                    e.setHeader("Unable to resolve constant `{}`", pretty);
+                }
+                return;
+            }
+
+            auto newCurrent = current.data(ctx)->findMember(ctx, member);
+            if (!newCurrent.exists()) {
+                if (auto e = ctx.beginError(stringLoc, core::errors::Resolver::LazyResolve)) {
+                    auto prettyCurrent = current == core::Symbols::root() ? "" : "::" + current.data(ctx)->show(ctx);
+                    auto pretty = fmt::format("{}::{}", prettyCurrent, part);
+                    e.setHeader("Unable to resolve constant `{}`", pretty);
                 }
-                current = newCurrent;
+                return;
             }
+            current = newCurrent;
         }
 
         ENFORCE(current.exists(), "Loop invariant violated");

test/testdata/packager/non_forcing_constants/bar/__package.rb

@@ -0,0 +1,9 @@
+# typed: strict
+# enable-packager: true
+
+class Project::Bar < PackageSpec
+  import Project::Foo
+
+  export Bar
+  export_methods BarMethods
+end

test/testdata/packager/non_forcing_constants/bar/bar.rb

@@ -0,0 +1,9 @@
+# typed: strict
+
+class Bar
+  extend T::Sig
+  sig {params(value: Integer).void}
+  def initialize(value)
+    @value = T.let(value, Integer)
+  end
+end

test/testdata/packager/non_forcing_constants/bar/bar_methods.rb

@@ -0,0 +1,17 @@
+# typed: strict
+
+module BarMethods
+  extend T::Sig
+
+  sig {returns(Project::Foo::Foo)}
+  def build_foo
+    # Construct an imported class.
+    Project::Foo::Foo.new(10)
+  end
+
+  sig {returns(T::Boolean)}
+  def check_bar
+    # Call an imported method.
+    Project::Foo.good_check_is_bar(Bar.new(5))
+  end
+end

test/testdata/packager/non_forcing_constants/foo/__package.rb

@@ -0,0 +1,6 @@
+# typed: strict
+
+class Project::Foo < PackageSpec
+  export Foo
+  export_methods FooMethods
+end

test/testdata/packager/non_forcing_constants/foo/foo.rb

@@ -0,0 +1,9 @@
+# typed: strict
+
+class Foo
+  extend T::Sig
+  sig {params(value: Integer).void}
+  def initialize(value)
+    @value = T.let(value, Integer)
+  end
+end

test/testdata/packager/non_forcing_constants/foo/foo_methods.rb

@@ -0,0 +1,77 @@
+# typed: strict
+
+module FooMethods
+  extend T::Sig
+
+  sig {params(arg: T.untyped).returns(T::Boolean)}
+  def global_check_is_bar(arg)
+    if T::NonForcingConstants.non_forcing_is_a?(arg, "::Bar") # error: Unable to resolve constant `::Bar`
+      true
+    else
+      false
+    end
+  end
+
+  sig {params(arg: T.untyped).returns(T::Boolean)}
+  def bad_not_keyword_arg(arg)
+    if T::NonForcingConstants.non_forcing_is_a?(arg, "::Bar", "Project::Bar") # error: Too many positional arguments
+      true
+    else
+      false
+    end
+  end
+
+  sig {params(arg: T.untyped).returns(T::Boolean)}
+  def bad_wrong_keyword_arg(arg)
+    if T::NonForcingConstants.non_forcing_is_a?(arg, "::Bar", pakige: "Project::Bar") # error: Unrecognized keyword argument
+      true
+    else
+      false
+    end
+  end
+
+  sig {params(arg: T.untyped).returns(T::Boolean)}
+  def bad_package_not_a_string(arg)
+    if T::NonForcingConstants.non_forcing_is_a?(arg, "::Bar", package: 5) # error: Expected `T.nilable(String)` but found `Integer(5)`
+      true
+    else
+      false
+    end
+  end
+
+  sig {params(arg: T.untyped).returns(T::Boolean)}
+  def bad_check_for_global_bar(arg)
+    if T::NonForcingConstants.non_forcing_is_a?(arg, "::Bar", package: "Project::Bar") # error: should not be an absolute constant reference
+      true
+    else
+      false
+    end
+  end
+
+  sig {params(arg: T.untyped).returns(T::Boolean)}
+  def bad_non_existent_package(arg)
+    if T::NonForcingConstants.non_forcing_is_a?(arg, "Bar", package: "Project::Quux") # error: Unable to find package
+      true
+    else
+      false
+    end
+  end
+
+  sig {params(arg: T.untyped).returns(T::Boolean)}
+  def check_for_non_existing_bar(arg)
+    if T::NonForcingConstants.non_forcing_is_a?(arg, "Quux", package: "Project::Bar") # error: Unable to resolve constant `::<PackageRegistry>::Project_Bar_Package::Quux`
+      true
+    else
+      false
+    end
+  end
+
+  sig {params(arg: T.untyped).returns(T::Boolean)}
+  def good_check_is_bar(arg)
+    if T::NonForcingConstants.non_forcing_is_a?(arg, "Bar", package: "Project::Bar")
+      true
+    else
+      false
+    end
+  end
+end