From eb8c2773cb57628c09b36243d74cb752327819a6 Mon Sep 17 00:00:00 2001
From: ko1 <ko1@b2dd03c8-39d4-4d8f-98ff-823fe69b080e>
Date: Sat, 23 Dec 2017 12:48:24 +0000
Subject: RubyVM::InstructionSequence#each_child.

* iseq.c (iseqw_each_child): add RubyVM::InstructionSequence#each_child
  method for tools which want to manipulate ISeq.

* test/ruby/test_iseq.rb: add a test for this method.


git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@61425 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
---
 NEWS                   |  6 +++++
 iseq.c                 | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++
 test/ruby/test_iseq.rb | 41 +++++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+)

diff --git a/NEWS b/NEWS
index 1cfbedad37..7e55a04b2d 100644
--- a/NEWS
+++ b/NEWS
@@ -163,6 +163,12 @@ with all sufficient information, see the ChangeLog file or Redmine
 
   * Support new 5 emoji-related Unicode character properties
 
+* RubyVM::InstructionSequence
+
+  * New method:
+
+    * RubyVM::InstructionSequence#each_child
+
 * String
 
   * String#-@ deduplicates unfrozen strings.  Already-frozen
diff --git a/iseq.c b/iseq.c
index 7282614806..e25c423b83 100644
--- a/iseq.c
+++ b/iseq.c
@@ -1786,6 +1786,46 @@ rb_iseq_disasm(const rb_iseq_t *iseq)
     return str;
 }
 
+static VALUE
+rb_iseq_all_children(const rb_iseq_t *iseq)
+{
+    unsigned int i;
+    VALUE *code = rb_iseq_original_iseq(iseq);
+    VALUE all_children = rb_obj_hide(rb_ident_hash_new());
+    VALUE child;
+
+    if (iseq->body->catch_table) {
+	for (i = 0; i < iseq->body->catch_table->size; i++) {
+	    const struct iseq_catch_table_entry *entry = &iseq->body->catch_table->entries[i];
+	    child = (VALUE)entry->iseq;
+	    if (child) {
+		rb_hash_aset(all_children, child, Qtrue);
+	    }
+	}
+    }
+    for (i=0; i<iseq->body->iseq_size;) {
+	VALUE insn = code[i];
+	int len = insn_len(insn);
+	const char *types = insn_op_types(insn);
+	int j;
+
+	for (j=0; types[j]; j++) {
+	    switch (types[j]) {
+	      case TS_ISEQ:
+		child = code[i+j+1];
+		if (child) {
+		    rb_hash_aset(all_children, child, Qtrue);
+		}
+		break;
+	      default:
+		break;
+	    }
+	}
+	i += len;
+    }
+    return all_children;
+}
+
 /*
  *  call-seq:
  *     iseq.disasm -> str
@@ -1810,6 +1850,30 @@ iseqw_disasm(VALUE self)
     return rb_iseq_disasm(iseqw_check(self));
 }
 
+static int
+iseqw_each_child_i(VALUE key, VALUE value, VALUE dummy)
+{
+    rb_yield(iseqw_new((const rb_iseq_t *)key));
+    return ST_CONTINUE;
+}
+
+/*
+ *  call-seq:
+ *     iseq.each_child{|child_iseq| ...} -> iseq
+ *
+ *  Iterate all direct child instruction sequences.
+ *  Iteration order is implementation/version defined
+ *  so that people should not rely on the order.
+ */
+static VALUE
+iseqw_each_child(VALUE self)
+{
+    const rb_iseq_t *iseq = iseqw_check(self);
+    VALUE all_children = rb_iseq_all_children(iseq);
+    rb_hash_foreach(all_children, iseqw_each_child_i, Qnil);
+    return self;
+}
+
 /*
  *  Returns the instruction sequence containing the given proc or method.
  *
@@ -2615,6 +2679,7 @@ Init_ISeq(void)
     rb_define_method(rb_cISeq, "label", iseqw_label, 0);
     rb_define_method(rb_cISeq, "base_label", iseqw_base_label, 0);
     rb_define_method(rb_cISeq, "first_lineno", iseqw_first_lineno, 0);
+    rb_define_method(rb_cISeq, "each_child", iseqw_each_child, 0);
 
 #if 0 /* TBD */
     rb_define_private_method(rb_cISeq, "marshal_dump", iseqw_marshal_dump, 0);
diff --git a/test/ruby/test_iseq.rb b/test/ruby/test_iseq.rb
index 4de2aa7169..2f23f78e23 100644
--- a/test/ruby/test_iseq.rb
+++ b/test/ruby/test_iseq.rb
@@ -280,4 +280,45 @@ class TestISeq < Test::Unit::TestCase
       assert_match /:#{name}@/, ISeq.of(m).inspect, name
     end
   end
+
+  def test_each_child
+    iseq = ISeq.compile <<-EOS
+    class C
+      def foo
+        begin
+        rescue
+          p :rescue
+        ensure
+          p :ensure
+        end
+      end
+      def bar
+        1.times{
+          2.times{
+          }
+        }
+      end
+    end
+    class D < C
+    end
+    EOS
+
+    collect_iseq = lambda{|iseq|
+      iseqs = []
+      iseq.each_child{|child_iseq|
+        iseqs << collect_iseq.call(child_iseq)
+      }
+      ["#{iseq.label}@#{iseq.first_lineno}", *iseqs.sort_by{|k, *| k}]
+    }
+
+    expected = ["<compiled>@1",
+                  ["<class:C>@1",
+                    ["bar@10", ["block in bar@11",
+                            ["block (2 levels) in bar@12"]]],
+                    ["foo@2", ["ensure in foo@2"],
+                              ["rescue in foo@4"]]],
+                  ["<class:D>@17"]]
+
+    assert_equal expected, collect_iseq.call(iseq)
+  end
 end
-- 
cgit v1.2.3