add a wrapper around keys in WeakMap to ensure correct ruby reference comparison

& hashing semantics

Author: Nicolas Laurent

src/main/java/org/truffleruby/core/hash/CompareByRubyIdentityWrapper.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2013, 2021 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.
+ */
+package org.truffleruby.core.hash;
+
+import org.truffleruby.core.basicobject.BasicObjectNodes.ReferenceEqualNode;
+import org.truffleruby.core.hash.HashingNodes.ToHashByIdentity;
+
+/** Wraps a value so that it will compared and hashed according to Ruby identity semantics. These semantics differ from
+ * Java semantics notably for primitives (e.g. all the NaN are different according to Ruby, but Java compares them
+ * equal). */
+public final class CompareByRubyIdentityWrapper {
+
+    public final Object value;
+
+    public CompareByRubyIdentityWrapper(Object value) {
+        this.value = value;
+    }
+
+    @Override
+    public int hashCode() {
+        return ToHashByIdentity.getUncached().execute(value);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        return obj instanceof CompareByRubyIdentityWrapper &&
+                ReferenceEqualNode
+                        .getUncached()
+                        .executeReferenceEqual(value, ((CompareByRubyIdentityWrapper) obj).value);
+    }
+}

src/main/java/org/truffleruby/core/objectspace/WeakMapNodes.java

@@ -15,8 +15,9 @@
 import org.truffleruby.builtins.Primitive;
 import org.truffleruby.builtins.UnaryCoreMethodNode;
 import org.truffleruby.builtins.YieldingCoreMethodNode;
-import org.truffleruby.collections.WeakValueCache;
+import org.truffleruby.collections.WeakValueCache.WeakMapEntry;
 import org.truffleruby.core.array.RubyArray;
+import org.truffleruby.core.hash.CompareByRubyIdentityWrapper;
 import org.truffleruby.core.klass.RubyClass;
 import org.truffleruby.core.proc.RubyProc;
 import org.truffleruby.language.Nil;
@@ -27,6 +28,8 @@
 import com.oracle.truffle.api.dsl.Specialization;
 import org.truffleruby.language.objects.AllocationTracing;
 
+import java.util.Collection;
+
 /** Note that WeakMap uses identity comparison semantics. See top comment in src/main/ruby/truffleruby/core/weakmap.rb
  * for more information. */
 @CoreModule(value = "ObjectSpace::WeakMap", isClass = true)
@@ -57,7 +60,7 @@ public abstract static class MemberNode extends CoreMethodArrayArgumentsNode {
 
         @Specialization
         protected boolean isMember(RubyWeakMap map, Object key) {
-            return map.storage.get(key) != null;
+            return map.storage.get(new CompareByRubyIdentityWrapper(key)) != null;
         }
     }
 
@@ -66,7 +69,7 @@ public abstract static class GetIndexNode extends CoreMethodArrayArgumentsNode {
 
         @Specialization
         protected Object get(RubyWeakMap map, Object key) {
-            Object value = map.storage.get(key);
+            Object value = map.storage.get(new CompareByRubyIdentityWrapper(key));
             return value == null ? nil : value;
         }
     }
@@ -76,7 +79,7 @@ public abstract static class SetIndexNode extends CoreMethodArrayArgumentsNode {
 
         @Specialization
         protected Object set(RubyWeakMap map, Object key, Object value) {
-            map.storage.put(key, value);
+            map.storage.put(new CompareByRubyIdentityWrapper(key), value);
             return value;
         }
     }
@@ -144,8 +147,8 @@ protected RubyWeakMap each(RubyWeakMap map, Nil block) {
         @Specialization
         protected RubyWeakMap each(RubyWeakMap map, RubyProc block) {
 
-            for (WeakValueCache.WeakMapEntry<?, ?> e : entries(map.storage)) {
-                yield(block, e.getKey(), e.getValue());
+            for (MapEntry entry : entries(map.storage)) {
+                yield(block, entry.key, entry.value);
             }
 
             return map;
@@ -154,17 +157,39 @@ protected RubyWeakMap each(RubyWeakMap map, RubyProc block) {
 
     @TruffleBoundary
     private static Object[] keys(WeakMapStorage storage) {
-        return storage.keys().toArray();
+        final Collection<CompareByRubyIdentityWrapper> keyWrappers = storage.keys();
+        final Object[] keys = new Object[keyWrappers.size()];
+        int i = 0;
+        for (CompareByRubyIdentityWrapper keyWrapper : keyWrappers) {
+            keys[i++] = keyWrapper.value;
+        }
+        return keys;
     }
 
     @TruffleBoundary
     private static Object[] values(WeakMapStorage storage) {
         return storage.values().toArray();
     }
 
+    private static class MapEntry {
+        final Object key;
+        final Object value;
+
+        private MapEntry(Object key, Object value) {
+            this.key = key;
+            this.value = value;
+        }
+    }
+
     @TruffleBoundary
-    private static WeakValueCache.WeakMapEntry<?, ?>[] entries(WeakMapStorage storage) {
-        return storage.entries().toArray(new WeakValueCache.WeakMapEntry<?, ?>[0]);
+    private static MapEntry[] entries(WeakMapStorage storage) {
+        final Collection<WeakMapEntry<CompareByRubyIdentityWrapper, Object>> wrappedEntries = storage.entries();
+        final MapEntry[] entries = new MapEntry[wrappedEntries.size()];
+        int i = 0;
+        for (WeakMapEntry<CompareByRubyIdentityWrapper, Object> wrappedEntry : wrappedEntries) {
+            entries[i++] = new MapEntry(wrappedEntry.getKey().value, wrappedEntry.getValue());
+        }
+        return entries;
     }
 
     private static RubyWeakMap eachNoBlockProvided(YieldingCoreMethodNode node, RubyWeakMap map) {

src/main/java/org/truffleruby/core/objectspace/WeakMapStorage.java

@@ -10,6 +10,7 @@
 package org.truffleruby.core.objectspace;
 
 import org.truffleruby.collections.WeakValueCache;
+import org.truffleruby.core.hash.CompareByRubyIdentityWrapper;
 
-public final class WeakMapStorage extends WeakValueCache<Object, Object> {
+public final class WeakMapStorage extends WeakValueCache<CompareByRubyIdentityWrapper, Object> {
 }