1 files changed, 80 insertions, 3 deletions

diff --git a/test/ruby/test_weakmap.rb b/test/ruby/test_weakmap.rb
index c72e7310db..4f5823ecf4 100644
--- a/test/ruby/test_weakmap.rb
+++ b/test/ruby/test_weakmap.rb
@@ -78,10 +78,26 @@ class TestWeakMap < Test::Unit::TestCase
       @wm[i] = Object.new
       @wm.inspect
     end
-    assert_match(/\A\#<#{@wm.class.name}:[^:]++:(?:\s\d+\s=>\s\#<(?:Object|collected):[^:<>]*+>(?:,|>\z))+/,
+    assert_match(/\A\#<#{@wm.class.name}:0x[\da-f]+(?::(?: \d+ => \#<(?:Object|collected):0x[\da-f]+>,?)+)?>\z/,
                  @wm.inspect)
   end
 
+  def test_delete
+    k1 = "foo"
+    x1 = Object.new
+    @wm[k1] = x1
+    assert_equal x1, @wm[k1]
+    assert_equal x1, @wm.delete(k1)
+    assert_nil @wm[k1]
+    assert_nil @wm.delete(k1)
+
+    fallback =  @wm.delete(k1) do |key|
+      assert_equal k1, key
+      42
+    end
+    assert_equal 42, fallback
+  end
+
   def test_each
     m = __callee__[/test_(.*)/, 1]
     x1 = Object.new
@@ -177,14 +193,17 @@ class TestWeakMap < Test::Unit::TestCase
     end;
   end
 
-  def test_compaction_bug_19529
+  def test_compaction
+    omit "compaction is not supported on this platform" unless GC.respond_to?(:compact)
+
+    # [Bug #19529]
     obj = Object.new
     100.times do |i|
       GC.compact
       @wm[i] = obj
     end
 
-    assert_separately(%w(--disable-gems), <<-'end;')
+    assert_ruby_status([], <<-'end;')
       wm = ObjectSpace::WeakMap.new
       obj = Object.new
       100.times do
@@ -193,6 +212,34 @@ class TestWeakMap < Test::Unit::TestCase
       end
       GC.compact
     end;
+
+    assert_separately(%w(-robjspace), <<-'end;')
+      wm = ObjectSpace::WeakMap.new
+      key = Object.new
+      val = Object.new
+      wm[key] = val
+
+      GC.verify_compaction_references(expand_heap: true, toward: :empty)
+
+      assert_equal(val, wm[key])
+    end;
+
+    assert_ruby_status(["-W0"], <<-'end;')
+      wm = ObjectSpace::WeakMap.new
+
+      ary = 10_000.times.map do
+        o = Object.new
+        wm[o] = 1
+        o
+      end
+
+      GC.verify_compaction_references(expand_heap: true, toward: :empty)
+    end;
+  end
+
+  def test_gc_compact_stress
+    omit "compaction doesn't work well on s390x" if RUBY_PLATFORM =~ /s390x/ # https://github.com/ruby/ruby/pull/5077
+    EnvUtil.under_gc_compact_stress { ObjectSpace::WeakMap.new }
   end
 
   def test_replaced_values_bug_19531
@@ -211,4 +258,34 @@ class TestWeakMap < Test::Unit::TestCase
 
     assert_equal b, @wm[1]
   end
+
+  def test_use_after_free_bug_20688
+    assert_normal_exit(<<~RUBY)
+      weakmap = ObjectSpace::WeakMap.new
+      10_000.times { weakmap[Object.new] = Object.new }
+    RUBY
+  end
+
+  def test_generational_gc
+    EnvUtil.without_gc do
+      wmap = ObjectSpace::WeakMap.new
+
+      (GC::INTERNAL_CONSTANTS[:RVALUE_OLD_AGE] - 1).times { GC.start }
+
+      retain = []
+      50.times do
+        k = Object.new
+        wmap[k] = true
+        retain << k
+      end
+
+      GC.start # WeakMap promoted, other objects still young
+
+      retain.clear
+
+      GC.start(full_mark: false)
+
+      wmap.keys.each(&:itself) # call method on keys to cause crash
+    end
+  end
 end