* ext/objspace/objspace.c: add ObjectSpace#reachable_objects_from.

  This method returns an array of objects referenced by given object.
  If given object is special objects such as true/false/nil/Fixnum etc
  then it returns nil. See rdoc for details.
  [ruby-core:39772]
* test/objspace/test_objspace.rb: add a test for this method.
* gc.c: add rb_objspace_reachable_objects_from().
  To make this function, add several member `mark_func_data'
  to rb_objspace_t.  If mark_func_data is not null, then
  gc_mark() calls mark_func_data::mark_func.
* gc.h: export rb_objspace_reachable_objects_from().



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

5 files changed, 202 insertions, 31 deletions

diff --git a/ChangeLog b/ChangeLog
index a230327cb2..f28b38af7e 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,20 @@
+Fri Oct 05 15:26:18 2012  Koichi Sasada  <ko1@atdot.net>
+
+	* ext/objspace/objspace.c: add ObjectSpace#reachable_objects_from.
+	  This method returns an array of objects referenced by given object.
+	  If given object is special objects such as true/false/nil/Fixnum etc
+	  then it returns nil. See rdoc for details.
+	  [ruby-core:39772]
+
+	* test/objspace/test_objspace.rb: add a test for this method.
+
+	* gc.c: add rb_objspace_reachable_objects_from().
+	  To make this function, add several member `mark_func_data'
+	  to rb_objspace_t.  If mark_func_data is not null, then
+	  gc_mark() calls mark_func_data::mark_func.
+
+	* gc.h: export rb_objspace_reachable_objects_from().
+
 Fri Oct  5 16:04:33 2012  Narihiro Nakamura  <authornari@gmail.com>
 
 	* gc.c (chain_finalized_object): remove to check a mark flag since
diff --git a/ext/objspace/objspace.c b/ext/objspace/objspace.c
index 33106382ca..1842209641 100644
--- a/ext/objspace/objspace.c
+++ b/ext/objspace/objspace.c
@@ -37,6 +37,7 @@ size_t rb_ary_memsize(VALUE);
 size_t rb_io_memsize(const rb_io_t *);
 size_t rb_generic_ivar_memsize(VALUE);
 size_t rb_objspace_data_type_memsize(VALUE obj);
+VALUE rb_objspace_reachable_objects_from(VALUE obj);
 
 static size_t
 memsize_of(VALUE obj)
@@ -626,6 +627,47 @@ count_tdata_objects(int argc, VALUE *argv, VALUE self)
     return hash;
 }
 
