Merge branch 'take-cheeze-values_at'

7 files changed, 115 insertions, 37 deletions

diff --git a/include/mruby.h b/include/mruby.h
index 9cf578836..db3b06aa8 100644
--- a/include/mruby.h
+++ b/include/mruby.h
@@ -381,6 +381,7 @@ mrb_value mrb_yield_with_class(mrb_state *mrb, mrb_value b, mrb_int argc, const
 
 void mrb_gc_protect(mrb_state *mrb, mrb_value obj);
 mrb_value mrb_to_int(mrb_state *mrb, mrb_value val);
+#define mrb_int(mrb, val) mrb_fixnum(mrb_to_int(mrb, val))
 void mrb_check_type(mrb_state *mrb, mrb_value x, enum mrb_vtype t);
 
 typedef enum call_type {
diff --git a/include/mruby/range.h b/include/mruby/range.h
index 828ec2691..61beb2319 100644
--- a/include/mruby/range.h
+++ b/include/mruby/range.h
@@ -27,6 +27,7 @@ struct RRange {
 
 mrb_value mrb_range_new(mrb_state*, mrb_value, mrb_value, mrb_bool);
 mrb_bool mrb_range_beg_len(mrb_state *mrb, mrb_value range, mrb_int *begp, mrb_int *lenp, mrb_int len);
+mrb_value mrb_get_values_at(mrb_state *mrb, mrb_value obj, mrb_int olen, mrb_int argc, const mrb_value *argv, mrb_value (*func)(mrb_state*, mrb_value, mrb_int));
 
 #if defined(__cplusplus)
 }  /* extern "C" { */
diff --git a/mrbgems/mruby-array-ext/src/array.c b/mrbgems/mruby-array-ext/src/array.c
index ad6cead2c..ae9d8296e 100644
--- a/mrbgems/mruby-array-ext/src/array.c
+++ b/mrbgems/mruby-array-ext/src/array.c
@@ -1,6 +1,7 @@
 #include "mruby.h"
 #include "mruby/value.h"
 #include "mruby/array.h"
+#include "mruby/range.h"
 
 /*
  *  call-seq:
@@ -122,6 +123,17 @@ mrb_ary_at(mrb_state *mrb, mrb_value ary)
   return mrb_ary_entry(ary, pos);
 }
 
+static mrb_value
+mrb_ary_values_at(mrb_state *mrb, mrb_value self)
+{
+  mrb_int argc;
+  mrb_value *argv;
+
+  mrb_get_args(mrb, "*", &argv, &argc);
+
+  return mrb_get_values_at(mrb, self, RARRAY_LEN(self), argc, argv, mrb_ary_ref);
+}
+
 void
 mrb_mruby_array_ext_gem_init(mrb_state* mrb)
 {
@@ -132,6 +144,7 @@ mrb_mruby_array_ext_gem_init(mrb_state* mrb)
   mrb_define_method(mrb, a, "assoc",  mrb_ary_assoc,  MRB_ARGS_REQ(1));
   mrb_define_method(mrb, a, "at",     mrb_ary_at,     MRB_ARGS_REQ(1));
   mrb_define_method(mrb, a, "rassoc", mrb_ary_rassoc, MRB_ARGS_REQ(1));
+  mrb_define_method(mrb, a, "values_at", mrb_ary_values_at, MRB_ARGS_ANY());
 }
 
 void
diff --git a/mrbgems/mruby-array-ext/test/array.rb b/mrbgems/mruby-array-ext/test/array.rb
index b0ad94b38..1fa7cfc04 100644
--- a/mrbgems/mruby-array-ext/test/array.rb
+++ b/mrbgems/mruby-array-ext/test/array.rb
@@ -280,3 +280,13 @@ assert("Array#select!") do
   assert_equal [4, 5], a.select! { |val| val > 3 }
   assert_equal [4, 5], a
 end
+
+assert('Array#values_at') do
+  a = %w{red green purple white none}
+
+  assert_equal %w{red purple none}, a.values_at(0, 2, 4)
+  assert_equal ['green', 'white', nil, nil], a.values_at(1, 3, 5, 7)
+  assert_equal ['none', 'white', 'white', nil], a.values_at(-1, -2, -2, -7)
+  assert_equal ['none', nil, nil, 'red', 'green', 'purple'], a.values_at(4..6, 0...3)
+  assert_raise(TypeError) { a.values_at 'tt' }
+end
diff --git a/mrbgems/mruby-struct/src/struct.c b/mrbgems/mruby-struct/src/struct.c
index a4d70ae1a..a15655dbb 100644
--- a/mrbgems/mruby-struct/src/struct.c
+++ b/mrbgems/mruby-struct/src/struct.c
@@ -12,6 +12,7 @@
 #include "mruby/class.h"
 #include "mruby/variable.h"
 #include "mruby/hash.h"
+#include "mruby/range.h"
 
 #define RSTRUCT_LEN(st) RARRAY_LEN(st)
 #define RSTRUCT_PTR(st) RARRAY_PTR(st)
@@ -531,7 +532,7 @@ mrb_struct_init_copy(mrb_state *mrb, mrb_value copy)
 }
 
 static mrb_value
-mrb_struct_aref_id(mrb_state *mrb, mrb_value s, mrb_sym id)
+struct_aref_sym(mrb_state *mrb, mrb_value s, mrb_sym id)
 {
   mrb_value *ptr, members, *ptr_members;
   mrb_int i, len;
@@ -549,6 +550,21 @@ mrb_struct_aref_id(mrb_state *mrb, mrb_value s, mrb_sym id)
   return mrb_nil_value();       /* not reached */
 }
 
+static mrb_value
+struct_aref_int(mrb_state *mrb, mrb_value s, mrb_int i)
+{
+  if (i < 0) i = RSTRUCT_LEN(s) + i;
+  if (i < 0)
+      mrb_raisef(mrb, E_INDEX_ERROR,
+                 "offset %S too small for struct(size:%S)",
+                 mrb_fixnum_value(i), mrb_fixnum_value(RSTRUCT_LEN(s)));
+  if (RSTRUCT_LEN(s) <= i)
+    mrb_raisef(mrb, E_INDEX_ERROR,
+               "offset %S too large for struct(size:%S)",
+               mrb_fixnum_value(i), mrb_fixnum_value(RSTRUCT_LEN(s)));
+  return RSTRUCT_PTR(s)[i];
+}
+
 /* 15.2.18.4.2  */
 /*
  *  call-seq:
@@ -569,10 +585,11 @@ mrb_struct_aref_id(mrb_state *mrb, mrb_value s, mrb_sym id)
  *     joe[0]        #=> "Joe Smith"
  */
 mrb_value
-mrb_struct_aref_n(mrb_state *mrb, mrb_value s, mrb_value idx)
+mrb_struct_aref(mrb_state *mrb, mrb_value s)
 {
-  mrb_int i;
+  mrb_value idx;
 
+  mrb_get_args(mrb, "o", &idx);
   if (mrb_string_p(idx)) {
     mrb_value sym = mrb_check_intern_str(mrb, idx);
 
@@ -582,33 +599,13 @@ mrb_struct_aref_n(mrb_state *mrb, mrb_value s, mrb_value idx)
     idx = sym;
   }
   if (mrb_symbol_p(idx)) {
-    return mrb_struct_aref_id(mrb, s, mrb_symbol(idx));
+    return struct_aref_sym(mrb, s, mrb_symbol(idx));
   }
-
-  i = mrb_fixnum(idx);
-  if (i < 0) i = RSTRUCT_LEN(s) + i;
-  if (i < 0)
-      mrb_raisef(mrb, E_INDEX_ERROR,
-                 "offset %S too small for struct(size:%S)",
-                 mrb_fixnum_value(i), mrb_fixnum_value(RSTRUCT_LEN(s)));
-  if (RSTRUCT_LEN(s) <= i)
-    mrb_raisef(mrb, E_INDEX_ERROR,
-               "offset %S too large for struct(size:%S)",
-               mrb_fixnum_value(i), mrb_fixnum_value(RSTRUCT_LEN(s)));
-  return RSTRUCT_PTR(s)[i];
-}
-
-mrb_value
-mrb_struct_aref(mrb_state *mrb, mrb_value s)
-{
-  mrb_value idx;
-
-  mrb_get_args(mrb, "o", &idx);
-  return mrb_struct_aref_n(mrb, s, idx);
+  return struct_aref_int(mrb, s, mrb_int(mrb, idx));
 }
 
 static mrb_value
-mrb_struct_aset_id(mrb_state *mrb, mrb_value s, mrb_sym id, mrb_value val)
+mrb_struct_aset_sym(mrb_state *mrb, mrb_value s, mrb_sym id, mrb_value val)
 {
   mrb_value members, *ptr, *ptr_members;
   mrb_int i, len;
@@ -663,8 +660,8 @@ mrb_struct_aset(mrb_state *mrb, mrb_value s)
 
   mrb_get_args(mrb, "oo", &idx, &val);
 
-  if (mrb_string_p(idx) || mrb_symbol_p(idx)) {
-    return mrb_struct_aset_id(mrb, s, mrb_obj_to_sym(mrb, idx), val);
+  if (mrb_symbol_p(idx)) {
+    return mrb_struct_aset_sym(mrb, s, mrb_symbol(idx), val);
   }
 
   i = mrb_fixnum(idx);
@@ -828,6 +825,17 @@ mrb_struct_to_h(mrb_state *mrb, mrb_value self)
   return ret;
 }
 
+static mrb_value
+mrb_struct_values_at(mrb_state *mrb, mrb_value self)
+{
+  mrb_int argc;
+  mrb_value *argv;
+
+  mrb_get_args(mrb, "*", &argv, &argc);
+
+  return mrb_get_values_at(mrb, self, RSTRUCT_LEN(self), argc, argv, struct_aref_int);
+}
+
 /*
  *  A <code>Struct</code> is a convenient way to bundle a number of
  *  attributes together, using accessor methods, without having to write
@@ -866,6 +874,7 @@ mrb_mruby_struct_gem_init(mrb_state* mrb)
   mrb_define_method(mrb, st,        "to_a",           mrb_struct_to_a,        MRB_ARGS_NONE());
   mrb_define_method(mrb, st,        "values",         mrb_struct_to_a,        MRB_ARGS_NONE());
   mrb_define_method(mrb, st,        "to_h",           mrb_struct_to_h,        MRB_ARGS_NONE());
+  mrb_define_method(mrb, st,        "values_at",      mrb_struct_values_at,   MRB_ARGS_NONE());
 }
 
 void
diff --git a/mrbgems/mruby-struct/test/struct.rb b/mrbgems/mruby-struct/test/struct.rb
index d654830ca..b3b9ce33f 100644
--- a/mrbgems/mruby-struct/test/struct.rb
+++ b/mrbgems/mruby-struct/test/struct.rb
@@ -123,3 +123,10 @@ assert('Struct#to_h') do
   s = Struct.new(:white, :red, :green).new('ruuko', 'yuzuki', 'hitoe')
   assert_equal(:white => 'ruuko', :red => 'yuzuki', :green => 'hitoe') { s.to_h }
 end
+
+assert('Struct#values_at') do
+  a = Struct.new(:blue, :purple).new('aki', 'io')
+  assert_equal ['aki'], a.values_at(0)
+  assert_equal ['io', 'aki'], a.values_at(1, 0)
+  assert_raise(IndexError) { a.values_at 2 }
+end
diff --git a/src/range.c b/src/range.c
index 4ab6708e1..3e5af1894 100644
--- a/src/range.c
+++ b/src/range.c
@@ -8,6 +8,7 @@
 #include "mruby/class.h"
 #include "mruby/range.h"
 #include "mruby/string.h"
+#include "mruby/array.h"
 
 #define RANGE_CLASS (mrb_class_get(mrb, "Range"))
 
@@ -234,28 +235,29 @@ mrb_range_include(mrb_state *mrb, mrb_value range)
 }
 
 mrb_bool
-mrb_range_beg_len(mrb_state *mrb, mrb_value range, mrb_int *begp, mrb_int *lenp, mrb_int len)
+range_beg_len(mrb_state *mrb, mrb_value range, mrb_int *begp, mrb_int *lenp, mrb_int len, mrb_bool trunc)
 {
   mrb_int beg, end, b, e;
   struct RRange *r = mrb_range_ptr(range);
 
-  if (mrb_type(range) != MRB_TT_RANGE) {
-    mrb_raise(mrb, E_TYPE_ERROR, "expected Range.");
-  }
+  if (mrb_type(range) != MRB_TT_RANGE) return FALSE;
 
-  beg = b = mrb_fixnum(r->edges->beg);
-  end = e = mrb_fixnum(r->edges->end);
+  beg = b = mrb_int(mrb, r->edges->beg);
+  end = e = mrb_int(mrb, r->edges->end);
 
   if (beg < 0) {
     beg += len;
     if (beg < 0) return FALSE;
   }
 
-  if (beg > len) return FALSE;
-  if (end > len) end = len;
+  if (trunc) {
+    if (beg > len) return FALSE;
+    if (end > len) end = len;
+  }
 
   if (end < 0) end += len;
-  if (!r->excl && end < len) end++;  /* include end point */
+  if (!r->excl && (!trunc || end < len))
+    end++;                      /* include end point */
   len = end - beg;
   if (len < 0) len = 0;
 
@@ -264,6 +266,12 @@ mrb_range_beg_len(mrb_state *mrb, mrb_value range, mrb_int *begp, mrb_int *lenp,
   return TRUE;
 }
 
+mrb_bool
+mrb_range_beg_len(mrb_state *mrb, mrb_value range, mrb_int *begp, mrb_int *lenp, mrb_int len)
+{
+  return range_beg_len(mrb, range, begp, lenp, len, TRUE);
+}
+
 /* 15.2.14.4.12(x) */
 /*
  * call-seq:
@@ -371,6 +379,35 @@ range_initialize_copy(mrb_state *mrb, mrb_value copy)
   return copy;
 }
 
+mrb_value
+mrb_get_values_at(mrb_state *mrb, mrb_value obj, mrb_int olen, mrb_int argc, const mrb_value *argv, mrb_value (*func)(mrb_state*, mrb_value, mrb_int))
+{
+  mrb_int i, j, beg, len;
+  mrb_value result;
+  result = mrb_ary_new(mrb);
+
+  for (i = 0; i < argc; ++i) {
+    if (mrb_fixnum_p(argv[i])) {
+      mrb_ary_push(mrb, result, func(mrb, obj, mrb_fixnum(argv[i])));
+    }
+    else if (range_beg_len(mrb, argv[i], &beg, &len, olen, FALSE)) {
+      mrb_int const end = olen < beg + len ? olen : beg + len;
+      for (j = beg; j < end; ++j) {
+        mrb_ary_push(mrb, result, func(mrb, obj, j));
+      }
+
+      for (; j < beg + len; ++j) {
+        mrb_ary_push(mrb, result, mrb_nil_value());
+      }
+    }
+    else {
+      mrb_raisef(mrb, E_TYPE_ERROR, "invalid values selector: %S", argv[i]);
+    }
+  }
+
+  return result;
+}
+
 void
 mrb_init_range(mrb_state *mrb)
 {