Implement `RubyVM::AST.of` [Feature #14836]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@65542 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

ast.c

@@ -5,6 +5,7 @@
 #include "internal.h"
 #include "node.h"
 #include "vm_core.h"
+#include "iseq.h"
 
 static VALUE rb_mAST;
 static VALUE rb_cNode;
@@ -125,6 +126,59 @@ rb_ast_s_parse_file(VALUE module, VALUE path)
     return obj;
 }
 
+static VALUE node_children(rb_ast_t*, NODE*);
+
+static VALUE
+node_find(VALUE self, const int node_id)
+{
+    VALUE ary;
+    long i;
+    struct ASTNodeData *data;
+    TypedData_Get_Struct(self, struct ASTNodeData, &rb_node_type, data);
+
+    if (nd_node_id(data->node) == node_id) return self;
+
+    ary = node_children(data->ast, data->node);
+
+    for (i = 0; i < RARRAY_LEN(ary); i++) {
+        VALUE child = RARRAY_AREF(ary, i);
+
+        if (CLASS_OF(child) == rb_cNode) {
+            VALUE result = node_find(child, node_id);
+            if (RTEST(result)) return result;
+        }
+    }
+
+    return Qnil;
+}
+
+static VALUE
+rb_ast_s_of(VALUE module, VALUE body)
+{
+    VALUE path, node;
+    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;
+        }
+    }
+    else {
+        iseq = rb_method_iseq(body);
+    }
+
+    if (!iseq) return Qnil;
+
+    path = rb_iseq_path(iseq);
+    node_id = iseq->body->location.node_id;
+    node = rb_ast_s_parse_file(module, path);
+
+    return node_find(node, node_id);
+}
+
 static VALUE
 rb_ast_node_alloc(VALUE klass)
 {
@@ -619,6 +673,7 @@ Init_ast(void)
     rb_undef_alloc_func(rb_cNode);
     rb_define_singleton_method(rb_mAST, "parse", rb_ast_s_parse, 1);
     rb_define_singleton_method(rb_mAST, "parse_file", rb_ast_s_parse_file, 1);
+    rb_define_singleton_method(rb_mAST, "of", rb_ast_s_of, 1);
     rb_define_method(rb_cNode, "type", rb_ast_node_type, 0);
     rb_define_method(rb_cNode, "first_lineno", rb_ast_node_first_lineno, 0);
     rb_define_method(rb_cNode, "first_column", rb_ast_node_first_column, 0);

compile.c

@@ -9077,6 +9077,7 @@ ibf_load_iseq_each(const struct ibf_load *load, rb_iseq_t *iseq, ibf_offset_t of
     RB_OBJ_WRITE(iseq, &load_body->location.base_label,    ibf_load_location_str(load, body->location.base_label));
     RB_OBJ_WRITE(iseq, &load_body->location.label,         ibf_load_location_str(load, body->location.label));
     load_body->location.first_lineno = body->location.first_lineno;
+    load_body->location.node_id = body->location.node_id;
     load_body->location.code_location = body->location.code_location;
 
     load_body->is_entries           = ZALLOC_N(union iseq_inline_storage_entry, body->is_size);

iseq.c

@@ -371,7 +371,7 @@ rb_iseq_pathobj_set(const rb_iseq_t *iseq, VALUE path, VALUE realpath)
 }
 
 static rb_iseq_location_t *
-iseq_location_setup(rb_iseq_t *iseq, VALUE name, VALUE path, VALUE realpath, VALUE first_lineno, const rb_code_location_t *code_location)
+iseq_location_setup(rb_iseq_t *iseq, VALUE name, VALUE path, VALUE realpath, VALUE first_lineno, const rb_code_location_t *code_location, const int node_id)
 {
     rb_iseq_location_t *loc = &iseq->body->location;
 
@@ -380,6 +380,7 @@ iseq_location_setup(rb_iseq_t *iseq, VALUE name, VALUE path, VALUE realpath, VAL
     RB_OBJ_WRITE(iseq, &loc->base_label, name);
     loc->first_lineno = first_lineno;
     if (code_location) {
+        loc->node_id = node_id;
 	loc->code_location = *code_location;
     }
     else {
@@ -420,7 +421,7 @@ set_relation(rb_iseq_t *iseq, const rb_iseq_t *piseq)
 
 static VALUE
 prepare_iseq_build(rb_iseq_t *iseq,
-		   VALUE name, VALUE path, VALUE realpath, VALUE first_lineno, const rb_code_location_t *code_location,
+		   VALUE name, VALUE path, VALUE realpath, VALUE first_lineno, const rb_code_location_t *code_location, const int node_id,
 		   const rb_iseq_t *parent, enum iseq_type type,
 		   const rb_compile_option_t *option)
 {
@@ -435,7 +436,7 @@ prepare_iseq_build(rb_iseq_t *iseq,
     set_relation(iseq, parent);
 
     name = rb_fstring(name);
-    iseq_location_setup(iseq, name, path, realpath, first_lineno, code_location);
+    iseq_location_setup(iseq, name, path, realpath, first_lineno, code_location, node_id);
     if (iseq != body->local_iseq) {
 	RB_OBJ_WRITE(iseq, &body->location.base_label, body->local_iseq->body->location.label);
     }
@@ -704,7 +705,7 @@ rb_iseq_new_with_opt(const rb_ast_body_t *ast, VALUE name, VALUE path, VALUE rea
     new_opt = option ? *option : COMPILE_OPTION_DEFAULT;
     if (ast && ast->compile_option) rb_iseq_make_compile_option(&new_opt, ast->compile_option);
 
-    prepare_iseq_build(iseq, name, path, realpath, first_lineno, node ? &node->nd_loc : NULL, parent, type, &new_opt);
+    prepare_iseq_build(iseq, name, path, realpath, first_lineno, node ? &node->nd_loc : NULL, node ? nd_node_id(node) : -1, parent, type, &new_opt);
 
     rb_iseq_compile_node(iseq, node);
     finish_iseq_build(iseq);
@@ -721,7 +722,7 @@ rb_iseq_new_ifunc(const struct vm_ifunc *ifunc, VALUE name, VALUE path, VALUE re
     rb_iseq_t *iseq = iseq_alloc();
 
     if (!option) option = &COMPILE_OPTION_DEFAULT;
-    prepare_iseq_build(iseq, name, path, realpath, first_lineno, NULL, parent, type, option);
+    prepare_iseq_build(iseq, name, path, realpath, first_lineno, NULL, -1, parent, type, option);
 
     rb_iseq_compile_ifunc(iseq, ifunc);
     finish_iseq_build(iseq);
@@ -780,7 +781,7 @@ iseq_load(VALUE data, const rb_iseq_t *parent, VALUE opt)
     rb_iseq_t *iseq = iseq_alloc();
 
     VALUE magic, version1, version2, format_type, misc;
-    VALUE name, path, realpath, first_lineno, code_location;
+    VALUE name, path, realpath, first_lineno, code_location, node_id;
     VALUE type, body, locals, params, exception;
 
     st_data_t iseq_type;
@@ -821,6 +822,8 @@ iseq_load(VALUE data, const rb_iseq_t *parent, VALUE opt)
 	rb_raise(rb_eTypeError, "unsupport type: :%"PRIsVALUE, rb_sym2str(type));
     }
 
+    node_id = rb_hash_aref(misc, ID2SYM(rb_intern("node_id")));
+
     code_location = rb_hash_aref(misc, ID2SYM(rb_intern("code_location")));
     if (RB_TYPE_P(code_location, T_ARRAY) && RARRAY_LEN(code_location) == 4) {
 	tmp_loc.beg_pos.lineno = NUM2INT(rb_ary_entry(code_location, 0));
@@ -831,7 +834,7 @@ iseq_load(VALUE data, const rb_iseq_t *parent, VALUE opt)
 
     make_compile_option(&option, opt);
     option.peephole_optimization = FALSE; /* because peephole optimization can modify original iseq */
-    prepare_iseq_build(iseq, name, path, realpath, first_lineno, &tmp_loc,
+    prepare_iseq_build(iseq, name, path, realpath, first_lineno, &tmp_loc, NUM2INT(node_id),
 		       parent, (enum iseq_type)iseq_type, &option);
 
     rb_iseq_build_from_ary(iseq, misc, locals, params, exception, body);
@@ -2730,6 +2733,7 @@ iseq_data_to_ary(const rb_iseq_t *iseq)
     rb_hash_aset(misc, ID2SYM(rb_intern("arg_size")), INT2FIX(iseq_body->param.size));
     rb_hash_aset(misc, ID2SYM(rb_intern("local_size")), INT2FIX(iseq_body->local_table_size));
     rb_hash_aset(misc, ID2SYM(rb_intern("stack_max")), INT2FIX(iseq_body->stack_max));
+    rb_hash_aset(misc, ID2SYM(rb_intern("node_id")), INT2FIX(iseq_body->location.node_id));
     rb_hash_aset(misc, ID2SYM(rb_intern("code_location")),
 	    rb_ary_new_from_args(4,
 		INT2FIX(iseq_body->location.code_location.beg_pos.lineno),

node.h

@@ -164,6 +164,7 @@ typedef struct RNode {
 	VALUE value;
     } u3;
     rb_code_location_t nd_loc;
+    int node_id;
 } NODE;
 
 #define RNODE(obj)  (R_CAST(RNode)(obj))
@@ -201,6 +202,8 @@ typedef struct RNode {
 #define nd_set_last_lineno(n, v) ((n)->nd_loc.end_pos.lineno = (v))
 #define nd_last_loc(n) ((n)->nd_loc.end_pos)
 #define nd_set_last_loc(n, v) (nd_last_loc(n) = (v))
+#define nd_node_id(n) ((n)->node_id)
+#define nd_set_node_id(n,id) ((n)->node_id = (id))
 
 #define nd_head  u1.node
 #define nd_alen  u2.argc

parse.y

@@ -234,6 +234,7 @@ struct parser_params {
     ID cur_arg;
 
     rb_ast_t *ast;
+    int node_id;
 
     unsigned int command_start:1;
     unsigned int eofp: 1;
@@ -339,6 +340,14 @@ static NODE* node_newnode(struct parser_params *, enum node_type, VALUE, VALUE,
 
 static NODE *nd_set_loc(NODE *nd, const YYLTYPE *loc);
 
+static int
+parser_get_node_id(struct parser_params *p)
+{
+    int node_id = p->node_id;
+    p->node_id++;
+    return node_id;
+}
+
 #ifndef RIPPER
 static inline void
 set_line_body(NODE *body, int line)
@@ -8294,6 +8303,7 @@ node_newnode(struct parser_params *p, enum node_type type, VALUE a0, VALUE a1, V
     rb_node_init(n, type, a0, a1, a2);
 
     nd_set_loc(n, loc);
+    nd_set_node_id(n, parser_get_node_id(p));
     return n;
 }
 
@@ -10774,6 +10784,7 @@ parser_initialize(struct parser_params *p)
     p->command_start = TRUE;
     p->ruby_sourcefile_string = Qnil;
     p->lex.lpar_beg = -1; /* make lambda_beginning_p() == FALSE at first */
+    p->node_id = 0;
 #ifdef RIPPER
     p->delayed = Qnil;
     p->result = Qnil;

test/ruby/test_ast.rb

@@ -170,6 +170,18 @@ class TestAst < Test::Unit::TestCase
     end
   end
 
+  def test_of
+    proc = Proc.new { 1 + 2 }
+    method = self.method(__method__)
+
+    node_proc = RubyVM::AST.of(proc)
+    node_method = RubyVM::AST.of(method)
+
+    assert_instance_of(RubyVM::AST::Node, node_proc)
+    assert_instance_of(RubyVM::AST::Node, node_method)
+    assert_raise(TypeError) { RubyVM::AST.of("1 + 2") }
+  end
+
   def test_scope_local_variables
     node = RubyVM::AST.parse("x = 0")
     lv, _, body = *node.children

vm_core.h

@@ -298,6 +298,7 @@ typedef struct rb_iseq_location_struct {
     VALUE base_label;   /* String */
     VALUE label;        /* String */
     VALUE first_lineno; /* TODO: may be unsigned short */
+    int node_id;
     rb_code_location_t code_location;
 } rb_iseq_location_t;