jruby · Aug 24, 2015 · Aug 24, 2015
diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/core/ObjectSpaceNodes.java b/truffle/src/main/java/org/jruby/truffle/nodes/core/ObjectSpaceNodes.java
@@ -9,14 +9,9 @@
  */
 package org.jruby.truffle.nodes.core;
 
-import com.oracle.truffle.api.CompilerDirectives;
-import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
-import com.oracle.truffle.api.dsl.Specialization;
-import com.oracle.truffle.api.frame.VirtualFrame;
-import com.oracle.truffle.api.object.DynamicObject;
-import com.oracle.truffle.api.source.SourceSection;
 import org.jruby.truffle.nodes.RubyGuards;
 import org.jruby.truffle.nodes.dispatch.RespondToNode;
+import org.jruby.truffle.nodes.objectstorage.ReadHeadObjectFieldNode;
 import org.jruby.truffle.runtime.ModuleOperations;
 import org.jruby.truffle.runtime.NotProvided;
 import org.jruby.truffle.runtime.RubyContext;
@@ -28,63 +23,82 @@
 import org.jruby.truffle.runtime.object.StopVisitingObjectsException;
 import org.jruby.util.Memo;
 
+import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.ImportStatic;
+import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.object.DynamicObject;
+import com.oracle.truffle.api.source.SourceSection;
+
 @CoreClass(name = "ObjectSpace")
 public abstract class ObjectSpaceNodes {
 
     @CoreMethod(names = "_id2ref", isModuleFunction = true, required = 1)
+    @ImportStatic(ObjectIDOperations.class)
     public abstract static class ID2RefNode extends CoreMethodArrayArgumentsNode {
 
         public ID2RefNode(RubyContext context, SourceSection sourceSection) {
             super(context, sourceSection);
         }
 
-        @Specialization
-        public Object id2Ref(int id) {
-            return id2Ref((long) id);
+        @Specialization(guards = "id == NIL")
+        public Object id2RefNil(long id) {
+            return nil();
         }
 
-        @TruffleBoundary
-        @Specialization
-        public Object id2Ref(final long id) {
-            if (id == ObjectIDOperations.NIL) {
-                return nil();
-            } else if (id == ObjectIDOperations.TRUE) {
-                return true;
-            } else if (id == ObjectIDOperations.FALSE) {
-                return false;
-            } else if (ObjectIDOperations.isSmallFixnumID(id)) {
-                return ObjectIDOperations.toFixnum(id);
-            } else {
-                final Memo<Object> result = new Memo<Object>(nil());
+        @Specialization(guards = "id == TRUE")
+        public boolean id2RefTrue(long id) {
+            return true;
+        }
 
-                new ObjectGraph(getContext()).visitObjects(new ObjectGraphVisitor() {
+        @Specialization(guards = "id == FALSE")
+        public boolean id2RefFalse(long id) {
+            return false;
+        }
 
-                    @Override
-                    public boolean visit(DynamicObject object) throws StopVisitingObjectsException {
-                        if (ObjectIDOperations.verySlowGetObjectID(object) == id) {
-                            result.set(object);
-                            throw new StopVisitingObjectsException();
-                        }
+        @Specialization(guards = "isSmallFixnumID(id)")
+        public long id2RefSmallInt(long id) {
+            return ObjectIDOperations.toFixnum(id);
+        }
 
-                        return true;
-                    }
+        @Specialization(guards = "isBasicObjectID(id)")
+        public Object id2Ref(final VirtualFrame frame, final long id,
+                @Cached("createReadObjectIDNode()") final ReadHeadObjectFieldNode readObjectIdNode) {
+            CompilerDirectives.transferToInterpreter();
 
-                });
+            final Memo<Object> result = new Memo<Object>(nil());
 
-                return result.get();
-            }
+            new ObjectGraph(getContext()).visitObjects(new ObjectGraphVisitor() {
+                @Override
+                public boolean visit(DynamicObject object) throws StopVisitingObjectsException {
+                    final Object objectID = readObjectIdNode.execute(object);
+                    if (objectID != nil() && (long) objectID == id) {
+                        result.set(object);
+                        throw new StopVisitingObjectsException();
+                    }
+                    return true;
+                }
+            });
+
+            return result.get();
         }
 
-        @Specialization(guards = {"isRubyBignum(id)", "isLargeFixnumID(id)"})
+        @Specialization(guards = { "isRubyBignum(id)", "isLargeFixnumID(id)" })
         public Object id2RefLargeFixnum(DynamicObject id) {
             return Layouts.BIGNUM.getValue(id).longValue();
         }
 
-        @Specialization(guards = {"isRubyBignum(id)", "isFloatID(id)"})
+        @Specialization(guards = { "isRubyBignum(id)", "isFloatID(id)" })
         public double id2RefFloat(DynamicObject id) {
             return Double.longBitsToDouble(Layouts.BIGNUM.getValue(id).longValue());
         }
 
+        protected ReadHeadObjectFieldNode createReadObjectIDNode() {
+            return new ReadHeadObjectFieldNode(Layouts.OBJECT_ID_IDENTIFIER);
+        }
+
         protected boolean isLargeFixnumID(DynamicObject id) {
             return ObjectIDOperations.isLargeFixnumID(Layouts.BIGNUM.getValue(id));
         }

diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/rubinius/ObjectPrimitiveNodes.java b/truffle/src/main/java/org/jruby/truffle/nodes/rubinius/ObjectPrimitiveNodes.java
@@ -34,9 +34,6 @@ public abstract class ObjectPrimitiveNodes {
     @RubiniusPrimitive(name = "object_id")
     public abstract static class ObjectIDPrimitiveNode extends RubiniusPrimitiveNode {
 
-        @Child private ReadHeadObjectFieldNode readObjectIdNode;
-        @Child private WriteHeadObjectFieldNode writeObjectIdNode;
-
         public ObjectIDPrimitiveNode(RubyContext context, SourceSection sourceSection) {
             super(context, sourceSection);
         }
@@ -71,7 +68,7 @@ public long objectIDSmallFixnumOverflow(long value) throws ArithmeticException {
         @Specialization
         public Object objectID(long value,
                                @Cached("createCountingProfile()") ConditionProfile smallProfile) {
-            if (smallProfile.profile(isSmallFixnum(value))) {
+            if (smallProfile.profile(ObjectIDOperations.isSmallFixnum(value))) {
                 return ObjectIDOperations.smallFixnumToID(value);
             } else {
                 return ObjectIDOperations.largeFixnumToID(getContext(), value);
@@ -84,30 +81,26 @@ public Object objectID(double value) {
         }
 
         @Specialization
-        public Object objectID(DynamicObject object) {
-            if (readObjectIdNode == null) {
-                CompilerDirectives.transferToInterpreter();
-                readObjectIdNode = insert(new ReadHeadObjectFieldNode(Layouts.OBJECT_ID_IDENTIFIER));
-            }
-
+        public long objectID(DynamicObject object,
+                @Cached("createReadObjectIDNode()") ReadHeadObjectFieldNode readObjectIdNode,
+                @Cached("createWriteObjectIDNode()") WriteHeadObjectFieldNode writeObjectIdNode) {
             final Object id = readObjectIdNode.execute(object);
 
             if (id == nil()) {
-                if (writeObjectIdNode == null) {
-                    CompilerDirectives.transferToInterpreter();
-                    writeObjectIdNode = insert(new WriteHeadObjectFieldNode(Layouts.OBJECT_ID_IDENTIFIER));
-                }
-
-                final Long newId = getContext().getNextObjectID();
+                final long newId = getContext().getNextObjectID();
                 writeObjectIdNode.execute(object, newId);
                 return newId;
             }
 
-            return id;
+            return (long) id;
+        }
+
+        protected ReadHeadObjectFieldNode createReadObjectIDNode() {
+            return new ReadHeadObjectFieldNode(Layouts.OBJECT_ID_IDENTIFIER);
         }
 
-        protected boolean isSmallFixnum(long fixnum) {
-            return ObjectIDOperations.isSmallFixnum(fixnum);
+        protected WriteHeadObjectFieldNode createWriteObjectIDNode() {
+            return new WriteHeadObjectFieldNode(Layouts.OBJECT_ID_IDENTIFIER);
         }
 
     }

diff --git a/truffle/src/main/java/org/jruby/truffle/runtime/RubyContext.java b/truffle/src/main/java/org/jruby/truffle/runtime/RubyContext.java
@@ -382,11 +382,10 @@ public Object execute(Source source, Encoding defaultEncoding, TranslatorDriver.
     }
 
     public long getNextObjectID() {
-        // TODO(CS): We can theoretically run out of long values
-
         final long id = nextObjectID.getAndAdd(2);
 
         if (id < 0) {
+            CompilerDirectives.transferToInterpreter();
             nextObjectID.set(Long.MIN_VALUE);
             throw new RuntimeException("Object IDs exhausted");
         }

diff --git a/truffle/src/main/java/org/jruby/truffle/runtime/object/ObjectIDOperations.java b/truffle/src/main/java/org/jruby/truffle/runtime/object/ObjectIDOperations.java
@@ -97,6 +97,10 @@ public static boolean isFloatID(BigInteger id) {
         return !id.and(FLOAT_FLAG).equals(BigInteger.ZERO);
     }
 
+    public static boolean isBasicObjectID(long id) {
+        return id >= FIRST_OBJECT_ID && id % 2 == 0;
+    }
+
     private static BigInteger unsignedBigInteger(long value) {
         BigInteger big = BigInteger.valueOf(value);
         if (value < 0) {