[GR-59866] Fix WeakKeyMap

PullRequest: truffleruby/4482

Author: andrykonchin

CHANGELOG.md

@@ -10,6 +10,7 @@ Compatibility:
 
 * Implement `StringScanner#{peek_byte,scan_byte,scan_integer,named_captures}` methods (#3788, @andrykonchin).
 * Support String patterns in `StringScanner#{exist?,scan_until,skip_until,check_until,search_full}` methods (@andrykonchin).
+* Implement `ObjectSpace::WeakKeyMap` (#3681, @andrykonchin).
 
 Performance:
 

spec/ruby/core/objectspace/weakkeymap/element_reference_spec.rb

@@ -18,14 +18,18 @@
 
     it "compares keys with #eql? semantics" do
       map = ObjectSpace::WeakKeyMap.new
-      map[[1.0]] = "x"
+      key = [1.0]
+      map[key] = "x"
       map[[1]].should == nil
       map[[1.0]].should == "x"
+      key.should == [1.0] # keep the key alive until here to keep the map entry
 
       map = ObjectSpace::WeakKeyMap.new
-      map[[1]] = "x"
+      key = [1]
+      map[key] = "x"
       map[[1.0]].should == nil
       map[[1]].should == "x"
+      key.should == [1] # keep the key alive until here to keep the map entry
 
       map = ObjectSpace::WeakKeyMap.new
       key1, key2 = %w[a a].map(&:upcase)
@@ -39,7 +43,8 @@
       x.should_receive(:hash).and_return(0)
 
       map = ObjectSpace::WeakKeyMap.new
-      map['foo'] = :bar
+      key = 'foo'
+      map[key] = :bar
       map[x].should == nil
     end
 

spec/ruby/core/objectspace/weakkeymap/element_set_spec.rb

@@ -43,6 +43,8 @@ def should_accept(map, key, value)
 
       map.getkey("a").should.equal? key
       map.getkey("a").should_not.frozen?
+
+      key.should == "a" # keep the key alive until here to keep the map entry
     end
 
     context "a key cannot be garbage collected" do

spec/ruby/core/objectspace/weakkeymap/getkey_spec.rb

@@ -9,6 +9,9 @@
       map[key1] = true
       map.getkey(key2).should equal(key1)
       map.getkey("X").should == nil
+
+      key1.should == "A" # keep the key alive until here to keep the map entry
+      key2.should == "A" # keep the key alive until here to keep the map entry
     end
 
     it "returns nil when a key cannot be garbage collected" do

spec/ruby/core/objectspace/weakkeymap/inspect_spec.rb

@@ -12,6 +12,10 @@
       map.inspect.should =~ /\A\#<ObjectSpace::WeakKeyMap:0x\h+ size=2>\z/
       map[key3] = 3
       map.inspect.should =~ /\A\#<ObjectSpace::WeakKeyMap:0x\h+ size=2>\z/
+
+      key1.should == "foo" # keep the key alive until here to keep the map entry
+      key2.should == "bar" # keep the key alive until here to keep the map entry
+      key3.should == "bar" # keep the key alive until here to keep the map entry
     end
   end
 end

spec/truffle/objectspace/weak_key_map_spec.rb

@@ -0,0 +1,58 @@
+# truffleruby_primitives: true
+#   (to nil out references to make unreachable)
+
+# Copyright (c) 2020, 2025 Oracle and/or its affiliates. All rights reserved. This
+# code is released under a tri EPL/GPL/LGPL license. You can use it,
+# redistribute it and/or modify it under the terms of the:
+#
+# Eclipse Public License version 2.0, or
+# GNU General Public License version 2, or
+# GNU Lesser General Public License version 2.1.
+
+require_relative '../../ruby/spec_helper'
+
+describe "ObjectSpace::WeakKeyMap" do
+
+  it "gets rid of unreferenced keys" do
+    map = ObjectSpace::WeakKeyMap.new
+    key = "a".upcase
+    value = "x"
+    map[key] = value
+    key = nil
+
+    Primitive.gc_force
+
+    map[key].should == nil
+    map.key?(key).should == false
+    map.getkey(key).should == nil
+  end
+
+  it "has iterators methods that exclude unreferenced objects" do
+    # This spec does not pass on MRI because the garbage collector is presumably too conservative and will not get rid
+    # of the references eagerly enough.
+
+    map = ObjectSpace::WeakKeyMap.new
+    k1, k2 = %w[a b].map(&:upcase)
+    v1, v2 = %w[x y].map(&:upcase)
+    map[k1] = v1
+    map[k2] = v2
+    k2 = nil
+
+    Primitive.gc_force
+
+    map.key?(k2).should == false
+    map[k2].should == nil
+
+    map.key?(k1).should == true
+    map[k1].should == v1
+  end
+
+  it "supports frozen objects" do
+    map = ObjectSpace::WeakKeyMap.new
+    k = "x".freeze
+    v = "y".freeze
+    map[k] = v
+    Primitive.gc_force
+    map[k].should == v
+  end
+end

src/main/java/org/truffleruby/collections/ConcurrentWeakKeysMap.java

@@ -94,10 +94,14 @@ public Value remove(Key key) {
         return map.remove(buildWeakReference(key));
     }
 
+    /** It's important a WeakReference that is built returns true for {@code ref.equals(ref)} even if the reference is
+     * cleared. */
     protected WeakReference<Key> buildWeakReference(Key key) {
         return new WeakKeyReference<>(key);
     }
 
+    /** It's important a WeakReference that is built returns true for {@code ref.equals(ref)} even if the reference is
+     * cleared. */
     protected WeakReference<Key> buildWeakReference(Key key, ReferenceQueue<Key> referenceQueue) {
         return new WeakKeyReference<>(key, referenceQueue);
     }
@@ -108,14 +112,16 @@ protected WeakReference<Key> buildWeakReference(Key key, ReferenceQueue<Key> ref
      * It is possible that the map still contains {@link WeakReference} instances whose key has been nulled out after a
      * call to this method (the reference not having been enqueued yet)! */
     protected void removeStaleEntries() {
-        WeakKeyReference<?> ref;
-        while ((ref = (WeakKeyReference<?>) referenceQueue.poll()) != null) {
+        WeakReference<?> ref;
+        while ((ref = (WeakReference<?>) referenceQueue.poll()) != null) {
             // Here ref.get() is null, so it will not remove a new key-value pair with the same key
-            // as that is a different WeakKeyReference instance.
+            // as that is a different WeakReference instance.
             map.remove(ref);
         }
     }
 
+    /** A default implementation of a key that wraps in a user-provided key. Compares keys by
+     * {@link Object#equals(Object)} */
     protected static final class WeakKeyReference<Key> extends WeakReference<Key> {
         private final int hashCode;