[Bug #19969] Compact st_table after deleted if possible

3 files changed, 57 insertions, 11 deletions

diff --git a/hash.c b/hash.c
index 2b4c8d3400..641b78891b 100644
--- a/hash.c
+++ b/hash.c
@@ -1423,6 +1423,16 @@ rb_hash_foreach(VALUE hash, rb_foreach_func *func, VALUE farg)
     hash_verify(hash);
 }
 
+void rb_st_compact_table(st_table *tab);
+
+static void
+compact_after_delete(VALUE hash)
+{
+    if (RHASH_ITER_LEV(hash) == 0 && RHASH_ST_TABLE_P(hash)) {
+        rb_st_compact_table(RHASH_ST_TABLE(hash));
+    }
+}
+
 static VALUE
 hash_alloc_flags(VALUE klass, VALUE flags, VALUE ifnone, bool st)
 {
@@ -2397,6 +2407,7 @@ rb_hash_delete_m(VALUE hash, VALUE key)
     val = rb_hash_delete_entry(hash, key);
 
     if (!UNDEF_P(val)) {
+        compact_after_delete(hash);
         return val;
     }
     else {
@@ -2517,6 +2528,7 @@ rb_hash_delete_if(VALUE hash)
     rb_hash_modify_check(hash);
     if (!RHASH_TABLE_EMPTY_P(hash)) {
         rb_hash_foreach(hash, delete_if_i, hash);
+        compact_after_delete(hash);
     }
     return hash;
 }
@@ -2580,6 +2592,7 @@ rb_hash_reject(VALUE hash)
     result = hash_dup_with_compare_by_id(hash);
     if (!RHASH_EMPTY_P(hash)) {
         rb_hash_foreach(result, delete_if_i, result);
+        compact_after_delete(result);
     }
     return result;
 }
@@ -2639,6 +2652,7 @@ rb_hash_except(int argc, VALUE *argv, VALUE hash)
         key = argv[i];
         rb_hash_delete(result, key);
     }
+    compact_after_delete(result);
 
     return result;
 }
@@ -2734,6 +2748,7 @@ rb_hash_select(VALUE hash)
     result = hash_dup_with_compare_by_id(hash);
     if (!RHASH_EMPTY_P(hash)) {
         rb_hash_foreach(result, keep_if_i, result);
+        compact_after_delete(result);
     }
     return result;
 }
@@ -2823,6 +2838,7 @@ rb_hash_clear(VALUE hash)
     }
     else {
         st_clear(RHASH_ST_TABLE(hash));
+        compact_after_delete(hash);
     }
 
     return hash;
@@ -3253,6 +3269,7 @@ rb_hash_transform_keys_bang(int argc, VALUE *argv, VALUE hash)
         rb_ary_clear(pairs);
         rb_hash_clear(new_keys);
     }
+    compact_after_delete(hash);
     return hash;
 }
 
@@ -3303,6 +3320,7 @@ rb_hash_transform_values(VALUE hash)
 
     if (!RHASH_EMPTY_P(hash)) {
         rb_hash_stlike_foreach_with_replace(result, transform_values_foreach_func, transform_values_foreach_replace, result);
+        compact_after_delete(result);
     }
 
     return result;
@@ -4267,6 +4285,7 @@ rb_hash_compact(VALUE hash)
     VALUE result = rb_hash_dup(hash);
     if (!RHASH_EMPTY_P(hash)) {
         rb_hash_foreach(result, delete_if_nil, result);
+        compact_after_delete(result);
     }
     else if (rb_hash_compare_by_id_p(hash)) {
         result = rb_hash_compare_by_id(result);
diff --git a/st.c b/st.c
index 755e230c9a..2825ed1eb6 100644
--- a/st.c
+++ b/st.c
@@ -717,6 +717,8 @@ count_collision(const struct st_hash_type *type)
 #error "REBUILD_THRESHOLD should be >= 2"
 #endif
 
+static void rebuild_table_with(st_table *new_tab, st_table *tab);
+
 /* Rebuild table TAB.  Rebuilding removes all deleted bins and entries
    and can change size of the table entries and bins arrays.
    Rebuilding is implemented by creation of a new table or by
@@ -724,14 +726,6 @@ count_collision(const struct st_hash_type *type)
 static void
 rebuild_table(st_table *tab)
 {
-    st_index_t i, ni;
-    unsigned int size_ind;
-    st_table *new_tab;
-    st_table_entry *new_entries;
-    st_table_entry *curr_entry_ptr;
-    st_index_t *bins;
-    st_index_t bin_ind;
-
     if ((2 * tab->num_entries <= get_allocated_entries(tab)
          && REBUILD_THRESHOLD * tab->num_entries > get_allocated_entries(tab))
         || tab->num_entries < (1 << MINIMAL_POWER2)) {
@@ -739,17 +733,30 @@ rebuild_table(st_table *tab)
         tab->num_entries = 0;
         if (tab->bins != NULL)
             initialize_bins(tab);
-        new_tab = tab;
-        new_entries = tab->entries;
+        rebuild_table_with(tab, tab);
     }
     else {
+        st_table *new_tab;
         /* This allocation could trigger GC and compaction. If tab is the
          * gen_iv_tbl, then tab could have changed in size due to objects being
          * freed and/or moved. Do not store attributes of tab before this line. */
         new_tab = st_init_table_with_size(tab->type,
                                           2 * tab->num_entries - 1);
-        new_entries = new_tab->entries;
+        rebuild_table_with(new_tab, tab);
     }
+}
+
+static void
+rebuild_table_with(st_table *new_tab, st_table *tab)
+{
+    st_index_t i, ni;
+    unsigned int size_ind;
+    st_table_entry *new_entries;
+    st_table_entry *curr_entry_ptr;
+    st_index_t *bins;
+    st_index_t bin_ind;
+
+    new_entries = new_tab->entries;
 
     ni = 0;
     bins = new_tab->bins;
@@ -2296,4 +2303,15 @@ rb_st_nth_key(st_table *tab, st_index_t index)
     }
 }
 
+void
+rb_st_compact_table(st_table *tab)
+{
+    st_index_t num = tab->num_entries;
+    if (REBUILD_THRESHOLD * num <= get_allocated_entries(tab)) {
+        /* Compaction: */
+        st_table *new_tab = st_init_table_with_size(tab->type, 2 * num);
+        rebuild_table_with(new_tab, tab);
+    }
+}
+
 #endif
diff --git a/test/ruby/test_hash.rb b/test/ruby/test_hash.rb
index f95876d855..35d30095c2 100644
--- a/test/ruby/test_hash.rb
+++ b/test/ruby/test_hash.rb
@@ -1792,6 +1792,15 @@ class TestHash < Test::Unit::TestCase
     assert_no_memory_leak([], prepare, code, bug9187)
   end
 
+  def test_memory_size_after_delete
+    require 'objspace'
+    h = {}
+    1000.times {|i| h[i] = true}
+    big = ObjectSpace.memsize_of(h)
+    1000.times {|i| h.delete(i)}
+    assert_operator ObjectSpace.memsize_of(h), :<, big/10
+  end
+
   def test_wrapper
     bug9381 = '[ruby-core:59638] [Bug #9381]'