Skip to content

Commit d798aa5

Browse files
marcandremarcandre
committed
* enumerator: New method #size; constructor accepts size
* include/ruby/intern.h: RETURN_SIZED_ENUMERATOR for support of sized enumerators
1 parent 2228421 commit d798aa5

File tree

3 files changed

+74
-14
lines changed

3 files changed

+74
-14
lines changed

‎enumerator.c

Lines changed: 60 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,8 @@ struct enumerator {
119119
VALUE lookahead;
120120
VALUE feedvalue;
121121
VALUE stop_exc;
122+
VALUE size;
123+
VALUE (*size_fn)(ANYARGS);
122124
};
123125

124126
static VALUE rb_cGenerator, rb_cYielder;
@@ -148,6 +150,7 @@ enumerator_mark(void *p)
148150
rb_gc_mark(ptr->lookahead);
149151
rb_gc_mark(ptr->feedvalue);
150152
rb_gc_mark(ptr->stop_exc);
153+
rb_gc_mark(ptr->size);
151154
}
152155

153156
#define enumerator_free RUBY_TYPED_DEFAULT_FREE
@@ -216,7 +219,7 @@ obj_to_enum(int argc, VALUE *argv, VALUE obj)
216219
--argc;
217220
meth = *argv++;
218221
}
219-
return rb_enumeratorize(obj, meth, argc, argv);
222+
return rb_enumeratorize(obj, meth, argc, argv, 0);
220223
}
221224

222225
static VALUE
@@ -232,7 +235,7 @@ enumerator_allocate(VALUE klass)
232235
}
233236

234237
static VALUE
235-
enumerator_init(VALUE enum_obj, VALUE obj, VALUE meth, int argc, VALUE *argv)
238+
enumerator_init(VALUE enum_obj, VALUE obj, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS), VALUE size)
236239
{
237240
struct enumerator *ptr;
238241

@@ -250,13 +253,15 @@ enumerator_init(VALUE enum_obj, VALUE obj, VALUE meth, int argc, VALUE *argv)
250253
ptr->lookahead = Qundef;
251254
ptr->feedvalue = Qundef;
252255
ptr->stop_exc = Qfalse;
256+
ptr->size = size;
257+
ptr->size_fn = size_fn;
253258

254259
return enum_obj;
255260
}
256261

257262
/*
258263
* call-seq:
259-
* Enumerator.new { |yielder| ... }
264+
* Enumerator.new(size = nil) { |yielder| ... }
260265
* Enumerator.new(obj, method = :each, *args)
261266
*
262267
* Creates a new Enumerator object, which can be used as an
@@ -276,6 +281,10 @@ enumerator_init(VALUE enum_obj, VALUE obj, VALUE meth, int argc, VALUE *argv)
276281
*
277282
* p fib.take(10) # => [1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
278283
*
284+
* The optional parameter can be used to specify how to calculate the size
285+
* in a lazy fashion (see Enumerable#size). It can either be a value or
286+
* a callable object.
287+
*
279288
* The block form can be used to create a lazy enumeration that only processes
280289
* elements as-needed. The generic pattern for this is:
281290
*
@@ -349,14 +358,23 @@ static VALUE
349358
enumerator_initialize(int argc, VALUE *argv, VALUE obj)
350359
{
351360
VALUE recv, meth = sym_each;
361+
VALUE size = Qnil;
352362

353-
if (argc == 0) {
354-
if (!rb_block_given_p())
355-
rb_check_arity(argc, 1, UNLIMITED_ARGUMENTS);
356-
363+
if (rb_block_given_p()) {
364+
rb_check_arity(argc, 0, 1);
357365
recv = generator_init(generator_allocate(rb_cGenerator), rb_block_proc());
366+
if (argc) {
367+
if (NIL_P(argv[0]) || rb_obj_is_proc(argv[0]) ||
368+
(TYPE(argv[0]) == T_FLOAT && RFLOAT_VALUE(argv[0]) == INFINITY)) {
369+
size = argv[0];
370+
} else {
371+
size = rb_to_int(argv[0]);
372+
}
373+
argc = 0;
374+
}
358375
}
359376
else {
377+
rb_check_arity(argc, 1, UNLIMITED_ARGUMENTS);
360378
rb_warn("Enumerator.new without a block is deprecated; use Object#to_enum");
361379
recv = *argv++;
362380
if (--argc) {
@@ -365,7 +383,7 @@ enumerator_initialize(int argc, VALUE *argv, VALUE obj)
365383
}
366384
}
367385

368-
return enumerator_init(obj, recv, meth, argc, argv);
386+
return enumerator_init(obj, recv, meth, argc, argv, 0, size);
369387
}
370388

371389
/* :nodoc: */
@@ -393,14 +411,16 @@ enumerator_init_copy(VALUE obj, VALUE orig)
393411
ptr1->fib = 0;
394412
ptr1->lookahead = Qundef;
395413
ptr1->feedvalue = Qundef;
414+
ptr1->size = ptr0->size;
415+
ptr1->size_fn = ptr0->size_fn;
396416

