[GR-18163] Allow immediate and frozen key/values in WeakMap (#2267)

PullRequest: truffleruby/2463

Author: Nicolas Laurent

CHANGELOG.md

@@ -69,6 +69,7 @@ Compatibility:
 * Implemented `Kernel#raise` `cause` parameter.
 * Improved compatibility of `Signal.trap` and `Kernel#trap` (#2287, @chrisseaton).
 * Implemented `GC.stat(:total_allocated_objects)` as `0` (#2292, @chrisseaton).
+* `ObjectSpace::WeakMap` now supports immediate and frozen values as both keys and values (#2267).
 
 Performance:
 

spec/tags/core/objectspace/weakmap/element_set_tags.txt

@@ -71,4 +71,20 @@
     map.each_value { |v| a << v }
     a.should == ["X"]
   end
+
+  it "supports constant objects" do
+    map = ObjectSpace::WeakMap.new
+    map[1] = 2
+    Primitive.gc_force
+    map[1].should == 2
+  end
+
+  it "supports frozen objects" do
+    map = ObjectSpace::WeakMap.new
+    x = "x".freeze
+    y = "y".freeze
+    map[x] = y
+    Primitive.gc_force
+    map[x].should == y
+  end
 end

spec/truffle/objectspace/weakmap_spec.rb

@@ -117,13 +117,20 @@ protected boolean equal(VirtualFrame frame, Object a, Object b) {
     /** This node is not trivial because primitives must be compared by value and never by identity. Also, this node
      * must consider (byte) n and (short) n and (int) n and (long) n equal, as well as (float) n and (double) n. So even
      * if a and b have different classes they might still be equal if they are primitives. */
+    @GenerateUncached
+    @GenerateNodeFactory
     @CoreMethod(names = { "equal?", "==" }, required = 1)
-    public abstract static class ReferenceEqualNode extends CoreMethodArrayArgumentsNode {
+    @NodeChild(value = "arguments", type = RubyNode[].class)
+    public abstract static class ReferenceEqualNode extends RubySourceNode {
 
         public static ReferenceEqualNode create() {
             return ReferenceEqualNodeFactory.create(null);
         }
 
+        public static ReferenceEqualNode getUncached() {
+            return ReferenceEqualNodeFactory.getUncached();
+        }
+
         public abstract boolean executeReferenceEqual(Object a, Object b);
 
         @Specialization

src/main/java/org/truffleruby/core/basicobject/BasicObjectNodes.java

@@ -9,25 +9,23 @@
  */
 package org.truffleruby.core.cast;
 
+import com.oracle.truffle.api.dsl.CachedContext;
+import com.oracle.truffle.api.dsl.GenerateUncached;
+import org.truffleruby.RubyContext;
+import org.truffleruby.RubyLanguage;
 import org.truffleruby.core.numeric.RubyBignum;
-import org.truffleruby.language.RubyContextSourceNode;
-import org.truffleruby.language.RubyNode;
+import org.truffleruby.language.RubyBaseNode;
 import org.truffleruby.language.dispatch.DispatchNode;
 
 import com.oracle.truffle.api.dsl.Cached;
-import com.oracle.truffle.api.dsl.NodeChild;
 import com.oracle.truffle.api.dsl.Specialization;
 
 /** See {@link ToIntNode} for a comparison of different integer conversion nodes. */
-@NodeChild(value = "child", type = RubyNode.class)
-public abstract class ToRubyIntegerNode extends RubyContextSourceNode {
+@GenerateUncached
+public abstract class ToRubyIntegerNode extends RubyBaseNode {
 
     public static ToRubyIntegerNode create() {
-        return ToRubyIntegerNodeGen.create(null);
-    }
-
-    public static ToRubyIntegerNode create(RubyNode child) {
-        return ToRubyIntegerNodeGen.create(child);
+        return ToRubyIntegerNodeGen.create();
     }
 
     public abstract Object execute(Object object);
@@ -49,7 +47,8 @@ protected RubyBignum coerceRubyBignum(RubyBignum value) {
 
     @Specialization(guards = "!isRubyInteger(object)")
     protected Object coerceObject(Object object,
+            @CachedContext(RubyLanguage.class) RubyContext context,
             @Cached DispatchNode toIntNode) {
-        return toIntNode.call(getContext().getCoreLibrary().truffleTypeModule, "rb_to_int_fallback", object);
+        return toIntNode.call(context.getCoreLibrary().truffleTypeModule, "rb_to_int_fallback", object);
     }
 }

src/main/java/org/truffleruby/core/cast/ToRubyIntegerNode.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/hash/CompareByRubyIdentityWrapper.java

@@ -13,18 +13,22 @@
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.CachedContext;
 import com.oracle.truffle.api.dsl.Fallback;
+import com.oracle.truffle.api.dsl.GenerateUncached;
 import com.oracle.truffle.api.dsl.Specialization;
 import org.truffleruby.RubyContext;
 import org.truffleruby.RubyLanguage;
 import org.truffleruby.core.basicobject.BasicObjectNodes.ObjectIDNode;
 import org.truffleruby.core.cast.ToRubyIntegerNode;
+import org.truffleruby.core.hash.HashingNodesFactory.HashCastResultNodeGen;
+import org.truffleruby.core.hash.HashingNodesFactory.ToHashByIdentityNodeGen;
 import org.truffleruby.core.numeric.BigIntegerOps;
 import org.truffleruby.core.numeric.RubyBignum;
 import org.truffleruby.core.string.RubyString;
 import org.truffleruby.core.string.StringNodes;
 import org.truffleruby.core.symbol.RubySymbol;
 import org.truffleruby.core.symbol.SymbolNodes;
 import org.truffleruby.core.string.ImmutableRubyString;
+import org.truffleruby.language.RubyBaseNode;
 import org.truffleruby.language.RubyContextNode;
 import org.truffleruby.language.dispatch.DispatchNode;
 
@@ -137,7 +141,16 @@ private int castResult(Object value) {
 
     }
 
-    public abstract static class ToHashByIdentity extends RubyContextNode {
+    @GenerateUncached
+    public abstract static class ToHashByIdentity extends RubyBaseNode {
+
+        public static ToHashByIdentity create() {
+            return ToHashByIdentityNodeGen.create();
+        }
+
+        public static ToHashByIdentity getUncached() {
+            return ToHashByIdentityNodeGen.getUncached();
+        }
 
         public abstract int execute(Object key);
 
@@ -147,19 +160,15 @@ protected int toHashByIdentity(Object hashed,
                 @Cached HashCastResultNode hashCastResultNode) {
             return hashCastResultNode.execute(objectIDNode.execute(hashed));
         }
-
     }
 
-
-    public abstract static class HashCastResultNode extends RubyContextNode {
-
-        @Child private ToRubyIntegerNode toRubyInteger;
-        @Child private HashCastResultNode hashCastResultNode;
+    @GenerateUncached
+    public abstract static class HashCastResultNode extends RubyBaseNode {
 
         public abstract int execute(Object key);
 
         public static HashCastResultNode create() {
-            return HashingNodesFactory.HashCastResultNodeGen.create();
+            return HashCastResultNodeGen.create();
         }
 
         @Specialization
@@ -177,24 +186,11 @@ protected int castBignum(RubyBignum hashed) {
             return BigIntegerOps.hashCode(hashed);
         }
 
-        @Fallback
-        protected int castOther(Object hashed) {
-            if (toRubyInteger == null) {
-                CompilerDirectives.transferToInterpreterAndInvalidate();
-                toRubyInteger = insert(ToRubyIntegerNode.create());
-            }
-            final Object coercedHashedObject = toRubyInteger.execute(hashed);
-            return castCoerced(coercedHashedObject);
+        @Specialization(guards = "!isRubyInteger(hashed)")
+        protected int castOther(Object hashed,
+                @Cached ToRubyIntegerNode toRubyInteger,
+                @Cached HashCastResultNode hashCastResult) {
+            return hashCastResult.execute(toRubyInteger.execute(hashed));
         }
-
-        private int castCoerced(Object coerced) {
-            if (hashCastResultNode == null) {
-                CompilerDirectives.transferToInterpreterAndInvalidate();
-                hashCastResultNode = insert(HashCastResultNode.create());
-            }
-            return hashCastResultNode.execute(coerced);
-        }
-
     }
-
 }

src/main/java/org/truffleruby/core/hash/HashingNodes.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/WeakMapNodes.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> {
 }

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

@@ -9,29 +9,11 @@
 # GNU Lesser General Public License version 2.1.
 
 module ObjectSpace
-  # WeakMap uses identity comparison semantics. The implementation assumes that the Java representation of objects
-  # do compare (equals() and hashCode()) using object identity. This is the case for instances of RubyDynamicObject.
-  #
-  # However, results are unspecified if used with instances of TruffleObject that override equals() and
-  # hashCode().
-  #
-  # Note that, following MRI, we disallow immediate (primitive) values and frozen objects as keys or value of
-  # WeakMap. We could easily transcend this limitation as we do not modify objects like MRI does (it sets a finalizer).
-  # However, keeping the limitation enables the above assumption of identity comparison.
+  # WeakMap uses Ruby identity semantics to compare and hash keys.
   class WeakMap
     include Enumerable
 
-    private def check_key_or_value(kv, type)
-      if Primitive.immediate_value?(kv)
-        raise ArgumentError, "WeakMap #{type} can't be an instance of #{kv.class}"
-      elsif Truffle::KernelOperations.value_frozen?(kv)
-        raise FrozenError, "WeakMap #{type} can't be a frozen object"
-      end
-    end
-
     def []=(key, value)
-      check_key_or_value(key, 'key')
-      check_key_or_value(value, 'value')
       Primitive.weakmap_aset(self, key, value)
     end