From f0af96c130737f8ed07eeed73691428cd55a346f Mon Sep 17 00:00:00 2001
From: "Yukihiro \"Matz\" Matsumoto" <matz@ruby.or.jp>
Date: Tue, 1 Feb 2022 22:37:49 +0900
Subject: [PATCH] mruby-class-ext: add `Class#subclasses` method.

---
 mrbgems/mruby-class-ext/src/class.c   | 55 +++++++++++++++++++++++++++
 mrbgems/mruby-class-ext/test/class.rb | 13 +++++++
 2 files changed, 68 insertions(+)
 create mode 100644 mrbgems/mruby-class-ext/test/class.rb

diff --git a/mrbgems/mruby-class-ext/src/class.c b/mrbgems/mruby-class-ext/src/class.c
index 290ddf4b3..373af8312 100644
--- a/mrbgems/mruby-class-ext/src/class.c
+++ b/mrbgems/mruby-class-ext/src/class.c
@@ -1,6 +1,7 @@
 #include "mruby.h"
 #include "mruby/class.h"
 #include "mruby/string.h"
+#include "mruby/array.h"
 #include "mruby/proc.h"
 
 static mrb_value
@@ -56,6 +57,57 @@ mrb_mod_module_exec(mrb_state *mrb, mrb_value self)
   return mrb_yield_cont(mrb, blk, self, argc, argv);
 }
 
+struct subclass_args {
+  struct RClass *c;
+  mrb_value ary;
+};
+
+static int
+add_subclasses(mrb_state *mrb, struct RBasic *obj, void *data)
+{
+  struct subclass_args *args = (struct subclass_args*)data;
+  if (obj->tt == MRB_TT_CLASS) {
+    struct RClass *c = (struct RClass*)obj;
+    if (mrb_class_real(c->super) == args->c) {
+      mrb_ary_push(mrb, args->ary, mrb_obj_value(obj));
+    }
+  }
+  return MRB_EACH_OBJ_OK;
+}
+
+/*
+ *  call-seq:
+ *     subclasses -> array
+ *
+ *  Returns an array of classes where the receiver is the
+ *  direct superclass of the class, excluding singleton classes.
+ *  The order of the returned array is not defined.
+ *
+ *     class A; end
+ *     class B < A; end
+ *     class C < B; end
+ *     class D < A; end
+ *
+ *     A.subclasses        #=> [D, B]
+ *     B.subclasses        #=> [C]
+ *     C.subclasses        #=> []
+ */
+static mrb_value
+mrb_class_subclasses(mrb_state *mrb, mrb_value self)
+{
+  struct RClass *c;
+  mrb_value ary;
+
+  c = mrb_class_ptr(self);
+  ary = mrb_ary_new(mrb);
+
+  if (c->flags & MRB_FL_CLASS_IS_INHERITED) {
+    struct subclass_args arg = {c, ary};
+    mrb_objspace_each_objects(mrb, add_subclasses, &arg);
+  }
+  return ary;
+}
+
 void
 mrb_mruby_class_ext_gem_init(mrb_state *mrb)
 {
@@ -65,6 +117,9 @@ mrb_mruby_class_ext_gem_init(mrb_state *mrb)
   mrb_define_method(mrb, mod, "singleton_class?", mrb_mod_singleton_class_p, MRB_ARGS_NONE());
   mrb_define_method(mrb, mod, "module_exec", mrb_mod_module_exec, MRB_ARGS_ANY()|MRB_ARGS_BLOCK());
   mrb_define_method(mrb, mod, "class_exec", mrb_mod_module_exec, MRB_ARGS_ANY()|MRB_ARGS_BLOCK());
+
+  struct RClass *cls = mrb->module_class;
+  mrb_define_method(mrb, cls, "subclasses", mrb_class_subclasses, MRB_ARGS_NONE());
 }
 
 void
diff --git a/mrbgems/mruby-class-ext/test/class.rb b/mrbgems/mruby-class-ext/test/class.rb
new file mode 100644
index 000000000..2043b39cb
--- /dev/null
+++ b/mrbgems/mruby-class-ext/test/class.rb
@@ -0,0 +1,13 @@
+assert 'Class#subclasses' do
+  a = Class.new
+  b = Class.new(a)
+  c = Class.new(b)
+  d = Class.new(a)
+
+  a_sub = a.subclasses
+  assert_equal(2, a_sub.size)
+  assert_true(a_sub.include?(b))
+  assert_true(a_sub.include?(d))
+  assert_equal([c], b.subclasses)
+  assert_equal([], c.subclasses)
+end