summaryrefslogtreecommitdiff
path: root/enumerator.c
diff options
context:
space:
mode:
authorJeremy Evans <code@jeremyevans.net>2019-08-08 13:09:17 -0700
committerJeremy Evans <code@jeremyevans.net>2019-09-03 11:30:49 -0700
commite94ac03eb0d07af3cbff20194acf479bf8851c39 (patch)
tree5a87f9a350b938f01408cfd6236d0a9a5fcda322 /enumerator.c
parentfb67d4fc77f211cfec6b2065787d9212f4e2154b (diff)
Make Enumerator::Lazy#with_index be lazy
Previously, Enumerator::Lazy#with_index was not defined, so it picked up the default implementation from Enumerator, which was not lazy. Based on earlier patch from nobu. Fixes [Bug #7877]
Notes
Notes: Merged: https://github.com/ruby/ruby/pull/2421
Diffstat (limited to 'enumerator.c')
-rw-r--r--enumerator.c66
1 files changed, 66 insertions, 0 deletions
diff --git a/enumerator.c b/enumerator.c
index 96daad2..d4022ae 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -3597,6 +3597,71 @@ arith_seq_size(VALUE self)
return len;
}
+static VALUE
+lazy_with_index_func(RB_BLOCK_CALL_FUNC_ARGLIST(val, offset))
+{
+ VALUE yielder, memo, result;
+ VALUE e = rb_enum_values_pack(argc - 1, argv + 1);
+ long idx;
+
+ yielder = argv[0];
+ memo = rb_attr_get(yielder, id_memo);
+ if (NIL_P(memo))
+ memo = offset;
+ idx = NUM2LONG(memo);
+ result = rb_assoc_new(e, memo);
+ rb_funcall(yielder, idLTLT, 1, result);
+ rb_ivar_set(yielder, id_memo, LONG2NUM(++idx));
+ return Qnil;
+}
+
+static VALUE
+lazy_with_index_iter(RB_BLOCK_CALL_FUNC_ARGLIST(val, offset))
+{
+ VALUE yielder, memo, result;
+ VALUE e = rb_enum_values_pack(argc - 1, argv + 1);
+ long idx;
+
+ yielder = argv[0];
+ memo = rb_attr_get(yielder, id_memo);
+ if (NIL_P(memo))
+ memo = offset;
+ idx = NUM2LONG(memo);
+ result = rb_yield(rb_assoc_new(e, memo));
+ rb_funcall(yielder, idLTLT, 1, result);
+ rb_ivar_set(yielder, id_memo, LONG2NUM(++idx));
+ return Qnil;
+}
+
+/*
+ * call-seq:
+ * lazy.with_index(offset = 0) {|(*args), idx| ... }
+ * lazy.with_index(offset = 0)
+ *
+ * Iterates the given block for each element with an index, which
+ * starts from +offset+. If no block is given, returns a new
+ * lazy enumerator that includes the index, starting from +offset+
+ *
+ * +offset+:: the starting index to use
+ *
+ * see Enumerator#with_index.
+ */
+static VALUE
+lazy_with_index(int argc, VALUE *argv, VALUE obj)
+{
+ VALUE memo;
+
+ rb_scan_args(argc, argv, "01", &memo);
+ if (NIL_P(memo))
+ memo = LONG2NUM(0);
+
+ return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
+ rb_block_given_p() ?
+ lazy_with_index_iter : lazy_with_index_func,
+ memo),
+ rb_ary_new_from_values(argc, argv), 0);
+}
+
void
InitVM_Enumerator(void)
{
@@ -3654,6 +3719,7 @@ InitVM_Enumerator(void)
rb_define_method(rb_cLazy, "slice_when", lazy_super, -1);
rb_define_method(rb_cLazy, "chunk_while", lazy_super, -1);
rb_define_method(rb_cLazy, "uniq", lazy_uniq, 0);
+ rb_define_method(rb_cLazy, "with_index", lazy_with_index, -1);
#if 0 /* for RDoc */
rb_define_method(rb_cLazy, "to_a", lazy_to_a, 0);