[Feature #19787] Add Enumerable#uniq_map, Enumerable::Lazy#uniq_map, Array#uniq_map and Array#uniq_map!

array.c

@@ -6204,6 +6204,92 @@ rb_ary_uniq(VALUE ary)
     return uniq;
 }
 
+/*
+ *  call-seq:
+ *    array.uniq_map {|element| ... } -> new_array
+ *    array.uniq_map -> new_enumerator
+ *
+ *  Calls the block, if given, with each element of +self+;
+ *  returns a new \Array whose elements are the return values from
+ *  the block, omitting duplicates while retaining the first occurrence.
+ *
+ *  With a block given, calls the block for each element;
+ *  identifies (using method <tt>eql?</tt>) and removes
+ *  duplicate elements returned by the block as a new array:
+ *
+ *    a = [0, 0, 1, 1, 2, 2]
+ *    a.uniq_map {|i| i * 2 } # => [0, 2, 4]
+ *
+ *  When no block given, returns a new \Enumerator:
+ *
+ *    a = [0, 0, 1, 1, 2, 2]
+ *    a.uniq_map #<Enumerator: [0, 0, 1, 1, 2, 2]:uniq_map>
+ */
+
+static VALUE
+rb_ary_uniq_map(VALUE ary)
+{
+    long i, val;
+    VALUE hash;
+
+    RETURN_SIZED_ENUMERATOR(ary, 0, 0, ary_enum_length);
+
+    hash = rb_obj_hide(rb_hash_new_capa(RARRAY_LEN(ary)));
+    for (i = 0; i < RARRAY_LEN(ary); i++) {
+        val = rb_yield(RARRAY_AREF(ary, i));
+        rb_hash_add_new_element(hash, val, val);
+    }
+
+    return rb_hash_values(hash);
+}
+
+/*
+ *  call-seq:
+ *    array.uniq_map! {|element| ... } -> self
+ *    array.uniq_map! -> Enumerator
+ *
+ *  Removes from +self+ duplicate elements returned by the block,
+ *  the first occurrence always being retained.
+ *
+ *  With a block given, calls the block for each element;
+ *  identifies (using method <tt>eql?</tt>) and removes
+ *  duplicate elements returned by the block:
+ *
+ *    a = [0, 0, 1, 1, 2, 2]
+ *    a.uniq_map! {|i| i * 2 } # => [0, 2, 4]
+ *
+ *  When no block given, returns a new \Enumerator:
+ *
+ *    a = [0, 0, 1, 1, 2, 2]
+ *    a.uniq_map! #<Enumerator: [0, 0, 1, 1, 2, 2]:uniq_map!>
+ */
+
+static VALUE
+rb_ary_uniq_map_bang(VALUE ary)
+{
+    long i1, i2, val;
+    VALUE hash;
+    long hash_size;
+
+    RETURN_SIZED_ENUMERATOR(ary, 0, 0, ary_enum_length);
+
+    rb_ary_modify(ary);
+    hash = rb_obj_hide(rb_hash_new_capa(RARRAY_LEN(ary)));
+    for (i1 = i2 = 0; i1 < RARRAY_LEN(ary); i1++) {
+        val = rb_yield(RARRAY_AREF(ary, i1));
+        if (!rb_hash_add_new_element(hash, val, val)) {
+            rb_ary_store(ary, i2, val);
+            i2++;
+        }
+    }
+
+    hash_size = RHASH_SIZE(hash);
+    ARY_SET_LEN(ary, hash_size);
+    ary_resize_capa(ary, hash_size);
+
+    return ary;
+}
+
 /*
  *  call-seq:
  *    array.compact! -> self or nil
@@ -8584,6 +8670,8 @@ rb_ary_deconstruct(VALUE ary)
  *
  *  - #map, #collect: Returns an array containing the block return-value for each element.
  *  - #map!, #collect!: Replaces each element with a block return-value.
+ *  - #uniq_map: Returns an array containing the block return-value for each element without duplicates.
+ *  - #uniq_map!: Replaces each element with a block return-value, removing duplicates.
  *  - #flatten: Returns an array that is a recursive flattening of +self+.
  *  - #flatten!: Replaces each nested array in +self+ with the elements from that array.
  *  - #inspect, #to_s: Returns a new String containing the elements.
@@ -8707,6 +8795,8 @@ Init_Array(void)
 
     rb_define_method(rb_cArray, "uniq", rb_ary_uniq, 0);
     rb_define_method(rb_cArray, "uniq!", rb_ary_uniq_bang, 0);
+    rb_define_method(rb_cArray, "uniq_map", rb_ary_uniq_map, 0);
+    rb_define_method(rb_cArray, "uniq_map!", rb_ary_uniq_map_bang, 0);
     rb_define_method(rb_cArray, "compact", rb_ary_compact, 0);
     rb_define_method(rb_cArray, "compact!", rb_ary_compact_bang, 0);
     rb_define_method(rb_cArray, "flatten", rb_ary_flatten, -1);

benchmark/array_uniq_map.yml

@@ -0,0 +1,6 @@
+prelude: a = Array.new(100) { |i| i % 2 }
+benchmark:
+  a.map { ... }.uniq: a.map { |i| i * 2 }.uniq
+  a.uniq_map { ... }: a.uniq_map { |i| i * 2 }
+  a.map! { ... }.uniq!: a.map! { |i| i * 2 }.uniq!
+  a.uniq_map! { ... }: a.uniq_map! { |i| i * 2 }

benchmark/enum_uniq_map.yml

@@ -0,0 +1,4 @@
+prelude: e = Array.new(100) { |i| i % 2 }.to_enum
+benchmark:
+  e.map { ... }.uniq: e.map { |i| i * 2 }.uniq
+  e.uniq_map { ... }: e.uniq_map { |i| i * 2 }

enum.c

@@ -4848,6 +4848,49 @@ enum_uniq(VALUE obj)
     return ret;
 }
 
+static VALUE
+uniq_map_iter(RB_BLOCK_CALL_FUNC_ARGLIST(i, hash))
+{
+    long val;
+
+    val = rb_yield_values2(argc, argv);
+    rb_hash_add_new_element(hash, val, val);
+
+    return Qnil;
+}
+
+/*
+ * call-seq:
+ *   uniq_map {|element| ... } -> array
+ *   uniq_map                  -> enumerator
+ *
+ * Returns an array containing only unique elements returned by the block.
+ *
+ * With a block given, calls the block with successive elements;
+ * returns an array which has no two elements +e0+ and +e1+ returned by
+ * the block such that <tt>e0.eql?(e1)</tt>:
+ *
+ *   [0, 1, 1].uniq_map {|i| i }                                  # => [0, 1]
+ *   (0..5).uniq_map {|i| i.odd? ? i * 2 : i }                    # => [0, 2, 6, 4, 10]
+ *   {foo: 0, bar: 1, baz: 1}.uniq_map {|_key, value| value * 2 } # => [0, 2]
+ *
+ * When no block is given, returns an \Enumerator.
+ *
+ */
+
+static VALUE
+enum_uniq_map(VALUE obj)
+{
+    VALUE hash;
+
+    RETURN_SIZED_ENUMERATOR(obj, 0, 0, enum_size);
+
+    hash = rb_obj_hide(rb_hash_new());
+    rb_block_call(obj, id_each, 0, 0, uniq_map_iter, hash);
+
+    return rb_hash_values(hash);
+}
+
 static VALUE
 compact_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, ary))
 {
@@ -4977,6 +5020,7 @@ enum_compact(VALUE obj)
  * - #map, #collect: Returns objects returned by the block.
  * - #filter_map: Returns truthy objects returned by the block.
  * - #flat_map, #collect_concat: Returns flattened objects returned by the block.
+ * - #uniq_map: Returns unique objects returned by the block.
  * - #grep: Returns elements selected by a given object
  *   or objects returned by a given block.
  * - #grep_v: Returns elements selected by a given object
@@ -5118,6 +5162,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, "uniq_map", enum_uniq_map, 0);
     rb_define_method(rb_mEnumerable, "compact", enum_compact, 0);
 
     id__alone = rb_intern_const("_alone");