+/*
+ *  call-seq:
+ *     ObjectSpace.reachable_objects_from(obj) -> array or nil
+ *
+ *  [MRI specific feature] Return all reachable objects from `obj'.
+ *
+ *  This method returns all reachable objects from `obj'.
+ *  If `obj' has references two or more references to same object `x',
+ *  them returned array only include one `x' object.
+ *  If `obj' is non-markable (non-heap management) object such as
+ *  true, false, nil, symbols and Fixnums (and Flonum) them it simply
+ *  returns nil.
+ *
+ *  With this method, you can find memory leaks.
+ *
+ *  This method is not expected to work except C Ruby.
+ *
+ *  Example:
+ *    ObjectSpace.reachable_objects_from(['a', 'b', 'c'])
+ *    #=> [Array, 'a', 'b', 'c']
+ *
+ *    ObjectSpace.reachable_objects_from(['a', 'a', 'a'])
+ *    #=> [Array, 'a', 'a', 'a'] # all 'a' strings have different object id
+ *
+ *    ObjectSpace.reachable_objects_from([v = 'a', v, v])
+ *    #=> [Array, 'a']
+ *
+ *    ObjectSpace.reachable_objects_from(1)
+ *    #=> nil # 1 is not markable (heap managed) object
+ *
+ *  Limitation: Current implementation can't acquire internal objects.
+ *              This means that you can't acquire complete object graph
+ *              (heap snapshot).  This is future work.
+ *
+ */
+static VALUE
+reachable_objects_from(VALUE self, VALUE obj)
+{
+    return rb_objspace_reachable_objects_from(obj);
+}
+
 /* objspace library extends ObjectSpace module and add several
  * methods to get internal statistic information about
  * object/memory management.
@@ -642,10 +684,11 @@ Init_objspace(void)
     VALUE rb_mObjSpace = rb_const_get(rb_cObject, rb_intern("ObjectSpace"));
 
     rb_define_module_function(rb_mObjSpace, "memsize_of", memsize_of_m, 1);
-    rb_define_module_function(rb_mObjSpace, "memsize_of_all",
-			      memsize_of_all_m, -1);
+    rb_define_module_function(rb_mObjSpace, "memsize_of_all", memsize_of_all_m, -1);
 
     rb_define_module_function(rb_mObjSpace, "count_objects_size", count_objects_size, -1);
     rb_define_module_function(rb_mObjSpace, "count_nodes", count_nodes, -1);
     rb_define_module_function(rb_mObjSpace, "count_tdata_objects", count_tdata_objects, -1);
+
+    rb_define_module_function(rb_mObjSpace, "reachable_objects_from", reachable_objects_from, 1);
 }
diff --git a/gc.c b/gc.c
index e87e029596..d34e70940b 100644
--- a/gc.c
+++ b/gc.c
@@ -273,6 +273,11 @@ typedef struct rb_objspace {
     struct gc_list *global_list;
     size_t count;
     int gc_stress;
+
+    struct mark_func_data_struct {
+	VALUE data;
+	void (*mark_func)(struct rb_objspace *objspace, VALUE v);
+    } *mark_func_data;
 } rb_objspace_t;
 
 #if defined(ENABLE_VM_OBJSPACE) && ENABLE_VM_OBJSPACE
@@ -1138,30 +1143,40 @@ struct os_each_struct {
 };
 
 static int
+internal_object_p(VALUE obj)
+{
+    RVALUE *p = (RVALUE *)obj;
+
+    if (p->as.basic.flags) {
+	switch (BUILTIN_TYPE(p)) {
+	  case T_NONE:
+	  case T_ICLASS:
+	  case T_NODE:
+	  case T_ZOMBIE:
+	    break;
+	  case T_CLASS:
+	    if (FL_TEST(p, FL_SINGLETON))
+	      break;
+	  default:
+	    if (!p->as.basic.klass) break;
+	    return 0;
+	}
+    }
+    return 1;
+}
+
+static int
 os_obj_of_i(void *vstart, void *vend, size_t stride, void *data)
 {
     struct os_each_struct *oes = (struct os_each_struct *)data;
     RVALUE *p = (RVALUE *)vstart, *pend = (RVALUE *)vend;
-    volatile VALUE v;
 
     for (; p != pend; p++) {
-	if (p->as.basic.flags) {
-	    switch (BUILTIN_TYPE(p)) {
-	      case T_NONE:
-	      case T_ICLASS:
-	      case T_NODE:
-	      case T_ZOMBIE:
-		continue;
-	      case T_CLASS:
-		if (FL_TEST(p, FL_SINGLETON))
-		    continue;
-	      default:
-		if (!p->as.basic.klass) continue;
-		v = (VALUE)p;
-		if (!oes->of || rb_obj_is_kind_of(v, oes->of)) {
-		    rb_yield(v);
-		    oes->num++;
-		}
+	volatile VALUE v = (VALUE)p;
+	if (!internal_object_p(v)) {
+	    if (!oes->of || rb_obj_is_kind_of(v, oes->of)) {
+		rb_yield(v);
+		oes->num++;
 	    }
 	}
     }
@@ -2521,17 +2536,31 @@ gc_mark_ptr(rb_objspace_t *objspace, VALUE ptr)
     return 1;
 }
 
+static int
+markable_object_p(rb_objspace_t *objspace, VALUE ptr)
+{
+    register RVALUE *obj = RANY(ptr);
+
+    if (rb_special_const_p(ptr)) return 0; /* special const not marked */
+    if (obj->as.basic.flags == 0) return 0 ;       /* free cell */
+
+    return 1;
+}
+
 static void
 gc_mark(rb_objspace_t *objspace, VALUE ptr)
 {
-    register RVALUE *obj;
-
-    obj = RANY(ptr);
-    if (rb_special_const_p(ptr)) return; /* special const not marked */
-    if (obj->as.basic.flags == 0) return;       /* free cell */
-    if (!gc_mark_ptr(objspace, ptr)) return;	/* already marked */
+    if (!markable_object_p(objspace, ptr)) {
+	return;
+    }
 
-    push_mark_stack(&objspace->mark_stack, ptr);
+    if (LIKELY(objspace->mark_func_data == 0)) {
+	if (!gc_mark_ptr(objspace, ptr)) return; /* already marked */
+	push_mark_stack(&objspace->mark_stack, ptr);
+    }
+    else {
+	objspace->mark_func_data->mark_func(objspace, ptr);
+    }
 }
 
 void
@@ -2548,10 +2577,16 @@ gc_mark_children(rb_objspace_t *objspace, VALUE ptr)
     goto marking;		/* skip */
 
   again:
