1 files changed, 100 insertions, 7 deletions

diff --git a/lib/error_highlight/base.rb b/lib/error_highlight/base.rb
index 8392979e24..b9c68b8eb8 100644
--- a/lib/error_highlight/base.rb
+++ b/lib/error_highlight/base.rb
@@ -1,12 +1,17 @@
 require_relative "version"
 
 module ErrorHighlight
-  # Identify the code fragment that seems associated with a given error
+  # Identify the code fragment at that a given exception occurred.
   #
-  # Arguments:
-  #  node: RubyVM::AbstractSyntaxTree::Node (script_lines should be enabled)
-  #  point_type: :name | :args
-  #  name: The name associated with the NameError/NoMethodError
+  # Options:
+  #
+  # point_type: :name | :args
+  #   :name (default) points the method/variable name that the exception occurred.
+  #   :args points the arguments of the method call that the exception occurred.
+  #
+  # backtrace_location: Thread::Backtrace::Location
+  #   It locates the code fragment of the given backtrace_location.
+  #   By default, it uses the first frame of backtrace_locations of the given exception.
   #
   # Returns:
   #  {
@@ -15,9 +20,65 @@ module ErrorHighlight
   #    last_lineno: Integer,
   #    last_column: Integer,
   #    snippet: String,
+  #    script_lines: [String],
   #  } | nil
-  def self.spot(...)
-    Spotter.new(...).spot
+  #
+  # Limitations:
+  #
+  # Currently, ErrorHighlight.spot only supports a single-line code fragment.
+  # Therefore, if the return value is not nil, first_lineno and last_lineno will have
+  # the same value. If the relevant code fragment spans multiple lines
+  # (e.g., Array#[] of +ary[(newline)expr(newline)]+), the method will return nil.
+  # This restriction may be removed in the future.
+  def self.spot(obj, **opts)
+    case obj
+    when Exception
+      exc = obj
+      loc = opts[:backtrace_location]
+      opts = { point_type: opts.fetch(:point_type, :name) }
+
+      unless loc
+        case exc
+        when TypeError, ArgumentError
+          opts[:point_type] = :args
+        end
+
+        locs = exc.backtrace_locations
+        return nil unless locs
+
+        loc = locs.first
+        return nil unless loc
+
+        opts[:name] = exc.name if NameError === obj
+      end
+
+      return nil unless Thread::Backtrace::Location === loc
+
+      node =
+        begin
+          RubyVM::AbstractSyntaxTree.of(loc, keep_script_lines: true)
+        rescue RuntimeError => error
+          # RubyVM::AbstractSyntaxTree.of raises an error with a message that
+          # includes "prism" when the ISEQ was compiled with the prism compiler.
+          # In this case, we'll set the node to `nil`. In the future, we will
+          # reparse with the prism parser and pass the parsed node to Spotter.
+          raise unless error.message.include?("prism")
+        end
+
+      Spotter.new(node, **opts).spot
+
+    when RubyVM::AbstractSyntaxTree::Node
+      Spotter.new(obj, **opts).spot
+
+    else
+      raise TypeError, "Exception is expected"
+    end
+
+  rescue SyntaxError,
+         SystemCallError, # file not found or something
+         ArgumentError # eval'ed code
+
+    return nil
   end
 
   class Spotter
@@ -46,9 +107,40 @@ module ErrorHighlight
       end
     end
 
+    OPT_GETCONSTANT_PATH = (RUBY_VERSION.split(".").map {|s| s.to_i } <=> [3, 2]) >= 0
+    private_constant :OPT_GETCONSTANT_PATH
+
     def spot
       return nil unless @node
 
+      if OPT_GETCONSTANT_PATH && @node.type == :COLON2
+        # In Ruby 3.2 or later, a nested constant access (like `Foo::Bar::Baz`)
+        # is compiled to one instruction (opt_getconstant_path).
+        # @node points to the node of the whole `Foo::Bar::Baz` even if `Foo`
+        # or `Foo::Bar` causes NameError.
+        # So we try to spot the sub-node that causes the NameError by using
+        # `NameError#name`.
+        subnodes = []
+        node = @node
+        while node.type == :COLON2
+          node2, const = node.children
+          subnodes << node if const == @name
+          node = node2
+        end
+        if node.type == :CONST || node.type == :COLON3
+          if node.children.first == @name
+            subnodes << node
+          end
+
+          # If we found only one sub-node whose name is equal to @name, use it
+          return nil if subnodes.size != 1
+          @node = subnodes.first
+        else
+          # Do nothing; opt_getconstant_path is used only when the const base is
+          # NODE_CONST (`Foo`) or NODE_COLON3 (`::Foo`)
+        end
+      end
+
       case @node.type
 
       when :CALL, :QCALL
@@ -122,6 +214,7 @@ module ErrorHighlight
           last_lineno: @end_lineno,
           last_column: @end_column,
           snippet: @snippet,
+          script_lines: @node.script_lines,
         }
       else
         return nil