Merge pull request #1183 from h2so5/proc-arity

Add Proc#arity

2 files changed, 27 insertions, 0 deletions

diff --git a/src/proc.c b/src/proc.c
index c32ce9c7a..8ddb06695 100644
--- a/src/proc.c
+++ b/src/proc.c
@@ -131,6 +131,23 @@ mrb_proc_iseq(mrb_state *mrb, struct RProc *p)
   return p->body.irep->iseq;
 }
 
+/* 15.2.17.4.2 */
+static mrb_value
+mrb_proc_arity(mrb_state *mrb, mrb_value self)
+{
+  struct RProc *p = mrb_proc_ptr(self);
+  mrb_code *iseq = mrb_proc_iseq(mrb, p);
+  mrb_aspec aspec = *iseq >> 6;
+  int ma, ra, pa, arity;
+  
+  ma = ARGS_GETREQ(aspec);
+  ra = ARGS_GETREST(aspec);
+  pa = ARGS_GETPOST(aspec);
+  arity = ra ? -(ma + pa + 1) : ma + pa;
+
+  return mrb_fixnum_value(arity);
+}
+
 /* 15.3.1.2.6  */
 /* 15.3.1.3.27 */
 /*
@@ -181,6 +198,7 @@ mrb_init_proc(mrb_state *mrb)
 
   mrb_define_method(mrb, mrb->proc_class, "initialize", mrb_proc_initialize, ARGS_NONE());
   mrb_define_method(mrb, mrb->proc_class, "initialize_copy", mrb_proc_init_copy, ARGS_REQ(1));
+  mrb_define_method(mrb, mrb->proc_class, "arity", mrb_proc_arity, ARGS_NONE());
 
   m = mrb_proc_new(mrb, call_irep);
   mrb_define_method_raw(mrb, mrb->proc_class, mrb_intern(mrb, "call"), m);
diff --git a/test/t/proc.rb b/test/t/proc.rb
index c0a1cf90f..4d5e5cab4 100644
--- a/test/t/proc.rb
+++ b/test/t/proc.rb
@@ -35,6 +35,15 @@ assert('Proc#[]', '15.2.17.4.1') do
   a == 1 and a2 == 5
 end
 
+assert('Proc#arity', '15.2.17.4.2') do
+  a = Proc.new {|x, y|}.arity
+  b = Proc.new {|x, *y, z|}.arity
+  c = Proc.new {|x=0, y|}.arity
+  d = Proc.new {|(x, y), z=0|}.arity
+
+  a == 2 and b == -3 and c == 1 and d == 1
+end
+
 assert('Proc#call', '15.2.17.4.3') do
   a = 0
   b = Proc.new { a += 1 }