-    obj = RANY(ptr);
-    if (rb_special_const_p(ptr)) return; /* special const not marked */
-    if (obj->as.basic.flags == 0) return;       /* free cell */
-    if (!gc_mark_ptr(objspace, ptr)) return;  /* already marked */
+    if (LIKELY(objspace->mark_func_data == 0)) {
+	obj = RANY(ptr);
+	if (rb_special_const_p(ptr)) return; /* special const not marked */
+	if (obj->as.basic.flags == 0) return;       /* free cell */
+	if (!gc_mark_ptr(objspace, ptr)) return;  /* already marked */
+    }
+    else {
+	gc_mark(objspace, ptr);
+	return;
+    }
 
   marking:
     if (FL_TEST(obj, FL_EXIVAR)) {
@@ -2953,6 +2988,8 @@ rb_gc_unregister_address(VALUE *addr)
 static int
 garbage_collect(rb_objspace_t *objspace)
 {
+    struct mark_func_data_struct *prev_mark_func_data;
+
     if (GC_NOTIFY) printf("start garbage_collect()\n");
 
     if (!heaps) {
@@ -2964,6 +3001,9 @@ garbage_collect(rb_objspace_t *objspace)
 
     gc_prof_timer_start(objspace);
 
+    prev_mark_func_data = objspace->mark_func_data;
+    objspace->mark_func_data = 0;
+
     rest_sweep(objspace);
 
     during_gc++;
@@ -2973,6 +3013,8 @@ garbage_collect(rb_objspace_t *objspace)
     gc_sweep(objspace);
     gc_prof_sweep_timer_stop(objspace);
 
+    objspace->mark_func_data = prev_mark_func_data;
+
     gc_prof_timer_stop(objspace, Qtrue);
     if (GC_NOTIFY) printf("end garbage_collect()\n");
     return TRUE;
@@ -3241,6 +3283,46 @@ rb_gc_set_params(void)
     }
 }
 
+static void
+collect_refs(rb_objspace_t *objspace, VALUE obj)
+{
+    if (markable_object_p(objspace, obj) && !internal_object_p(obj)) {
+	st_insert((st_table *)objspace->mark_func_data->data, obj, Qtrue);
+    }
+}
+
+static int
+collect_keys(st_data_t key, st_data_t value, st_data_t data)
+{
+    VALUE ary = (VALUE)data;
+    rb_ary_push(ary, (VALUE)key);
+    return ST_CONTINUE;
+}
+
+VALUE
+rb_objspace_reachable_objects_from(VALUE obj)
+{
+    rb_objspace_t *objspace = &rb_objspace;
+
+    if (markable_object_p(objspace, obj)) {
+	st_table *refs = st_init_numtable();
+	struct mark_func_data_struct mfd;
+	VALUE ret = rb_ary_new();
+	mfd.mark_func = collect_refs;
+	mfd.data = (VALUE)refs;
+	objspace->mark_func_data = &mfd;
+
+	gc_mark_children(objspace, obj);
+
+	objspace->mark_func_data = 0;
+	st_foreach(refs, collect_keys, (st_data_t)ret);
+	return ret;
+    }
+    else {
+	return Qnil;
+    }
+}
+
 /*
   ------------------------ Extended allocator ------------------------
 */
diff --git a/gc.h b/gc.h
index 96d22a6e09..741945349d 100644
--- a/gc.h
+++ b/gc.h
@@ -87,7 +87,10 @@ int ruby_get_stack_grow_direction(volatile VALUE *addr);
 #pragma GCC visibility push(default)
 #endif
 
+/* exports for objspace module */
 size_t rb_objspace_data_type_memsize(VALUE obj);
+VALUE rb_objspace_reachable_objects_from(VALUE obj);
+
 void rb_objspace_each_objects(
     int (*callback)(void *start, void *end, size_t stride, void *data),
     void *data);
diff --git a/test/objspace/test_objspace.rb b/test/objspace/test_objspace.rb
index 345b00f737..a18923137e 100644
--- a/test/objspace/test_objspace.rb
+++ b/test/objspace/test_objspace.rb
@@ -68,4 +68,30 @@ class TestObjSpace < Test::Unit::TestCase
     ObjectSpace.count_tdata_objects(arg)
     assert_equal(false, arg.empty?)
   end
+
+  def test_reachable_objects_from
+    assert_equal(nil, ObjectSpace.reachable_objects_from(nil))
+    assert_equal([Array, 'a', 'b', 'c'], ObjectSpace.reachable_objects_from(['a', 'b', 'c']))
+
+    assert_equal([Array, 'a', 'a', 'a'], ObjectSpace.reachable_objects_from(['a', 'a', 'a']))
+    assert_equal([Array, 'a', 'a'], ObjectSpace.reachable_objects_from(['a', v = 'a', v]))
+    assert_equal([Array, 'a'], ObjectSpace.reachable_objects_from([v = 'a', v, v]))
+
+    long_ary = Array.new(1_000){''}
+    max = 0
+
+    ObjectSpace.each_object{|o|
+      refs = ObjectSpace.reachable_objects_from(o)
+      max = [refs.size, max].max
+
+      unless refs.nil?
+        refs.each{|ro|
+          # check this referenced object is not internal object
+          assert_equal(false, ro.nil?)
+        }
+      end
+    }
+    STDERR.puts max
+    assert_equal(true, max >= 1_001) # 1000 elems + Array class
+  end
 end