2 files changed, 49 insertions, 34 deletions

diff --git a/hash.c b/hash.c
index 78fbfb866c..d025b1c30a 100644
--- a/hash.c
+++ b/hash.c
@@ -4080,24 +4080,17 @@ assoc_cmp(VALUE a, VALUE b)
     return !RTEST(rb_equal(a, b));
 }
 
-static VALUE
-lookup2_call(VALUE arg)
-{
-    VALUE *args = (VALUE *)arg;
-    return rb_hash_lookup2(args[0], args[1], Qundef);
-}
-
-struct reset_hash_type_arg {
-    VALUE hash;
-    const struct st_hash_type *orighash;
+struct assoc_arg {
+    st_table *tbl;
+    st_data_t key;
 };
 
 static VALUE
-reset_hash_type(VALUE arg)
+assoc_lookup(VALUE arg)
 {
-    struct reset_hash_type_arg *p = (struct reset_hash_type_arg *)arg;
-    HASH_ASSERT(RHASH_ST_TABLE_P(p->hash));
-    RHASH_ST_TABLE(p->hash)->type = p->orighash;
+    struct assoc_arg *p = (struct assoc_arg*)arg;
+    st_data_t data;
+    if (st_lookup(p->tbl, p->key, &data)) return (VALUE)data;
     return Qundef;
 }
 
@@ -4127,30 +4120,30 @@ assoc_i(VALUE key, VALUE val, VALUE arg)
 static VALUE
 rb_hash_assoc(VALUE hash, VALUE key)
 {
-    st_table *table;
-    const struct st_hash_type *orighash;
     VALUE args[2];
 
     if (RHASH_EMPTY_P(hash)) return Qnil;
 
-    ar_force_convert_table(hash, __FILE__, __LINE__);
-    HASH_ASSERT(RHASH_ST_TABLE_P(hash));
-    table = RHASH_ST_TABLE(hash);
-    orighash = table->type;
-
-    if (orighash != &identhash) {
-        VALUE value;
-        struct reset_hash_type_arg ensure_arg;
-        struct st_hash_type assochash;
-
-        assochash.compare = assoc_cmp;
-        assochash.hash = orighash->hash;
-        table->type = &assochash;
-        args[0] = hash;
-        args[1] = key;
-        ensure_arg.hash = hash;
-        ensure_arg.orighash = orighash;
-        value = rb_ensure(lookup2_call, (VALUE)&args, reset_hash_type, (VALUE)&ensure_arg);
+    if (RHASH_ST_TABLE_P(hash) && RHASH_ST_TABLE(hash)->type != &identhash) {
+        VALUE value = Qundef;
+        st_table assoctable = *RHASH_ST_TABLE(hash);
+        assoctable.type = &(struct st_hash_type){
+            .compare = assoc_cmp,
+            .hash = assoctable.type->hash,
+        };
+        VALUE arg = (VALUE)&(struct assoc_arg){
+            .tbl = &assoctable,
+            .key = (st_data_t)key,
+        };
+
+        if (RB_OBJ_FROZEN(hash)) {
+            value = assoc_lookup(arg);
+        }
+        else {
+            hash_iter_lev_inc(hash);
+            value = rb_ensure(assoc_lookup, arg, hash_foreach_ensure, hash);
+        }
+        hash_verify(hash);
         if (!UNDEF_P(value)) return rb_assoc_new(key, value);
     }
 
diff --git a/test/ruby/test_hash.rb b/test/ruby/test_hash.rb
index 970d0938bf..639707d8eb 100644
--- a/test/ruby/test_hash.rb
+++ b/test/ruby/test_hash.rb
@@ -178,6 +178,24 @@ class TestHash < Test::Unit::TestCase
     assert_equal('default', h['spurious'])
   end
 
+  def test_st_literal_memory_leak
+    assert_no_memory_leak([], "", "#{<<~"begin;"}\n#{<<~'end;'}", rss: true)
+    begin;
+      1_000_000.times do
+        # >8 element hashes are ST allocated rather than AR allocated
+        {a: 1, b: 2, c: 3, d: 4, e: 5, f: 6, g: 7, h: 8, i: 9}
+      end
+    end;
+  end
+
+  def test_try_convert
+    assert_equal({1=>2}, Hash.try_convert({1=>2}))
+    assert_equal(nil, Hash.try_convert("1=>2"))
+    o = Object.new
+    def o.to_hash; {3=>4} end
+    assert_equal({3=>4}, Hash.try_convert(o))
+  end
+
   def test_AREF # '[]'
     t = Time.now
     h = @cls[
@@ -351,6 +369,10 @@ class TestHash < Test::Unit::TestCase
       end
     end
     assert_equal(base.dup, h)
+
+    h = base.dup
+    assert_same h, h.delete_if {h.assoc(nil); true}
+    assert_empty h
   end
 
   def test_keep_if