enumerator.c

@@ -2689,6 +2689,54 @@ lazy_uniq(VALUE obj)
     return lazy_add_method(obj, 0, 0, Qnil, Qnil, funcs);
 }
 
+static int
+lazy_uniq_map_check(VALUE value, VALUE memos, long memo_index)
+{
+    VALUE hash = rb_ary_entry(memos, memo_index);
+
+    if (NIL_P(hash)) {
+        hash = rb_obj_hide(rb_hash_new());
+        rb_ary_store(memos, memo_index, hash);
+    }
+
+    return rb_hash_add_new_element(hash, value, value);
+}
+
+static struct MEMO *
+lazy_uniq_map_iter_proc(VALUE proc_entry, struct MEMO *result, VALUE memos, long memo_index)
+{
+    VALUE value = lazyenum_yield(proc_entry, result);
+
+    if (lazy_uniq_map_check(value, memos, memo_index)) return 0;
+
+    LAZY_MEMO_SET_VALUE(result, value);
+    LAZY_MEMO_RESET_PACKED(result);
+
+    return result;
+}
+
+static const lazyenum_funcs lazy_uniq_map_iter_funcs = {
+    lazy_uniq_map_iter_proc, 0,
+};
+
+/*
+ *  call-seq:
+ *     lazy.uniq_map                  -> lazy_enumerator
+ *     lazy.uniq_map { |item| block } -> lazy_enumerator
+ *
+ *  Like Enumerable#uniq_map, but chains operation to be lazy-evaluated.
+ */
+
+static VALUE
+lazy_uniq_map(VALUE obj)
+{
+    if (!rb_block_given_p()) {
+        rb_raise(rb_eArgError, "tried to call lazy uniq_map without a block");
+    }
+
+    return lazy_add_method(obj, 0, 0, Qnil, Qnil, &lazy_uniq_map_iter_funcs);
+}
+
 static struct MEMO *
 lazy_compact_proc(VALUE proc_entry, struct MEMO *result, VALUE memos, long memo_index)
 {
@@ -4492,6 +4540,7 @@ InitVM_Enumerator(void)
     rb_define_alias(rb_cLazy, "_enumerable_drop", "drop");
     rb_define_alias(rb_cLazy, "_enumerable_drop_while", "drop_while");
     rb_define_alias(rb_cLazy, "_enumerable_uniq", "uniq");
+    rb_define_alias(rb_cLazy, "_enumerable_uniq_map", "uniq_map");
     rb_define_private_method(rb_cLazy, "_enumerable_with_index", enumerator_with_index, -1);
 
     rb_funcall(rb_cLazy, id_private, 1, sym("_enumerable_map"));
@@ -4511,6 +4560,7 @@ InitVM_Enumerator(void)
     rb_funcall(rb_cLazy, id_private, 1, sym("_enumerable_drop"));
     rb_funcall(rb_cLazy, id_private, 1, sym("_enumerable_drop_while"));
     rb_funcall(rb_cLazy, id_private, 1, sym("_enumerable_uniq"));
+    rb_funcall(rb_cLazy, id_private, 1, sym("_enumerable_uniq_map"));
 
     rb_define_method(rb_cLazy, "initialize", lazy_initialize, -1);
     rb_define_method(rb_cLazy, "to_enum", lazy_to_enum, -1);
@@ -4539,6 +4589,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, "uniq_map", lazy_uniq_map, 0);
     rb_define_method(rb_cLazy, "compact", lazy_compact, 0);
     rb_define_method(rb_cLazy, "with_index", lazy_with_index, -1);
 
