1 files changed, 35 insertions, 31 deletions
diff --git a/lib/error_highlight/core_ext.rb b/lib/error_highlight/core_ext.rb
index 53e409dd8f..c3354f46cd 100644
--- a/lib/error_highlight/core_ext.rb
+++ b/lib/error_highlight/core_ext.rb
@@ -3,36 +3,38 @@ require_relative "formatter"
 module ErrorHighlight
   module CoreExt
     private def generate_snippet
-      locs = backtrace_locations
-      return "" unless locs
-
-      loc = locs.first
-      return "" unless loc
-
-      begin
-        node = RubyVM::AbstractSyntaxTree.of(loc, keep_script_lines: true)
-        opts = {}
-
-        case self
-        when NoMethodError, NameError
-          opts[:point_type] = :name
-          opts[:name] = name
-        when TypeError, ArgumentError
-          opts[:point_type] = :args
+      if ArgumentError === self && message =~ /\A(?:wrong number of arguments|missing keyword[s]?|unknown keyword[s]?|no keywords accepted)\b/
+        locs = self.backtrace_locations
+        return "" if locs.size < 2
+        callee_loc, caller_loc = locs
+        callee_spot = ErrorHighlight.spot(self, backtrace_location: callee_loc, point_type: :name)
+        caller_spot = ErrorHighlight.spot(self, backtrace_location: caller_loc, point_type: :name)
+        if caller_spot && callee_spot &&
+            caller_loc.path == callee_loc.path &&
+            caller_loc.lineno == callee_loc.lineno &&
+            caller_spot == callee_spot
+          callee_loc = callee_spot = nil
         end
-
-        spot = ErrorHighlight.spot(node, **opts)
-
-      rescue SyntaxError
-      rescue SystemCallError # file not found or something
-      rescue ArgumentError   # eval'ed code
-      end
-
-      if spot
+        ret = +"\n"
+        [["caller", caller_loc, caller_spot], ["callee", callee_loc, callee_spot]].each do |header, loc, spot|
+          out = nil
+          if loc
+            out = "    #{ header }: #{ loc.path }:#{ loc.lineno }"
+            if spot
+              _, _, snippet, highlight = ErrorHighlight.formatter.message_for(spot).lines
+              out += "\n    | #{ snippet }      #{ highlight }"
+            else
+              # do nothing
+            end
+          end
+          ret << "\n" + out if out
+        end
+        ret
+      else
+        spot = ErrorHighlight.spot(self)
+        return "" unless spot
         return ErrorHighlight.formatter.message_for(spot)
       end
-
-      ""
     end
 
     if Exception.method_defined?(:detailed_message)
@@ -65,8 +67,10 @@ module ErrorHighlight
 
   NameError.prepend(CoreExt)
 
-  # The extension for TypeError/ArgumentError is temporarily disabled due to many test failures
-
-  #TypeError.prepend(CoreExt)
-  #ArgumentError.prepend(CoreExt)
+  if Exception.method_defined?(:detailed_message)
+    # ErrorHighlight is enabled for TypeError and ArgumentError only when Exception#detailed_message is available.
+    # This is because changing ArgumentError#message is highly incompatible.
+    TypeError.prepend(CoreExt)
+    ArgumentError.prepend(CoreExt)
+  end
 end