5 files changed, 62 insertions, 23 deletions

diff --git a/ast.c b/ast.c
index 05bfc755a3..e180213ef7 100644
--- a/ast.c
+++ b/ast.c
@@ -196,14 +196,15 @@ static VALUE
 ast_s_of(rb_execution_context_t *ec, VALUE module, VALUE body, VALUE keep_script_lines)
 {
     VALUE path, node, lines = Qnil;
+    const rb_iseq_t *iseq;
     int node_id;
 
     if (rb_frame_info_p(body)) {
-        rb_frame_info_get(body, &path, &lines, &node_id);
-        if (NIL_P(path) && NIL_P(lines)) return Qnil;
+        iseq = rb_get_iseq_from_frame_info(body);
+        node_id = rb_get_node_id_from_frame_info(body);
     }
     else {
-        const rb_iseq_t *iseq = NULL;
+        iseq = NULL;
 
         if (rb_obj_is_proc(body)) {
             iseq = vm_proc_iseq(body);
@@ -213,17 +214,20 @@ ast_s_of(rb_execution_context_t *ec, VALUE module, VALUE body, VALUE keep_script
         else {
             iseq = rb_method_iseq(body);
         }
-        if (!iseq) {
-            return Qnil;
+        if (iseq) {
+            node_id = iseq->body->location.node_id;
         }
-        if (rb_iseq_from_eval_p(iseq)) {
-            rb_raise(rb_eArgError, "cannot get AST for method defined in eval");
-        }
-        path = rb_iseq_path(iseq);
-        lines = iseq->body->variable.script_lines;
-        node_id = iseq->body->location.node_id;
     }
 
+    if (!iseq) {
+        return Qnil;
+    }
+    if (rb_iseq_from_eval_p(iseq)) {
+        rb_raise(rb_eArgError, "cannot get AST for method defined in eval");
+    }
+    path = rb_iseq_path(iseq);
+    lines = iseq->body->variable.script_lines;
+
     if (!NIL_P(lines) || !NIL_P(lines = script_lines(path))) {
         node = rb_ast_parse_array(lines, keep_script_lines);
     }
diff --git a/internal/vm.h b/internal/vm.h
index e902a81411..bfb593176e 100644
--- a/internal/vm.h
+++ b/internal/vm.h
@@ -111,7 +111,8 @@ VALUE rb_backtrace_to_str_ary(VALUE obj);
 VALUE rb_backtrace_to_location_ary(VALUE obj);
 void rb_backtrace_each(VALUE (*iter)(VALUE recv, VALUE str), VALUE output);
 int rb_frame_info_p(VALUE obj);
-void rb_frame_info_get(VALUE obj, VALUE *path, VALUE *script_lines, int *node_id);
+int rb_get_node_id_from_frame_info(VALUE obj);
+const struct rb_iseq_struct *rb_get_iseq_from_frame_info(VALUE obj);
 
 MJIT_SYMBOL_EXPORT_BEGIN
 VALUE rb_ec_backtrace_object(const struct rb_execution_context_struct *ec);
diff --git a/lib/error_highlight/core_ext.rb b/lib/error_highlight/core_ext.rb
index ebb6788b02..0bf1992f49 100644
--- a/lib/error_highlight/core_ext.rb
+++ b/lib/error_highlight/core_ext.rb
@@ -29,7 +29,7 @@ module ErrorHighlight
 
         spot = ErrorHighlight.spot(node, **opts)
 
-      rescue SystemCallError, SyntaxError
+      rescue SystemCallError, SyntaxError, ArgumentError
       end
 
       if spot
diff --git a/test/ruby/test_ast.rb b/test/ruby/test_ast.rb
index 953c8435c3..2ab9f4e294 100644
--- a/test/ruby/test_ast.rb
+++ b/test/ruby/test_ast.rb
@@ -185,7 +185,7 @@ class TestAst < Test::Unit::TestCase
     end
   end
 
-  def test_of
+  def test_of_proc_and_method
     proc = Proc.new { 1 + 2 }
     method = self.method(__method__)
 
@@ -194,7 +194,6 @@ class TestAst < Test::Unit::TestCase
 
     assert_instance_of(RubyVM::AbstractSyntaxTree::Node, node_proc)
     assert_instance_of(RubyVM::AbstractSyntaxTree::Node, node_method)
-    assert_raise(TypeError) { RubyVM::AbstractSyntaxTree.of("1 + 2") }
 
     Tempfile.create(%w"test_of .rb") do |tmp|
       tmp.print "#{<<-"begin;"}\n#{<<-'end;'}"
@@ -211,7 +210,22 @@ class TestAst < Test::Unit::TestCase
     end
   end
 
-  def test_of_eval
+  def sample_backtrace_location
+    [caller_locations(0).first, __LINE__]
+  end
+
+  def test_of_backtrace_location
+    backtrace_location, lineno = sample_backtrace_location
+    node = RubyVM::AbstractSyntaxTree.of(backtrace_location)
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, node)
+    assert_equal(lineno, node.first_lineno)
+  end
+
+  def test_of_error
+    assert_raise(TypeError) { RubyVM::AbstractSyntaxTree.of("1 + 2") }
+  end
+
+  def test_of_proc_and_method_under_eval
     method = self.method(eval("def example_method_#{$$}; end"))
     assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }
 
@@ -229,6 +243,21 @@ class TestAst < Test::Unit::TestCase
 
     method = eval("Class.new{def example_method; end}.instance_method(:example_method)")
     assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }
+
+    method = eval("Class.new{def example_method; end}.instance_method(:example_method)")
+    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }
+  end
+
+  def test_of_backtrace_location_under_eval
+    m = Module.new do
+      eval(<<-END, nil, __FILE__, __LINE__)
+        def self.sample_backtrace_location
+          [caller_locations(0).first, __LINE__]
+        end
+      END
+    end
+    backtrace_location, lineno = m.sample_backtrace_location
+    assert_raise(ArgumentError) {  RubyVM::AbstractSyntaxTree.of(backtrace_location) }
   end
 
   def test_of_c_method
diff --git a/vm_backtrace.c b/vm_backtrace.c
index f2e9bed6ba..02c2ff4abd 100644
--- a/vm_backtrace.c
+++ b/vm_backtrace.c
@@ -333,20 +333,25 @@ location_node_id(rb_backtrace_location_t *loc)
 }
 #endif
 
-void
-rb_frame_info_get(VALUE obj, VALUE *path, VALUE *script_lines, int *node_id)
+int
+rb_get_node_id_from_frame_info(VALUE obj)
 {
 #ifdef USE_ISEQ_NODE_ID
     rb_backtrace_location_t *loc = location_ptr(obj);
-    const rb_iseq_t *iseq = location_iseq(loc);
-    *path = iseq ? rb_iseq_path(iseq) : Qnil;
-    *script_lines = iseq ? iseq->body->variable.script_lines : Qnil;
-    *node_id = location_node_id(loc);
+    return location_node_id(loc);
 #else
-    *path = Qnil;
+    return -1;
 #endif
 }
 
+const rb_iseq_t *
+rb_get_iseq_from_frame_info(VALUE obj)
+{
+    rb_backtrace_location_t *loc = location_ptr(obj);
+    const rb_iseq_t *iseq = location_iseq(loc);
+    return iseq;
+}
+
 static VALUE
 location_realpath(rb_backtrace_location_t *loc)
 {