@@ -4560,6 +4611,7 @@ InitVM_Enumerator(void)
     rb_hash_aset(lazy_use_super_method, sym("drop"), sym("_enumerable_drop"));
     rb_hash_aset(lazy_use_super_method, sym("drop_while"), sym("_enumerable_drop_while"));
     rb_hash_aset(lazy_use_super_method, sym("uniq"), sym("_enumerable_uniq"));
+    rb_hash_aset(lazy_use_super_method, sym("uniq_map"), sym("_enumerable_uniq_map"));
     rb_hash_aset(lazy_use_super_method, sym("with_index"), sym("_enumerable_with_index"));
     rb_obj_freeze(lazy_use_super_method);
     rb_vm_register_global_object(lazy_use_super_method);

spec/ruby/core/array/shared/collect.rb

@@ -90,6 +90,10 @@
       -> { ArraySpecs.frozen_array.send(@method) {} }.should raise_error(FrozenError)
     end
 
+    it "doesn't yield to the block on a frozen array" do
+      -> { ArraySpecs.frozen_array.send(@method) {raise RangeError, "shouldn't yield" } }.should raise_error(FrozenError)
+    end
+
     it "raises a FrozenError when empty" do
       -> { ArraySpecs.empty_frozen_array.send(@method) {} }.should raise_error(FrozenError)
     end