diff options
Diffstat (limited to 'array.c')
-rw-r--r-- | array.c | 79 |
1 files changed, 79 insertions, 0 deletions
@@ -1983,6 +1983,84 @@ rb_ary_delete_if(VALUE ary) return ary; } +static VALUE +take_i(VALUE val, VALUE *args, int argc, VALUE *argv) +{ + if (args[1]-- == 0) rb_iter_break(); + if (argc > 1) val = rb_ary_new4(argc, argv); + rb_ary_push(args[0], val); + return Qnil; +} + +static VALUE +take_items(VALUE obj, long n) +{ + VALUE result = rb_ary_new2(n); + VALUE args[2]; + + args[0] = result; args[1] = (VALUE)n; + rb_block_call(obj, rb_intern("each"), 0, 0, take_i, (VALUE)args); + return result; +} + + +/* + * call-seq: + * array.zip(arg, ...) -> an_array + * array.zip(arg, ...) {| arr | block } -> nil + * + * Converts any arguments to arrays, then merges elements of + * <i>self</i> with corresponding elements from each argument. This + * generates a sequence of <code>self.size</code> <em>n</em>-element + * arrays, where <em>n</em> is one more that the count of arguments. If + * the size of any argument is less than <code>enumObj.size</code>, + * <code>nil</code> values are supplied. If a block given, it is + * invoked for each output array, otherwise an array of arrays is + * returned. + * + * a = [ 4, 5, 6 ] + * b = [ 7, 8, 9 ] + * + * [1,2,3].zip(a, b) #=> [[1, 4, 7], [2, 5, 8], [3, 6, 9]] + * [1,2].zip(a,b) #=> [[1, 4, 7], [2, 5, 8]] + * a.zip([1,2],[8]) #=> [[4,1,8], [5,2,nil], [6,nil,nil]] + */ + +static VALUE +rb_ary_zip(argc, argv, ary) + int argc; + VALUE *argv; + VALUE ary; +{ + int i, j; + long len; + VALUE result = Qnil; + + len = RARRAY_LEN(ary); + for (i=0; i<argc; i++) { + argv[i] = take_items(argv[i], len); + } + if (!rb_block_given_p()) { + result = rb_ary_new2(len); + } + + for (i=0; i<RARRAY_LEN(ary); i++) { + VALUE tmp = rb_ary_new2(argc+1); + + rb_ary_push(tmp, rb_ary_elt(ary, i)); + for (j=0; j<argc; j++) { + rb_ary_push(tmp, rb_ary_elt(argv[j], i)); + } + if (NIL_P(result)) { + rb_yield(tmp); + } + else { + rb_ary_push(result, tmp); + } + } + return result; +} + /* * call-seq: * array.transpose -> an_array @@ -3299,6 +3377,7 @@ Init_Array(void) rb_define_method(rb_cArray, "delete_if", rb_ary_delete_if, 0); rb_define_method(rb_cArray, "reject", rb_ary_reject, 0); rb_define_method(rb_cArray, "reject!", rb_ary_reject_bang, 0); + rb_define_method(rb_cArray, "zip", rb_ary_zip, -1); rb_define_method(rb_cArray, "transpose", rb_ary_transpose, 0); rb_define_method(rb_cArray, "replace", rb_ary_replace, 1); rb_define_method(rb_cArray, "clear", rb_ary_clear, 0); |