From d094c3ef046aba0bb99fd08bcbc72ff87216e736 Mon Sep 17 00:00:00 2001
From: Marc-Andre Lafortune <github@marc-andre.ca>
Date: Mon, 15 Mar 2021 23:51:13 -0400
Subject: [PATCH] Avoid rehashing in Hash#select/reject [Bug #16996]

---
 hash.c                 | 32 +++++++-------------------------
 test/ruby/test_hash.rb | 25 +++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/hash.c b/hash.c
index b34351aab3..2dd61dcd3e 100644
--- a/hash.c
+++ b/hash.c
@@ -2547,15 +2547,6 @@ rb_hash_reject_bang(VALUE hash)
     return hash;
 }
 
-static int
-reject_i(VALUE key, VALUE value, VALUE result)
-{
-    if (!RTEST(rb_yield_values(2, key, value))) {
-	rb_hash_aset(result, key, value);
-    }
-    return ST_CONTINUE;
-}
-
 /*
  *  call-seq:
  *    hash.reject {|key, value| ... } -> new_hash
@@ -2586,9 +2577,9 @@ rb_hash_reject(VALUE hash)
 	    rb_warn("extra states are no longer copied: %+"PRIsVALUE, hash);
 	}
     }
-    result = rb_hash_new();
+    result = hash_copy(hash_alloc(rb_cHash), hash);
     if (!RHASH_EMPTY_P(hash)) {
-	rb_hash_foreach(hash, reject_i, result);
+	rb_hash_foreach(result, delete_if_i, result);
     }
     return result;
 }
@@ -2711,10 +2702,10 @@ rb_hash_fetch_values(int argc, VALUE *argv, VALUE hash)
 }
 
 static int
-select_i(VALUE key, VALUE value, VALUE result)
+keep_if_i(VALUE key, VALUE value, VALUE hash)
 {
-    if (RTEST(rb_yield_values(2, key, value))) {
-	rb_hash_aset(result, key, value);
+    if (!RTEST(rb_yield_values(2, key, value))) {
+	return ST_DELETE;
     }
     return ST_CONTINUE;
 }
@@ -2742,22 +2733,13 @@ rb_hash_select(VALUE hash)
     VALUE result;
 
     RETURN_SIZED_ENUMERATOR(hash, 0, 0, hash_enum_size);
-    result = rb_hash_new();
+    result = hash_copy(hash_alloc(rb_cHash), hash);
     if (!RHASH_EMPTY_P(hash)) {
-	rb_hash_foreach(hash, select_i, result);
+	rb_hash_foreach(result, keep_if_i, result);
     }
     return result;
 }
 
-static int
-keep_if_i(VALUE key, VALUE value, VALUE hash)
-{
-    if (!RTEST(rb_yield_values(2, key, value))) {
-	return ST_DELETE;
-    }
-    return ST_CONTINUE;
-}
-
 /*
  *  call-seq:
  *    hash.select! {|key, value| ... } -> self or nil
diff --git a/test/ruby/test_hash.rb b/test/ruby/test_hash.rb
index 62d8b3f836..25dfaf4f64 100644
--- a/test/ruby/test_hash.rb
+++ b/test/ruby/test_hash.rb
@@ -122,6 +122,31 @@ class TestHash < Test::Unit::TestCase
     assert_equal set2, set2.dup
   end
 
+  def assert_hash_does_not_rehash
+    obj = Object.new
+    class << obj
+      attr_accessor :hash_calls
+      def hash
+        @hash_calls += 1
+        super
+      end
+    end
+    obj.hash_calls = 0
+    hash = {obj => 42}
+    assert_equal(1, obj.hash_calls)
+    yield hash
+    assert_equal(1, obj.hash_calls)
+  end
+
+  def test_select_reject_will_not_rehash
+    assert_hash_does_not_rehash do |hash|
+      hash.select { true }
+    end
+    assert_hash_does_not_rehash do |hash|
+      hash.reject { false }
+    end
+  end
+
   def test_s_AREF
     h = @cls["a" => 100, "b" => 200]
     assert_equal(100, h['a'])