From 31e0382723bfb35cffe3ca485dd0577668cafa07 Mon Sep 17 00:00:00 2001
From: Kenichi Kamiya <kachick1@gmail.com>
Date: Sun, 28 Mar 2021 09:14:57 +0900
Subject: [PATCH] Keep non evaluated keys in `Hash#transform_keys!` [Bug
 #17735]

---
 hash.c                                     |  6 +++++-
 spec/ruby/core/hash/transform_keys_spec.rb | 12 +++++++++++-
 test/ruby/test_hash.rb                     |  8 ++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/hash.c b/hash.c
index fdecb1edad..d6d3752fb7 100644
--- a/hash.c
+++ b/hash.c
@@ -3268,8 +3268,8 @@ rb_hash_transform_keys_bang(int argc, VALUE *argv, VALUE hash)
     rb_hash_modify_check(hash);
     if (!RHASH_TABLE_EMPTY_P(hash)) {
         long i;
+        VALUE new_keys = hash_alloc(0);
         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, val;
 
@@ -3286,7 +3286,11 @@ rb_hash_transform_keys_bang(int argc, VALUE *argv, VALUE hash)
                 new_key = key;
             }
             val = RARRAY_AREF(pairs, i+1);
+            if (!hash_stlike_lookup(new_keys, key, NULL)) {
+                rb_hash_stlike_delete(hash, &key, NULL);
+            }
             rb_hash_aset(hash, new_key, val);
+            rb_hash_aset(new_keys, new_key, Qnil);
         }
     }
     return hash;
diff --git a/spec/ruby/core/hash/transform_keys_spec.rb b/spec/ruby/core/hash/transform_keys_spec.rb
index 2c5d4124e0..63403a25b0 100644
--- a/spec/ruby/core/hash/transform_keys_spec.rb
+++ b/spec/ruby/core/hash/transform_keys_spec.rb
@@ -84,7 +84,7 @@ describe "Hash#transform_keys!" do
     end
   end
 
-  ruby_version_is "2.5.1" do
+  ruby_version_is "2.5.1"..."3.1" do
     it "returns the processed keys if we broke from the block" do
       @hash.transform_keys! do |v|
         break if v == :c
@@ -94,6 +94,16 @@ describe "Hash#transform_keys!" do
     end
   end
 
+  ruby_version_is "3.1" do
+    it "returns the processed keys and non evaluated keys if we broke from the block" do
+      @hash.transform_keys! do |v|
+        break if v == :c
+        v.succ
+      end
+      @hash.should == { b: 1, c: 2, d: 4 }
+    end
+  end
+
   it "keeps later pair if new keys conflict" do
     @hash.transform_keys! { |_| :a }.should == { a: 4 }
   end
diff --git a/test/ruby/test_hash.rb b/test/ruby/test_hash.rb
index bd94b408bb..23df2410d7 100644
--- a/test/ruby/test_hash.rb
+++ b/test/ruby/test_hash.rb
@@ -1765,6 +1765,10 @@ class TestHash < Test::Unit::TestCase
     x.transform_keys! {|k| -k }
     assert_equal([-1, :a, 1, :b], x.flatten)
 
+    x = @cls[a: 1, b: 2, c: 3]
+    x.transform_keys! { |k| k == :b && break }
+    assert_equal({false => 1, b: 2, c: 3}, x)
+
     x = @cls[true => :a, false => :b]
     x.transform_keys! {|k| !k }
     assert_equal([false, :a, true, :b], x.flatten)
@@ -1800,6 +1804,10 @@ class TestHash < Test::Unit::TestCase
     assert_equal([1, 4, 9], y.values_at(:a, :b, :c))
     assert_same(x, y)
 
+    x = @cls[a: 1, b: 2, c: 3]
+    x.transform_values! { |v| v == 2 && break }
+    assert_equal({a: false, b: 2, c: 3}, x)
+
     x = @cls[a: 1, b: 2, c: 3]
     y = x.transform_values!.with_index {|v, i| "#{v}.#{i}" }
     assert_equal(%w(1.0  2.1  3.2), y.values_at(:a, :b, :c))