Make `#to_h` to take a block; [ruby-core:89088]

4 files changed, 42 insertions, 50 deletions

diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb
index c0995bb99..bb9e61bdd 100644
--- a/mrbgems/mruby-array-ext/mrblib/array.rb
+++ b/mrbgems/mruby-array-ext/mrblib/array.rb
@@ -932,4 +932,29 @@ class Array
       self.map { |row| row[column_index] }
     end
   end
+
+  ##
+  #  call-seq:
+  #    ary.to_h                ->   Hash
+  #    ary.to_h{|item| ... }   ->   Hash
+  #
+  # Returns the result of interpreting <i>aray</i> as an array of
+  # <tt>[key, value]</tt> pairs. If a block is given, it should
+  # return <tt>[key, value]</tt> pairs to construct a hash.
+  #
+  #     [[:foo, :bar], [1, 2]].to_h
+  #       # => {:foo => :bar, 1 => 2}
+  #     [1, 2].to_h{|x| [x, x*2]}
+  #       # => {1 => 2, 2 => 4}
+  #
+  def to_h(&blk)
+    h = {}
+    self.each do |v|
+      v = blk.call(v) if blk
+      raise TypeError, "wrong element type #{v.class}" unless Array === v
+      raise ArgumentError, "wrong array length (expected 2, was #{v.length})" unless v.length == 2
+      h[v[0]] = v[1]
+    end
+    h
+  end
 end
diff --git a/mrbgems/mruby-array-ext/src/array.c b/mrbgems/mruby-array-ext/src/array.c
index 792eb26de..b6d9c9c80 100644
--- a/mrbgems/mruby-array-ext/src/array.c
+++ b/mrbgems/mruby-array-ext/src/array.c
@@ -106,49 +106,6 @@ mrb_ary_values_at(mrb_state *mrb, mrb_value self)
   return mrb_get_values_at(mrb, self, RARRAY_LEN(self), argc, argv, mrb_ary_ref);
 }
 
-/*
- *  call-seq:
- *     ary.to_h   ->   Hash
- *
- *  Returns the result of interpreting <i>aray</i> as an array of
- *  <tt>[key, value]</tt> paris.
- *
- *      [[:foo, :bar], [1, 2]].to_h
- *        # => {:foo => :bar, 1 => 2}
- *
- */
-
-static mrb_value
-mrb_ary_to_h(mrb_state *mrb, mrb_value ary)
-{
-  mrb_int i;
-  mrb_value v, hash;
-
-  hash = mrb_hash_new_capa(mrb, 0);
-
-  for (i = 0; i < RARRAY_LEN(ary); ++i) {
-    mrb_value elt = RARRAY_PTR(ary)[i];
-    v = mrb_check_array_type(mrb, elt);
-
-    if (mrb_nil_p(v)) {
-      mrb_raisef(mrb, E_TYPE_ERROR, "wrong element type %S at %S (expected array)",
-                 mrb_str_new_cstr(mrb,  mrb_obj_classname(mrb, elt)),
-                 mrb_fixnum_value(i)
-      );
-    }
-
-    if (RARRAY_LEN(v) != 2) {
-      mrb_raisef(mrb, E_ARGUMENT_ERROR, "wrong array length at %S (expected 2, was %S)",
-                 mrb_fixnum_value(i),
-                 mrb_fixnum_value(RARRAY_LEN(v))
-      );
-    }
-
-    mrb_hash_set(mrb, hash, RARRAY_PTR(v)[0], RARRAY_PTR(v)[1]);
-  }
-
-  return hash;
-}
 
 /*
  *  call-seq:
@@ -237,7 +194,6 @@ mrb_mruby_array_ext_gem_init(mrb_state* mrb)
   mrb_define_method(mrb, a, "at",     mrb_ary_at,     MRB_ARGS_REQ(1));
   mrb_define_method(mrb, a, "rassoc", mrb_ary_rassoc, MRB_ARGS_REQ(1));
   mrb_define_method(mrb, a, "values_at", mrb_ary_values_at, MRB_ARGS_ANY());
-  mrb_define_method(mrb, a, "to_h",   mrb_ary_to_h, MRB_ARGS_REQ(0));
   mrb_define_method(mrb, a, "slice!", mrb_ary_slice_bang,   MRB_ARGS_ANY());
 }
 
diff --git a/mrbgems/mruby-enum-ext/mrblib/enum.rb b/mrbgems/mruby-enum-ext/mrblib/enum.rb
index a840ade3b..6cbacdf9e 100644
--- a/mrbgems/mruby-enum-ext/mrblib/enum.rb
+++ b/mrbgems/mruby-enum-ext/mrblib/enum.rb
@@ -803,13 +803,22 @@ module Enumerable
   #       # => {:hello => 0, :world => 1}
   #
 
-  def to_h
+  def to_h(&blk)
     h = {}
-    self.each do |*v|
-      v = v.__svalue
-      raise TypeError, "wrong element type #{v.class} (expected Array)" unless v.is_a? Array
-      raise ArgumentError, "element has wrong array length (expected 2, was #{v.size})" if v.size != 2
-      h[v[0]] = v[1]
+    if blk
+      self.each do |v|
+        v = blk.call(v)
+        raise TypeError, "wrong element type #{v.class} (expected Array)" unless v.is_a? Array
+        raise ArgumentError, "element has wrong array length (expected 2, was #{v.size})" if v.size != 2
+        h[v[0]] = v[1]
+      end
+    else
+      self.each do |*v|
+        v = v.__svalue
+        raise TypeError, "wrong element type #{v.class} (expected Array)" unless v.is_a? Array
+        raise ArgumentError, "element has wrong array length (expected 2, was #{v.size})" if v.size != 2
+        h[v[0]] = v[1]
+      end
     end
     h
   end
diff --git a/mrbgems/mruby-enum-ext/test/enum.rb b/mrbgems/mruby-enum-ext/test/enum.rb
index 46ed5f0f9..b2d7297a4 100644
--- a/mrbgems/mruby-enum-ext/test/enum.rb
+++ b/mrbgems/mruby-enum-ext/test/enum.rb
@@ -188,4 +188,6 @@ assert("Enumerable#to_h") do
   assert_equal h0, h
   # mruby-enum-ext also provides nil.to_h
   assert_equal Hash.new, nil.to_h
+
+  assert_equal({1=>4,3=>8}, c.new.to_h{|k,v|[k,v*2]})
 end