Make it possible to get AST::Node from Thread::Backtrace::Location

RubyVM::AST.of(Thread::Backtrace::Location) returns a node that
corresponds to the location. Typically, the node is a method call, but
not always.

This change also includes iseq's dump/load support of node_ids for each
instructions.

5 files changed, 107 insertions, 21 deletions

diff --git a/ast.c b/ast.c
index 78de316c61..04ba3c987f 100644
--- a/ast.c
+++ b/ast.c
@@ -197,23 +197,26 @@ ast_s_of(rb_execution_context_t *ec, VALUE module, VALUE body, VALUE save_script
 {
     VALUE path, node, lines;
     int node_id;
-    const rb_iseq_t *iseq = NULL;
 
-    if (rb_obj_is_proc(body)) {
-        iseq = vm_proc_iseq(body);
-
-        if (!rb_obj_is_iseq((VALUE)iseq)) {
-            iseq = NULL;
-        }
+    if (rb_frame_info_p(body)) {
+        rb_frame_info_get(body, &path, &node_id);
+        if (NIL_P(path)) return Qnil;
     }
     else {
-        iseq = rb_method_iseq(body);
-    }
+        const rb_iseq_t *iseq = NULL;
+
+        if (rb_obj_is_proc(body)) {
+            iseq = vm_proc_iseq(body);
 
-    if (!iseq) return Qnil;
+            if (!rb_obj_is_iseq((VALUE)iseq)) return Qnil;
+        }
+        else {
+            iseq = rb_method_iseq(body);
+        }
+        path = rb_iseq_path(iseq);
+        node_id = iseq->body->location.node_id;
+    }
 
-    path = rb_iseq_path(iseq);
-    node_id = iseq->body->location.node_id;
     if (!NIL_P(lines = script_lines(path))) {
         node = rb_ast_parse_array(lines, save_script_lines);
     }
diff --git a/compile.c b/compile.c
index 4b967a87a0..a653b2d8da 100644
--- a/compile.c
+++ b/compile.c
@@ -9704,13 +9704,13 @@ event_name_to_flag(VALUE sym)
 
 static int
 iseq_build_from_ary_body(rb_iseq_t *iseq, LINK_ANCHOR *const anchor,
-			 VALUE body, VALUE labels_wrapper)
+			 VALUE body, VALUE node_ids, VALUE labels_wrapper)
 {
     /* TODO: body should be frozen */
     long i, len = RARRAY_LEN(body);
     struct st_table *labels_table = DATA_PTR(labels_wrapper);
     int j;
-    int line_no = 0;
+    int line_no = 0, node_id = -1, insn_idx = 0;
     int ret = COMPILE_OK;
 
     /*
@@ -9744,6 +9744,10 @@ iseq_build_from_ary_body(rb_iseq_t *iseq, LINK_ANCHOR *const anchor,
 	    st_data_t insn_id;
 	    VALUE insn;
 
+            if (node_ids) {
+                node_id = NUM2INT(rb_ary_entry(node_ids, insn_idx++));
+            }
+
 	    insn = (argc < 0) ? Qnil : RARRAY_AREF(obj, 0);
 	    if (st_lookup(insn_table, (st_data_t)insn, &insn_id) == 0) {
 		/* TODO: exception */
@@ -9764,7 +9768,7 @@ iseq_build_from_ary_body(rb_iseq_t *iseq, LINK_ANCHOR *const anchor,
                 argv = compile_data_calloc2(iseq, sizeof(VALUE), argc);
 
                 // add element before operand setup to make GC root
-                NODE dummy_line_node = generate_dummy_line_node(line_no, -1);
+                NODE dummy_line_node = generate_dummy_line_node(line_no, node_id);
                 ADD_ELEM(anchor,
                          (LINK_ELEMENT*)new_insn_core(iseq, &dummy_line_node,
                                                       (enum ruby_vminsn_type)insn_id, argc, argv));
@@ -9848,7 +9852,7 @@ iseq_build_from_ary_body(rb_iseq_t *iseq, LINK_ANCHOR *const anchor,
 		}
 	    }
             else {
-                NODE dummy_line_node = generate_dummy_line_node(line_no, -1);
+                NODE dummy_line_node = generate_dummy_line_node(line_no, node_id);
                 ADD_ELEM(anchor,
                          (LINK_ELEMENT*)new_insn_core(iseq, &dummy_line_node,
                                                       (enum ruby_vminsn_type)insn_id, argc, NULL));
@@ -10054,6 +10058,14 @@ rb_iseq_build_from_ary(rb_iseq_t *iseq, VALUE misc, VALUE locals, VALUE params,
 #undef INT_PARAM
     }
 
+    VALUE node_ids = Qfalse;
+#ifdef EXPERIMENTAL_ISEQ_NODE_ID
+    node_ids = rb_hash_aref(misc, ID2SYM(rb_intern("node_ids")));
+    if (!RB_TYPE_P(node_ids, T_ARRAY)) {
+	rb_raise(rb_eTypeError, "node_ids is not an array");
+    }
+#endif
+
     if (RB_TYPE_P(arg_opt_labels, T_ARRAY)) {
 	len = RARRAY_LENINT(arg_opt_labels);
 	iseq->body->param.flags.has_opt = !!(len - 1 >= 0);
@@ -10103,7 +10115,7 @@ rb_iseq_build_from_ary(rb_iseq_t *iseq, VALUE misc, VALUE locals, VALUE params,
     iseq_build_from_ary_exception(iseq, labels_table, exception);
 
     /* body */
-    iseq_build_from_ary_body(iseq, anchor, body, labels_wrapper);
+    iseq_build_from_ary_body(iseq, anchor, body, node_ids, labels_wrapper);
 
     iseq->body->param.size = arg_size;
     iseq->body->local_table_size = local_size;
@@ -10928,6 +10940,9 @@ ibf_dump_insns_info_body(struct ibf_dump *dump, const rb_iseq_t *iseq)
     unsigned int i;
     for (i = 0; i < iseq->body->insns_info.size; i++) {
         ibf_dump_write_small_value(dump, entries[i].line_no);
+#ifdef EXPERIMENTAL_ISEQ_NODE_ID
+        ibf_dump_write_small_value(dump, entries[i].node_id);
+#endif
         ibf_dump_write_small_value(dump, entries[i].events);
     }
 
@@ -10943,6 +10958,9 @@ ibf_load_insns_info_body(const struct ibf_load *load, ibf_offset_t body_offset,
     unsigned int i;
     for (i = 0; i < size; i++) {
         entries[i].line_no = (int)ibf_load_small_value(load, &reading_pos);
+#ifdef EXPERIMENTAL_ISEQ_NODE_ID
+        entries[i].node_id = (int)ibf_load_small_value(load, &reading_pos);
+#endif
         entries[i].events = (rb_event_flag_t)ibf_load_small_value(load, &reading_pos);
     }
 
diff --git a/internal/vm.h b/internal/vm.h
index 689b4fa61f..8251b09e98 100644
--- a/internal/vm.h
+++ b/internal/vm.h
@@ -111,6 +111,8 @@ int rb_backtrace_p(VALUE obj);
 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, int *node_id);
 
 MJIT_SYMBOL_EXPORT_BEGIN
 VALUE rb_ec_backtrace_object(const struct rb_execution_context_struct *ec);
diff --git a/iseq.c b/iseq.c
index 1609432770..26c59b8cbb 100644
--- a/iseq.c
+++ b/iseq.c
@@ -2998,7 +2998,7 @@ iseq_data_to_ary(const rb_iseq_t *iseq)
 		INT2FIX(iseq_body->location.code_location.end_pos.lineno),
 		INT2FIX(iseq_body->location.code_location.end_pos.column)));
 #ifdef EXPERIMENTAL_ISEQ_NODE_ID
-    rb_hash_aset(misc, ID2SYM(rb_intern("node_ids_for_each_insn")), node_ids);
+    rb_hash_aset(misc, ID2SYM(rb_intern("node_ids")), node_ids);
 #endif
 
     /*
diff --git a/vm_backtrace.c b/vm_backtrace.c
index 01c9aeb355..d0fdda343f 100644
--- a/vm_backtrace.c
+++ b/vm_backtrace.c
@@ -34,7 +34,7 @@ id2str(ID id)
 #define ALL_BACKTRACE_LINES -1
 
 inline static int
-calc_lineno(const rb_iseq_t *iseq, const VALUE *pc)
+calc_pos(const rb_iseq_t *iseq, const VALUE *pc, int *lineno, int *node_id)
 {
     VM_ASSERT(iseq);
     VM_ASSERT(iseq->body);
@@ -46,7 +46,11 @@ calc_lineno(const rb_iseq_t *iseq, const VALUE *pc)
             VM_ASSERT(! iseq->body->local_table_size);
             return 0;
         }
-        return FIX2INT(iseq->body->location.first_lineno);
+        if (lineno) *lineno = FIX2INT(iseq->body->location.first_lineno);
+#ifdef EXPERIMENTAL_ISEQ_NODE_ID
+        if (node_id) *node_id = -1;
+#endif
+        return 1;
     }
     else {
         ptrdiff_t n = pc - iseq->body->iseq_encoded;
@@ -65,10 +69,32 @@ calc_lineno(const rb_iseq_t *iseq, const VALUE *pc)
             __builtin_trap();
         }
 #endif
-        return rb_iseq_line_no(iseq, pos);
+        if (lineno) *lineno = rb_iseq_line_no(iseq, pos);
+#ifdef EXPERIMENTAL_ISEQ_NODE_ID
+        if (node_id) *node_id = rb_iseq_node_id(iseq, pos);
+#endif
+        return 1;
     }
 }
 
+inline static int
+calc_lineno(const rb_iseq_t *iseq, const VALUE *pc)
+{
+    int lineno;
+    if (calc_pos(iseq, pc, &lineno, NULL)) return lineno;
+    return 0;
+}
+
+#ifdef EXPERIMENTAL_ISEQ_NODE_ID
+inline static int
+calc_node_id(const rb_iseq_t *iseq, const VALUE *pc)
+{
+    int node_id;
+    if (calc_pos(iseq, pc, NULL, &node_id)) return node_id;
+    return -1;
+}
+#endif
+
 int
 rb_vm_get_sourceline(const rb_control_frame_t *cfp)
 {
@@ -143,6 +169,12 @@ static const rb_data_type_t location_data_type = {
     0, 0, RUBY_TYPED_FREE_IMMEDIATELY
 };
 
+int
+rb_frame_info_p(VALUE obj)
+{
+    return rb_typeddata_is_kind_of(obj, &location_data_type);
+}
+
 static inline rb_backtrace_location_t *
 location_ptr(VALUE locobj)
 {
@@ -287,6 +319,37 @@ location_path_m(VALUE self)
     return location_path(location_ptr(self));
 }
 
+#ifdef EXPERIMENTAL_ISEQ_NODE_ID
+static int
+location_node_id(rb_backtrace_location_t *loc)
+{
+    switch (loc->type) {
+      case LOCATION_TYPE_ISEQ:
+	return calc_node_id(loc->body.iseq.iseq, loc->body.iseq.pc);
+      case LOCATION_TYPE_CFUNC:
+        if (loc->body.cfunc.prev_loc) {
+            return location_node_id(loc->body.cfunc.prev_loc);
+        }
+        return -1;
+      default:
+        rb_bug("location_node_id: unreachable");
+        UNREACHABLE;
+    }
+}
+#endif
+
+void
+rb_frame_info_get(VALUE obj, VALUE *path, int *node_id)
+{
+#ifdef EXPERIMENTAL_ISEQ_NODE_ID
+    rb_backtrace_location_t *loc = location_ptr(obj);
+    *path = location_path(loc);
+    *node_id = location_node_id(loc);
+#else
+    *path = Qnil;
+#endif
+}
+
 static VALUE
 location_realpath(rb_backtrace_location_t *loc)
 {