Add Enumerable#compact and Enumerator::Lazy#compact

4 files changed, 88 insertions, 0 deletions

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
@@ -2654,6 +2654,30 @@ lazy_uniq(VALUE obj)
 }
 
 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)
 {
     struct proc_entry *entry = proc_entry_ptr(proc_entry);
@@ -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