Transform hash keys by a hash [Feature #16274]

2 files changed, 78 insertions, 9 deletions

diff --git a/hash.c b/hash.c
index 82383e9a93..1fd3a335b3 100644
--- a/hash.c
+++ b/hash.c
@@ -3023,6 +3023,28 @@ rb_hash_each_pair(VALUE hash)
     return hash;
 }
 
+struct transform_keys_args{
+    VALUE trans;
+    VALUE result;
+    int block_given;
+};
+
+static int
+transform_keys_hash_i(VALUE key, VALUE value, VALUE transarg)
+{
+    struct transform_keys_args *p = (void *)transarg;
+    VALUE trans = p->trans, result = p->result;
+    VALUE new_key = rb_hash_lookup2(trans, key, Qundef);
+    if (new_key == Qundef) {
+        if (p->block_given)
+            new_key = rb_yield(key);
+        else
+            new_key = key;
+    }
+    rb_hash_aset(result, new_key, value);
+    return ST_CONTINUE;
+}
+
 static int
 transform_keys_i(VALUE key, VALUE value, VALUE result)
 {
@@ -3049,14 +3071,28 @@ transform_keys_i(VALUE key, VALUE value, VALUE result)
  *  If no block is given, an enumerator is returned instead.
  */
 static VALUE
-rb_hash_transform_keys(VALUE hash)
+rb_hash_transform_keys(int argc, VALUE *argv, VALUE hash)
 {
     VALUE result;
+    struct transform_keys_args transarg = {0};
 
-    RETURN_SIZED_ENUMERATOR(hash, 0, 0, hash_enum_size);
+    argc = rb_check_arity(argc, 0, 1);
+    if (argc > 0) {
+        transarg.trans = to_hash(argv[0]);
+        transarg.block_given = rb_block_given_p();
+    }
+    else {
+        RETURN_SIZED_ENUMERATOR(hash, 0, 0, hash_enum_size);
+    }
     result = rb_hash_new();
     if (!RHASH_EMPTY_P(hash)) {
-        rb_hash_foreach(hash, transform_keys_i, result);
+        if (transarg.trans) {
+            transarg.result = result;
+            rb_hash_foreach(hash, transform_keys_hash_i, (VALUE)&transarg);
+        }
+        else {
+            rb_hash_foreach(hash, transform_keys_i, result);
+        }
     }
 
     return result;
@@ -3082,17 +3118,40 @@ static VALUE rb_hash_flatten(int argc, VALUE *argv, VALUE hash);
  *  If no block is given, an enumerator is returned instead.
  */
 static VALUE
-rb_hash_transform_keys_bang(VALUE hash)
+rb_hash_transform_keys_bang(int argc, VALUE *argv, VALUE hash)
 {
-    RETURN_SIZED_ENUMERATOR(hash, 0, 0, hash_enum_size);
+    VALUE trans = 0;
+    int block_given = 0;
+
+    argc = rb_check_arity(argc, 0, 1);
+    if (argc > 0) {
+        trans = to_hash(argv[0]);
+        block_given = rb_block_given_p();
+    }
+    else {
+        RETURN_SIZED_ENUMERATOR(hash, 0, 0, hash_enum_size);
+    }
     rb_hash_modify_check(hash);
     if (!RHASH_TABLE_EMPTY_P(hash)) {
         long i;
         VALUE pairs = rb_hash_flatten(0, NULL, hash);
         rb_hash_clear(hash);
         for (i = 0; i < RARRAY_LEN(pairs); i += 2) {
-            VALUE key = RARRAY_AREF(pairs, i), new_key = rb_yield(key),
-                  val = RARRAY_AREF(pairs, i+1);
+            VALUE key = RARRAY_AREF(pairs, i), new_key, val;
+
+            if (!trans) {
+                new_key = rb_yield(key);
+            }
+            else if ((new_key = rb_hash_lookup2(trans, key, Qundef)) != Qundef) {
+                /* use the transformed key */
+            }
+            else if (block_given) {
+                new_key = rb_yield(key);
+            }
+            else {
+                new_key = key;
+            }
+            val = RARRAY_AREF(pairs, i+1);
             rb_hash_aset(hash, new_key, val);
         }
     }
@@ -6285,8 +6344,8 @@ Init_Hash(void)
     rb_define_method(rb_cHash, "each_pair", rb_hash_each_pair, 0);
     rb_define_method(rb_cHash, "each", rb_hash_each_pair, 0);
 
-    rb_define_method(rb_cHash, "transform_keys", rb_hash_transform_keys, 0);
-    rb_define_method(rb_cHash, "transform_keys!", rb_hash_transform_keys_bang, 0);
+    rb_define_method(rb_cHash, "transform_keys", rb_hash_transform_keys, -1);
+    rb_define_method(rb_cHash, "transform_keys!", rb_hash_transform_keys_bang, -1);
     rb_define_method(rb_cHash, "transform_values", rb_hash_transform_values, 0);
     rb_define_method(rb_cHash, "transform_values!", rb_hash_transform_values_bang, 0);
 
diff --git a/test/ruby/test_hash.rb b/test/ruby/test_hash.rb
index ccc3355930..733bccd138 100644
--- a/test/ruby/test_hash.rb
+++ b/test/ruby/test_hash.rb
@@ -1638,6 +1638,9 @@ class TestHash < Test::Unit::TestCase
 
     y = x.transform_keys.with_index {|k, i| "#{k}.#{i}" }
     assert_equal(%w(a.0 b.1 c.2), y.keys)
+
+    assert_equal({A: 1, B: 2, c: 3}, x.transform_keys({a: :A, b: :B, d: :D}))
+    assert_equal({A: 1, B: 2, "c" => 3}, x.transform_keys({a: :A, b: :B, d: :D}, &:to_s))
   end
 
   def test_transform_keys_bang
@@ -1660,6 +1663,13 @@ class TestHash < Test::Unit::TestCase
     x = @cls[true => :a, false => :b]
     x.transform_keys! {|k| !k }
     assert_equal([false, :a, true, :b], x.flatten)
+
+    x = @cls[a: 1, b: 2, c: 3]
+    x.transform_keys!({a: :A, b: :B, d: :D})
+    assert_equal({A: 1, B: 2, c: 3}, x)
+    x = @cls[a: 1, b: 2, c: 3]
+    x.transform_keys!({a: :A, b: :B, d: :D}, &:to_s)
+    assert_equal({A: 1, B: 2, "c" => 3}, x)
   end
 
   def test_transform_values