enum.c: Enumerable#tally

* enum.c (enum_tally): new methods Enumerable#tally, which group and count elements of the collection. [Feature #11076] git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@67020 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
ruby · Feb 7, 2019 · 673dc51 · vipulnsward · Mar 3, 2019 · nobu
1 parent e0f1b51
commit 673dc51
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 0 deletions.
diff --git a/NEWS b/NEWS
@@ -24,6 +24,12 @@ sufficient information, see the ChangeLog file or Redmine
 
 === Core classes updates (outstanding ones only)
 
+Enumerable::
+
+  New method::
+
+   * Added Enumerable#tally.  [Feature #11076]
+
 === Stdlib updates (outstanding ones only)
 
 CSV::

diff --git a/enum.c b/enum.c
@@ -935,6 +935,48 @@ enum_group_by(VALUE obj)
     return enum_hashify(obj, 0, 0, group_by_i);
 }
 
+static void
+tally_up(VALUE hash, VALUE group)
+{
+    VALUE tally = rb_hash_aref(hash, group);
+    if (NIL_P(tally)) {
+	tally = INT2FIX(1);
+    }
+    else if (FIXNUM_P(tally) && tally < INT2FIX(FIXNUM_MAX)) {
+	tally += INT2FIX(1) & ~FIXNUM_FLAG;
+    }
+    else {
+	tally = rb_big_plus(tally, INT2FIX(1));
+    }
+    rb_hash_aset(hash, group, tally);
+}
+
+static VALUE
+tally_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, hash))
+{
+    ENUM_WANT_SVALUE();
+    tally_up(hash, i);
+    return Qnil;
+}
+
+/*
+ *  call-seq:
+ *     enum.tally -> a_hash
+ *
+ *  Tallys the collection.  Returns a hash where the keys are the
+ *  elements and the values are numbers of elements in the collection
+ *  that correspond to the key.
+ *
+ *     (1..6).tally { |i| i%3 }   #=> {0=>2, 1=>2, 2=>2}
+ *
+ */
+
+static VALUE
+enum_tally(VALUE obj)
+{
+    return enum_hashify(obj, 0, 0, tally_i);
+}
+
 static VALUE
 first_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, params))
 {
@@ -4070,6 +4112,7 @@ Init_Enumerable(void)
     rb_define_method(rb_mEnumerable, "reduce", enum_inject, -1);
     rb_define_method(rb_mEnumerable, "partition", enum_partition, 0);
     rb_define_method(rb_mEnumerable, "group_by", enum_group_by, 0);
+    rb_define_method(rb_mEnumerable, "tally", enum_tally, 0);
     rb_define_method(rb_mEnumerable, "first", enum_first, -1);
     rb_define_method(rb_mEnumerable, "all?", enum_all, -1);
     rb_define_method(rb_mEnumerable, "any?", enum_any, -1);

diff --git a/test/ruby/test_enum.rb b/test/ruby/test_enum.rb
@@ -294,6 +294,11 @@ def test_group_by
     assert_equal(h, @obj.each_with_index.group_by(&cond))
   end
 
+  def test_tally
+    h = {1 => 2, 2 => 2, 3 => 1}
+    assert_equal(h, @obj.tally)
+  end
+
   def test_first
     assert_equal(1, @obj.first)
     assert_equal([1, 2, 3], @obj.first(3))