diff --git a/enum.c b/enum.c index da9f4348fa..e9cac3e092 100644 --- a/enum.c +++ b/enum.c @@ -4178,6 +4178,48 @@ enum_uniq(VALUE obj) return ret; } +static VALUE +compact_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, ary)) +{ + ENUM_WANT_SVALUE(); + + if (!NIL_P(i)) { + rb_ary_push(ary, i); + } + return Qnil; +} + +/* + * call-seq: + * enum.compact -> array + * + * Returns an array of all non-+nil+ elements from enumeration. + * + * def with_nils + * yield 1 + * yield 2 + * yield nil + * yield 3 + * end + * + * to_enum(:with_nils).compact + * # => [1, 2, 3] + * + * See also Array#compact. + */ + +static VALUE +enum_compact(VALUE obj) +{ + VALUE ary; + + ary = rb_ary_new(); + rb_block_call(obj, id_each, 0, 0, compact_i, ary); + + return ary; +} + + /* * The Enumerable mixin provides collection classes with several * traversal and searching methods, and with the ability to sort. The @@ -4251,6 +4293,7 @@ Init_Enumerable(void) rb_define_method(rb_mEnumerable, "chunk_while", enum_chunk_while, 0); rb_define_method(rb_mEnumerable, "sum", enum_sum, -1); rb_define_method(rb_mEnumerable, "uniq", enum_uniq, 0); + rb_define_method(rb_mEnumerable, "compact", enum_compact, 0); id_next = rb_intern_const("next"); } diff --git a/enumerator.c b/enumerator.c index b4a7cb588e..6d6a588af9 100644 --- a/enumerator.c +++ b/enumerator.c @@ -2653,6 +2653,30 @@ lazy_uniq(VALUE obj) return lazy_add_method(obj, 0, 0, Qnil, Qnil, funcs); } +static struct MEMO * +lazy_compact_proc(VALUE proc_entry, struct MEMO *result, VALUE memos, long memo_index) +{ + if (NIL_P(result->memo_value)) return 0; + return result; +} + +static const lazyenum_funcs lazy_compact_funcs = { + lazy_compact_proc, 0, +}; + +/* + * call-seq: + * lazy.compact -> lazy_enumerator + * + * Like Enumerable#compact, but chains operation to be lazy-evaluated. + */ + +static VALUE +lazy_compact(VALUE obj) +{ + return lazy_add_method(obj, 0, 0, Qnil, Qnil, &lazy_compact_funcs); +} + static struct MEMO * lazy_with_index_proc(VALUE proc_entry, struct MEMO* result, VALUE memos, long memo_index) { @@ -4098,6 +4122,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, "compact", lazy_compact, 0); rb_define_method(rb_cLazy, "with_index", lazy_with_index, -1); lazy_use_super_method = rb_hash_new_with_size(18); diff --git a/test/ruby/test_enum.rb b/test/ruby/test_enum.rb index 126b100b03..fc0220368b 100644 --- a/test/ruby/test_enum.rb +++ b/test/ruby/test_enum.rb @@ -1229,6 +1229,21 @@ class TestEnumerable < Test::Unit::TestCase assert_equal([1, [1, 2]], Foo.new.to_enum.uniq) end + def test_compact + class << (enum = Object.new) + include Enumerable + def each + yield 3 + yield nil + yield 7 + yield 9 + yield nil + end + end + + assert_equal([3, 7, 9], enum.compact) + end + def test_transient_heap_sort_by klass = Class.new do include Comparable diff --git a/test/ruby/test_enumerator.rb b/test/ruby/test_enumerator.rb index eb15d4b616..718bcb74b6 100644 --- a/test/ruby/test_enumerator.rb +++ b/test/ruby/test_enumerator.rb @@ -696,6 +696,11 @@ class TestEnumerator < Test::Unit::TestCase assert_equal([0, 1], u.force) end + def test_compact + u = [0, 1, nil, 2, 3, nil].to_enum.lazy.compact + assert_equal([0, 1, 2, 3], u.force) + end + def test_enum_chain_and_plus r = 1..5