397417
return obj;
398418
}
399419

400420
VALUE
401-
rb_enumeratorize(VALUE obj, VALUE meth, int argc, VALUE *argv)
421+
rb_enumeratorize(VALUE obj, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS))
402422
{
403-
return enumerator_init(enumerator_allocate(rb_cEnumerator), obj, meth, argc, argv);
423+
return enumerator_init(enumerator_allocate(rb_cEnumerator), obj, meth, argc, argv, size_fn, Qnil);
404424
}
405425

406426
static VALUE
@@ -942,6 +962,34 @@ enumerator_inspect(VALUE obj)
942962
return rb_exec_recursive(inspect_enumerator, obj, 0);
943963
}
944964

965+
/*
966+
* call-seq:
967+
* e.size -> int, Float::INFINITY or nil
968+
*
969+
* Returns the size of the enumerator, or +nil+ if it can't be calculated lazily.
970+
*
971+
* (1..100).to_a.permutation(4).size # => 94109400
972+
* loop.size # => Float::INFINITY
973+
* (1..100).drop_while.size # => nil
974+
*/
975+
976+
static VALUE
977+
enumerator_size(VALUE obj)
978+
{
979+
struct enumerator *e = enumerator_ptr(obj);
980+
981+
if (e->size_fn) {
982+
return (*e->size_fn)(e->obj, e->args);
983+
}
984+
if (rb_obj_is_proc(e->size)) {
985+
if(e->args)
986+
return rb_proc_call(e->size, e->args);
987+
else
988+
return rb_proc_call_with_block(e->size, 0, 0, Qnil);
989+
}
990+
return e->size;
991+
}
992+
945993
/*
946994
* Yielder
947995
*/
@@ -1253,7 +1301,7 @@ lazy_initialize(int argc, VALUE *argv, VALUE self)
12531301
rb_block_call(generator, id_initialize, 0, 0,
12541302
(rb_block_given_p() ? lazy_init_block_i : lazy_init_block),
12551303
obj);
1256-
enumerator_init(self, generator, meth, argc - offset, argv + offset);
1304+
enumerator_init(self, generator, meth, argc - offset, argv + offset, 0, Qnil);
12571305
rb_ivar_set(self, id_receiver, obj);
12581306

12591307
return self;
@@ -1749,6 +1797,7 @@ InitVM_Enumerator(void)
17491797
rb_define_method(rb_cEnumerator, "feed", enumerator_feed, 1);
17501798
rb_define_method(rb_cEnumerator, "rewind", enumerator_rewind, 0);
17511799
rb_define_method(rb_cEnumerator, "inspect", enumerator_inspect, 0);
1800+
rb_define_method(rb_cEnumerator, "size", enumerator_size, 0);
17521801

17531802
/* Lazy */
17541803
rb_cLazy = rb_define_class_under(rb_cEnumerator, "Lazy", rb_cEnumerator);

‎include/ruby/intern.h

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,12 +200,13 @@ VALUE rb_fiber_alive_p(VALUE);
200200
/* enum.c */
201201
VALUE rb_enum_values_pack(int, VALUE*);
202202
/* enumerator.c */
203-
VALUE rb_enumeratorize(VALUE, VALUE, int, VALUE *);
204-
#define RETURN_ENUMERATOR(obj, argc, argv) do { \
203+
VALUE rb_enumeratorize(VALUE, VALUE, int, VALUE *, VALUE (*)(ANYARGS));
204+
#define RETURN_SIZED_ENUMERATOR(obj, argc, argv, size_fn) do { \
205205
if (!rb_block_given_p()) \
206206
return rb_enumeratorize((obj), ID2SYM(rb_frame_this_func()),\
207-
(argc), (argv)); \
207+
(argc), (argv), (size_fn)); \
208208
} while (0)
209+
#define RETURN_ENUMERATOR(obj, argc, argv) RETURN_SIZED_ENUMERATOR(obj, argc, argv, 0)
209210
/* error.c */
210211
VALUE rb_exc_new(VALUE, const char*, long);
211212
VALUE rb_exc_new2(VALUE, const char*);

‎test/ruby/test_enumerator.rb

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -401,5 +401,15 @@ def test_yielder
401401

402402
assert_raise(LocalJumpError) { Enumerator::Yielder.new }
403403
end
404+
405+
def test_size
406+
assert_equal nil, Enumerator.new{}.size
407+
assert_equal 42, Enumerator.new(->{42}){}.size
408+
assert_equal 42, Enumerator.new(42){}.size
409+
assert_equal 1 << 70, Enumerator.new(1 << 70){}.size
410+
assert_equal Float::INFINITY, Enumerator.new(Float::INFINITY){}.size
411+
assert_equal nil, Enumerator.new(nil){}.size
412+
assert_raise(TypeError) { Enumerator.new("42"){} }
413+
end
404414
end
405415

0 commit comments

Comments
 (0)