[ruby/prism] Prism::StringQuery

Introduce StringQuery to provide methods to access some metadata
about the Ruby lexer.

https://github.com/ruby/prism/commit/d3f55b67b9

4 files changed, 77 insertions, 4 deletions

diff --git a/lib/prism.rb b/lib/prism.rb
index 66a64e7fd0..50b14a5486 100644
--- a/lib/prism.rb
+++ b/lib/prism.rb
@@ -25,6 +25,7 @@ module Prism
   autoload :Pattern, "prism/pattern"
   autoload :Reflection, "prism/reflection"
   autoload :Serialize, "prism/serialize"
+  autoload :StringQuery, "prism/string_query"
   autoload :Translation, "prism/translation"
   autoload :Visitor, "prism/visitor"
 
@@ -75,13 +76,13 @@ require_relative "prism/parse_result"
 # it's going to require the built library. Otherwise, it's going to require a
 # module that uses FFI to call into the library.
 if RUBY_ENGINE == "ruby" and !ENV["PRISM_FFI_BACKEND"]
-  require "prism/prism"
-
   # The C extension is the default backend on CRuby.
   Prism::BACKEND = :CEXT
-else
-  require_relative "prism/ffi"
 
+  require "prism/prism"
+else
   # The FFI backend is used on other Ruby implementations.
   Prism::BACKEND = :FFI
+
+  require_relative "prism/ffi"
 end
diff --git a/lib/prism/ffi.rb b/lib/prism/ffi.rb
index 0520f7cdd2..5caae440f4 100644
--- a/lib/prism/ffi.rb
+++ b/lib/prism/ffi.rb
@@ -73,6 +73,7 @@ module Prism
 
     callback :pm_parse_stream_fgets_t, [:pointer, :int, :pointer], :pointer
     enum :pm_string_init_result_t, %i[PM_STRING_INIT_SUCCESS PM_STRING_INIT_ERROR_GENERIC PM_STRING_INIT_ERROR_DIRECTORY]
+    enum :pm_string_query_t, [:PM_STRING_QUERY_ERROR, -1, :PM_STRING_QUERY_FALSE, :PM_STRING_QUERY_TRUE]
 
     load_exported_functions_from(
       "prism.h",
@@ -83,6 +84,9 @@ module Prism
       "pm_serialize_lex",
       "pm_serialize_parse_lex",
       "pm_parse_success_p",
+      "pm_string_query_local",
+      "pm_string_query_constant",
+      "pm_string_query_method_name",
       [:pm_parse_stream_fgets_t]
     )
 
@@ -492,4 +496,39 @@ module Prism
       values.pack(template)
     end
   end
+
+  # Here we are going to patch StringQuery to put in the class-level methods so
+  # that it can maintain a consistent interface
+  class StringQuery
+    class << self
+      # Mirrors the C extension's StringQuery::local? method.
+      def local?(string)
+        query(LibRubyParser.pm_string_query_local(string, string.bytesize, string.encoding.name))
+      end
+
+      # Mirrors the C extension's StringQuery::constant? method.
+      def constant?(string)
+        query(LibRubyParser.pm_string_query_constant(string, string.bytesize, string.encoding.name))
+      end
+
+      # Mirrors the C extension's StringQuery::method_name? method.
+      def method_name?(string)
+        query(LibRubyParser.pm_string_query_method_name(string, string.bytesize, string.encoding.name))
+      end
+
+      private
+
+      # Parse the enum result and return an appropriate boolean.
+      def query(result)
+        case result
+        when :PM_STRING_QUERY_ERROR
+          raise ArgumentError, "Invalid or non ascii-compatible encoding"
+        when :PM_STRING_QUERY_FALSE
+          false
+        when :PM_STRING_QUERY_TRUE
+          true
+        end
+      end
+    end
+  end
 end
diff --git a/lib/prism/prism.gemspec b/lib/prism/prism.gemspec
index 37aa979576..6123b71fc8 100644
--- a/lib/prism/prism.gemspec
+++ b/lib/prism/prism.gemspec
@@ -89,6 +89,7 @@ Gem::Specification.new do |spec|
     "lib/prism/polyfill/unpack1.rb",
     "lib/prism/reflection.rb",
     "lib/prism/serialize.rb",
+    "lib/prism/string_query.rb",
     "lib/prism/translation.rb",
     "lib/prism/translation/parser.rb",
     "lib/prism/translation/parser33.rb",
@@ -109,6 +110,7 @@ Gem::Specification.new do |spec|
     "rbi/prism/node.rbi",
     "rbi/prism/parse_result.rbi",
     "rbi/prism/reflection.rbi",
+    "rbi/prism/string_query.rbi",
     "rbi/prism/translation/parser.rbi",
     "rbi/prism/translation/parser33.rbi",
     "rbi/prism/translation/parser34.rbi",
@@ -129,6 +131,7 @@ Gem::Specification.new do |spec|
     "sig/prism/pattern.rbs",
     "sig/prism/reflection.rbs",
     "sig/prism/serialize.rbs",
+    "sig/prism/string_query.rbs",
     "sig/prism/visitor.rbs",
     "src/diagnostic.c",
     "src/encoding.c",
diff --git a/lib/prism/string_query.rb b/lib/prism/string_query.rb
new file mode 100644
index 0000000000..9011051d2b
--- /dev/null
+++ b/lib/prism/string_query.rb
@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Prism
+  # Query methods that allow categorizing strings based on their context for
+  # where they could be valid in a Ruby syntax tree.
+  class StringQuery
+    # The string that this query is wrapping.
+    attr_reader :string
+
+    # Initialize a new query with the given string.
+    def initialize(string)
+      @string = string
+    end
+
+    # Whether or not this string is a valid local variable name.
+    def local?
+      StringQuery.local?(string)
+    end
+
+    # Whether or not this string is a valid constant name.
+    def constant?
+      StringQuery.constant?(string)
+    end
+
+    # Whether or not this string is a valid method name.
+    def method_name?
+      StringQuery.method_name?(string)
+    end
+  end
+end