diff options
| -rw-r--r-- | benchmark/enum_lazy_flat_map.yml | 16 | ||||
| -rw-r--r-- | enumerator.c | 82 | ||||
| -rw-r--r-- | test/ruby/test_lazy_enumerator.rb | 4 |
3 files changed, 60 insertions, 42 deletions
diff --git a/benchmark/enum_lazy_flat_map.yml b/benchmark/enum_lazy_flat_map.yml new file mode 100644 index 0000000000..0ee390a441 --- /dev/null +++ b/benchmark/enum_lazy_flat_map.yml @@ -0,0 +1,16 @@ +prelude: | + num = (1..).lazy.take(100) + ary2 = [[1,2]].cycle.lazy.take(10) + ary10 = [[*1..10]].cycle.lazy.take(10) + ary20 = [[*1..20]].cycle.lazy.take(10) + ary50 = [[*1..50]].cycle.lazy.take(10) + ary100 = [[*1..100]].cycle.lazy.take(10) + +benchmark: + num3: num.flat_map {|x| x}.take(3).to_a + num10: num.flat_map {|x| x}.take(3).to_a + ary2: ary2.flat_map {|x| x}.take(3).to_a + ary10: ary10.flat_map {|x| x}.take(3).to_a + ary20: ary20.flat_map {|x| x}.take(3).to_a + ary50: ary50.flat_map {|x| x}.take(3).to_a + ary100: ary100.flat_map {|x| x}.take(3).to_a diff --git a/enumerator.c b/enumerator.c index d10f239922..c560c163a9 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1605,6 +1605,7 @@ lazy_init_block_i(RB_BLOCK_CALL_FUNC_ARGLIST(val, m)) #define LAZY_MEMO_BREAK_P(memo) ((memo)->memo_flags & LAZY_MEMO_BREAK) #define LAZY_MEMO_PACKED_P(memo) ((memo)->memo_flags & LAZY_MEMO_PACKED) #define LAZY_MEMO_SET_BREAK(memo) ((memo)->memo_flags |= LAZY_MEMO_BREAK) +#define LAZY_MEMO_RESET_BREAK(memo) ((memo)->memo_flags &= ~LAZY_MEMO_BREAK) #define LAZY_MEMO_SET_VALUE(memo, value) MEMO_V2_SET(memo, value) #define LAZY_MEMO_SET_PACKED(memo) ((memo)->memo_flags |= LAZY_MEMO_PACKED) #define LAZY_MEMO_RESET_PACKED(memo) ((memo)->memo_flags &= ~LAZY_MEMO_PACKED) @@ -2058,58 +2059,57 @@ lazy_map(VALUE obj) return lazy_add_method(obj, 0, 0, Qnil, Qnil, &lazy_map_funcs); } +struct flat_map_i_arg { + struct MEMO *result; + long index; +}; + static VALUE -lazy_flat_map_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, yielder)) +lazy_flat_map_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, y)) { - VALUE arg = rb_enum_values_pack(argc, argv); + struct flat_map_i_arg *arg = (struct flat_map_i_arg *)y; - return rb_funcallv(yielder, idLTLT, 1, &arg); + return lazy_yielder_yield(arg->result, arg->index, argc, argv); } -static VALUE -lazy_flat_map_each(VALUE obj, VALUE yielder) +static struct MEMO * +lazy_flat_map_proc(VALUE proc_entry, struct MEMO *result, VALUE memos, long memo_index) { - rb_block_call(obj, id_each, 0, 0, lazy_flat_map_i, yielder); - return Qnil; -} + VALUE value = lazyenum_yield_values(proc_entry, result); + VALUE ary = 0; + const long proc_index = memo_index + 1; + int break_p = LAZY_MEMO_BREAK_P(result); -static VALUE -lazy_flat_map_to_ary(VALUE obj, VALUE yielder) -{ - VALUE ary = rb_check_array_type(obj); - if (NIL_P(ary)) { - rb_funcall(yielder, idLTLT, 1, obj); + if (RB_TYPE_P(value, T_ARRAY)) { + ary = value; } - else { - long i; - for (i = 0; i < RARRAY_LEN(ary); i++) { - rb_funcall(yielder, idLTLT, 1, RARRAY_AREF(ary, i)); - } + else if (rb_respond_to(value, id_force) && rb_respond_to(value, id_each)) { + struct flat_map_i_arg arg = {.result = result, .index = proc_index}; + LAZY_MEMO_RESET_BREAK(result); + rb_block_call(value, id_each, 0, 0, lazy_flat_map_i, (VALUE)&arg); + if (break_p) LAZY_MEMO_SET_BREAK(result); + return 0; } - return Qnil; -} -static VALUE -lazy_flat_map_proc(RB_BLOCK_CALL_FUNC_ARGLIST(val, m)) -{ - VALUE result = rb_yield_values2(argc - 1, &argv[1]); - if (RB_TYPE_P(result, T_ARRAY)) { - long i; - for (i = 0; i < RARRAY_LEN(result); i++) { - rb_funcall(argv[0], idLTLT, 1, RARRAY_AREF(result, i)); - } - } - else { - if (rb_respond_to(result, id_force) && rb_respond_to(result, id_each)) { - lazy_flat_map_each(result, argv[0]); - } - else { - lazy_flat_map_to_ary(result, argv[0]); - } + if (ary || !NIL_P(ary = rb_check_array_type(value))) { + long i; + LAZY_MEMO_RESET_BREAK(result); + for (i = 0; i + 1 < RARRAY_LEN(ary); i++) { + lazy_yielder_yield(result, proc_index, 1, &RARRAY_AREF(ary, i)); + } + if (break_p) LAZY_MEMO_SET_BREAK(result); + if (i >= RARRAY_LEN(ary)) return 0; + value = RARRAY_AREF(ary, i); } - return Qnil; + LAZY_MEMO_SET_VALUE(result, value); + LAZY_MEMO_RESET_PACKED(result); + return result; } +static const lazyenum_funcs lazy_flat_map_funcs = { + lazy_flat_map_proc, 0, +}; + /* * call-seq: * lazy.collect_concat { |obj| block } -> a_lazy_enumerator @@ -2140,9 +2140,7 @@ lazy_flat_map(VALUE obj) rb_raise(rb_eArgError, "tried to call lazy flat_map without a block"); } - return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj, - lazy_flat_map_proc, 0), - Qnil, 0); + return lazy_add_method(obj, 0, 0, Qnil, Qnil, &lazy_flat_map_funcs); } static struct MEMO * diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 6e5c1714a5..3f5a05555b 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -160,6 +160,10 @@ class TestLazyEnumerator < Test::Unit::TestCase assert_equal([{?a=>97}, {?b=>98}, {?c=>99}], [?a, ?b, ?c].lazy.flat_map {|x| {x=>x.ord}}.force) end + def test_flat_map_take + assert_equal([1,2]*3, [[1,2]].cycle.lazy.take(3).flat_map {|x| x}.to_a) + end + def test_reject a = Step.new(1..6) assert_equal(4, a.reject {|x| x < 4}.first) |