Merge branch 'values_at' of https://github.com/take-cheeze/mruby into take-cheeze-values_at

1 files changed, 28 insertions, 0 deletions

diff --git a/src/range.c b/src/range.c
index 4ab6708e1..627b572d4 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"))
 
@@ -371,6 +372,33 @@ 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 (mrb_range_beg_len(mrb, argv[i], &beg, &len, olen)) {
+      mrb_int const end = RARRAY_LEN(obj) < beg + len ? RARRAY_LEN(obj) : 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_ARGUMENT_ERROR, "invalid values selector: %S", argv[i]);
+    }
+  }
+
+  return result;
+}
+
 void
 mrb_init_range(mrb_state *mrb)
 {