Improve performance of class variables

CHANGELOG.md

@@ -59,6 +59,7 @@ Performance:
 * Improve the performance of checks for recursion (#2189, @LillianZ).
 * Improve random number generation performance by avoiding synchronization (#2190, @ivoanjo).
 * We now create a single call target per block by default instead of two.
+* Some uses of class variables are now much better optimized (#2259, @chrisseaton).
 
 Changes:
 

src/main/java/org/truffleruby/RubyLanguage.java

@@ -82,6 +82,7 @@
 import org.truffleruby.language.RubyInlineParsingRequestNode;
 import org.truffleruby.language.RubyParsingRequestNode;
 import org.truffleruby.language.objects.RubyObjectType;
+import org.truffleruby.language.objects.classvariables.ClassVariableStorage;
 import org.truffleruby.options.LanguageOptions;
 import org.truffleruby.platform.Platform;
 import org.truffleruby.shared.Metrics;
@@ -211,6 +212,8 @@ public final class RubyLanguage extends TruffleLanguage<RubyContext> {
     public final Shape unboundMethodShape = createShape(RubyUnboundMethod.class);
     public final Shape weakMapShape = createShape(RubyWeakMap.class);
 
+    public final Shape classVariableShape = Shape.newBuilder().layout(ClassVariableStorage.class).build();
+
     public RubyLanguage() {
         coreMethodAssumptions = new CoreMethodAssumptions(this);
         coreStrings = new CoreStrings(this);

src/main/java/org/truffleruby/core/module/ModuleFields.java

@@ -21,6 +21,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.locks.ReentrantLock;
 
+import com.oracle.truffle.api.object.DynamicObjectLibrary;
 import org.jcodings.specific.UTF8Encoding;
 import org.truffleruby.RubyContext;
 import org.truffleruby.RubyLanguage;
@@ -46,6 +47,7 @@
 import org.truffleruby.language.methods.InternalMethod;
 import org.truffleruby.language.objects.ObjectGraph;
 import org.truffleruby.language.objects.ObjectGraphNode;
+import org.truffleruby.language.objects.classvariables.ClassVariableStorage;
 import org.truffleruby.language.objects.shared.SharedObjects;
 
 import com.oracle.truffle.api.Assumption;
@@ -95,7 +97,7 @@ public static void debugModuleChain(RubyModule module) {
 
     private final ConcurrentMap<String, InternalMethod> methods = new ConcurrentHashMap<>();
     private final ConcurrentMap<String, RubyConstant> constants = new ConcurrentHashMap<>();
-    private final ConcurrentMap<String, Object> classVariables = new ConcurrentHashMap<>();
+    private final ClassVariableStorage classVariables;
 
     /** The refinements (calls to Module#refine) nested under/contained in this namespace module (M). Represented as a
      * map of refined classes and modules (C) to refinement modules (R). */
@@ -121,6 +123,7 @@ public ModuleFields(
         this.rubyModule = rubyModule;
         this.methodsUnmodifiedAssumption = new CyclicAssumption("methods are unmodified");
         this.constantsUnmodifiedAssumption = new CyclicAssumption("constants are unmodified");
+        classVariables = new ClassVariableStorage(language);
         start = new PrependMarker(this);
     }
 
@@ -193,7 +196,12 @@ public void initCopy(RubyModule from) {
             this.constants.put(entry.getKey(), entry.getValue());
         }
 
-        this.classVariables.putAll(fromFields.classVariables);
+        for (Object key : fromFields.classVariables.getShape().getKeys()) {
+            DynamicObjectLibrary.getUncached().put(
+                    this.classVariables,
+                    key,
+                    DynamicObjectLibrary.getUncached().getOrDefault(fromFields.classVariables, key, null));
+        }
 
         if (fromFields.hasPrependedModules()) {
             // Then the parent is the first in the prepend chain
@@ -745,7 +753,7 @@ public InternalMethod getMethod(String name) {
         return methods.get(name);
     }
 
-    public ConcurrentMap<String, Object> getClassVariables() {
+    public ClassVariableStorage getClassVariables() {
         return classVariables;
     }
 
@@ -849,9 +857,7 @@ public void getAdjacentObjects(Set<Object> adjacent) {
             ObjectGraph.addProperty(adjacent, constant);
         }
 
-        for (Object value : classVariables.values()) {
-            ObjectGraph.addProperty(adjacent, value);
-        }
+        ObjectGraph.addProperty(adjacent, classVariables);
 
         for (InternalMethod method : methods.values()) {
             ObjectGraph.addProperty(adjacent, method);

src/main/java/org/truffleruby/core/module/ModuleNodes.java

@@ -19,6 +19,8 @@
 import java.util.concurrent.ConcurrentMap;
 
 import com.oracle.truffle.api.library.CachedLibrary;
+import com.oracle.truffle.api.object.DynamicObjectLibrary;
+import com.oracle.truffle.api.profiles.ConditionProfile;
 import org.jcodings.specific.UTF8Encoding;
 import org.truffleruby.RubyContext;
 import org.truffleruby.RubyLanguage;
@@ -60,7 +62,6 @@
 import org.truffleruby.core.string.StringUtils;
 import org.truffleruby.core.support.TypeNodes;
 import org.truffleruby.core.symbol.RubySymbol;
-import org.truffleruby.core.symbol.SymbolTable;
 import org.truffleruby.language.LexicalScope;
 import org.truffleruby.language.Nil;
 import org.truffleruby.language.NotProvided;
@@ -103,6 +104,10 @@
 import org.truffleruby.language.objects.ReadInstanceVariableNode;
 import org.truffleruby.language.objects.SingletonClassNode;
 import org.truffleruby.language.objects.WriteInstanceVariableNode;
+import org.truffleruby.language.objects.classvariables.CheckClassVariableNameNode;
+import org.truffleruby.language.objects.classvariables.ClassVariableStorage;
+import org.truffleruby.language.objects.classvariables.LookupClassVariableNode;
+import org.truffleruby.language.objects.classvariables.SetClassVariableNode;
 import org.truffleruby.language.yield.CallBlockNode;
 import org.truffleruby.parser.Identifiers;
 import org.truffleruby.parser.ParserContext;
@@ -804,14 +809,12 @@ protected RubyNode coerceToString(RubyNode name) {
             return NameToJavaStringNode.create(name);
         }
 
-        @TruffleBoundary
         @Specialization
-        protected boolean isClassVariableDefinedString(RubyModule module, String name) {
-            SymbolTable.checkClassVariableName(getContext(), name, module, this);
-
-            final Object value = ModuleOperations.lookupClassVariable(module, name);
-
-            return value != null;
+        protected boolean isClassVariableDefinedString(RubyModule module, String name,
+                @Cached CheckClassVariableNameNode checkClassVariableNameNode,
+                @Cached LookupClassVariableNode lookupClassVariableNode) {
+            checkClassVariableNameNode.execute(module, name);
+            return lookupClassVariableNode.execute(module, name) != null;
         }
 
     }
@@ -827,13 +830,14 @@ protected RubyNode coerceToString(RubyNode name) {
         }
 
         @Specialization
-        @TruffleBoundary
-        protected Object getClassVariable(RubyModule module, String name) {
-            SymbolTable.checkClassVariableName(getContext(), name, module, this);
+        protected Object getClassVariable(RubyModule module, String name,
+                @Cached CheckClassVariableNameNode checkClassVariableNameNode,
+                @Cached LookupClassVariableNode lookupClassVariableNode,
+                @Cached ConditionProfile undefinedProfile) {
+            checkClassVariableNameNode.execute(module, name);
+            final Object value = lookupClassVariableNode.execute(module, name);
 
-            final Object value = ModuleOperations.lookupClassVariable(module, name);
-
-            if (value == null) {
+            if (undefinedProfile.profile(value == null)) {
                 throw new RaiseException(
                         getContext(),
                         coreExceptions().nameErrorUninitializedClassVariable(module, name, this));
@@ -856,12 +860,11 @@ protected RubyNode coerceToString(RubyNode name) {
         }
 
         @Specialization
-        @TruffleBoundary
-        protected Object setClassVariable(RubyModule module, String name, Object value) {
-            SymbolTable.checkClassVariableName(getContext(), name, module, this);
-
-            ModuleOperations.setClassVariable(getLanguage(), getContext(), module, name, value, this);
-
+        protected Object setClassVariable(RubyModule module, String name, Object value,
+                @Cached CheckClassVariableNameNode checkClassVariableNameNode,
+                @Cached SetClassVariableNode setClassVariableNode) {
+            checkClassVariableNameNode.execute(module, name);
+            setClassVariableNode.execute(module, name, value);
             return value;
         }
 
@@ -873,12 +876,25 @@ public abstract static class ClassVariablesNode extends CoreMethodArrayArguments
         @TruffleBoundary
         @Specialization
         protected RubyArray getClassVariables(RubyModule module) {
-            final Map<String, Object> allClassVariables = ModuleOperations.getAllClassVariables(module);
-            final int size = allClassVariables.size();
+            final Map<String, Object> classVariables = new HashMap<>();
+
+            ModuleOperations.classVariableLookup(module, m -> {
+                final ClassVariableStorage classVariableStorage = m.fields.getClassVariables();
+
+                for (Object key : m.fields.getClassVariables().getShape().getKeys()) {
+                    classVariables.put(
+                            (String) key,
+                            DynamicObjectLibrary.getUncached().getOrDefault(classVariableStorage, key, null));
+                }
+
+                return null;
+            });
+
+            final int size = classVariables.size();
             final Object[] store = new Object[size];
 
             int i = 0;
-            for (String variable : allClassVariables.keySet()) {
+            for (String variable : classVariables.keySet()) {
                 store[i++] = getSymbol(variable);
             }
             return createArray(store);
@@ -1941,10 +1957,10 @@ protected RubyNode coerceToString(RubyNode name) {
             return NameToJavaStringNode.create(name);
         }
 
-        @TruffleBoundary
         @Specialization
-        protected Object removeClassVariableString(RubyModule module, String name) {
-            SymbolTable.checkClassVariableName(getContext(), name, module, this);
+        protected Object removeClassVariableString(RubyModule module, String name,
+                @Cached CheckClassVariableNameNode checkClassVariableNameNode) {
+            checkClassVariableNameNode.execute(module, name);
             return ModuleOperations.removeClassVariable(module.fields, getContext(), this, name);
         }
 

src/main/java/org/truffleruby/core/module/ModuleOperations.java

@@ -16,6 +16,7 @@
 import java.util.Map.Entry;
 import java.util.function.Function;
 
+import com.oracle.truffle.api.object.DynamicObjectLibrary;
 import org.truffleruby.RubyContext;
 import org.truffleruby.RubyLanguage;
 import org.truffleruby.collections.Memo;
@@ -27,6 +28,7 @@
 import org.truffleruby.language.control.RaiseException;
 import org.truffleruby.language.methods.DeclarationContext;
 import org.truffleruby.language.methods.InternalMethod;
+import org.truffleruby.language.objects.classvariables.ClassVariableStorage;
 import org.truffleruby.language.objects.shared.SharedObjects;
 import org.truffleruby.parser.Identifiers;
 
@@ -614,23 +616,6 @@ private static Assumption[] toArray(ArrayList<Assumption> assumptions) {
         return assumptions.toArray(EMPTY_ASSUMPTION_ARRAY);
     }
 
-    @TruffleBoundary
-    public static Map<String, Object> getAllClassVariables(RubyModule module) {
-        final Map<String, Object> classVariables = new HashMap<>();
-
-        classVariableLookup(module, module1 -> {
-            classVariables.putAll(module1.fields.getClassVariables());
-            return null;
-        });
-
-        return classVariables;
-    }
-
-    @TruffleBoundary
-    public static Object lookupClassVariable(RubyModule module, final String name) {
-        return classVariableLookup(module, module1 -> module1.fields.getClassVariables().get(name));
-    }
-
     @TruffleBoundary
     public static void setClassVariable(RubyLanguage language, RubyContext context, RubyModule module, String name,
             Object value, Node currentNode) {
@@ -643,48 +628,45 @@ public static void setClassVariable(RubyLanguage language, RubyContext context,
         if (!trySetClassVariable(module, name, value)) {
             synchronized (context.getClassVariableDefinitionLock()) {
                 if (!trySetClassVariable(module, name, value)) {
-                    /* This is double-checked locking, but it is safe because when writing to a ConcurrentHashMap "an
-                     * update operation for a given key bears a happens-before relation with any (non-null) retrieval
-                     * for that key reporting the updated value" (JavaDoc) so the value is guaranteed to be fully
-                     * published before it can be found in the map. */
-
-                    moduleFields.getClassVariables().put(name, value);
+                    DynamicObjectLibrary.getUncached().put(moduleFields.getClassVariables(), name, value);
                 }
             }
         }
     }
 
     private static boolean trySetClassVariable(RubyModule topModule, String name, Object value) {
-        final RubyModule found = classVariableLookup(topModule, module -> {
-            final ModuleFields moduleFields = module.fields;
-            if (moduleFields.getClassVariables().replace(name, value) != null) {
-                return module;
-            } else {
-                return null;
-            }
-        });
-        return found != null;
+        return classVariableLookup(
+                topModule,
+                module -> DynamicObjectLibrary.getUncached().putIfPresent(
+                        module.fields.getClassVariables(),
+                        name,
+                        value) ? module : null) != null;
     }
 
     @TruffleBoundary
     public static Object removeClassVariable(ModuleFields moduleFields, RubyContext context, Node currentNode,
             String name) {
         moduleFields.checkFrozen(context, currentNode);
 
-        final Object found = moduleFields.getClassVariables().remove(name);
+        final ClassVariableStorage classVariableStorage = moduleFields.getClassVariables();
+
+        final Object found = DynamicObjectLibrary.getUncached().getOrDefault(classVariableStorage, name, null);
+
         if (found == null) {
             throw new RaiseException(
                     context,
                     context.getCoreExceptions().nameErrorClassVariableNotDefined(
                             name,
                             moduleFields.rubyModule,
                             currentNode));
+        } else {
+            DynamicObjectLibrary.getUncached().removeKey(classVariableStorage, name);
+            return found;
         }
-        return found;
     }
 
     @TruffleBoundary
-    private static <R> R classVariableLookup(RubyModule module, Function<RubyModule, R> action) {
+    public static <R> R classVariableLookup(RubyModule module, Function<RubyModule, R> action) {
         // Look in the current module
         R result = action.apply(module);
         if (result != null) {

src/main/java/org/truffleruby/core/module/RubyModule.java

@@ -23,6 +23,7 @@
 import com.oracle.truffle.api.library.ExportMessage;
 import com.oracle.truffle.api.object.Shape;
 import com.oracle.truffle.api.source.SourceSection;
+import org.truffleruby.language.objects.classvariables.ClassVariableStorage;
 
 @ExportLibrary(InteropLibrary.class)
 public class RubyModule extends RubyDynamicObject implements ObjectGraphNode {
@@ -63,6 +64,10 @@ public String getName() {
         return fields.getName();
     }
 
+    public ClassVariableStorage getClassVariables() {
+        return fields.getClassVariables();
+    }
+
     @TruffleBoundary
     @Override
     public String toString() {

src/main/java/org/truffleruby/language/LexicalScope.java

@@ -12,7 +12,6 @@
 import org.truffleruby.core.module.RubyModule;
 
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
-import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 
 /** Instances of this class represent the Ruby lexical scope for constants, which is only changed by `class Name`,
  * `module Name` and `class << expr`. Other lexical scope features such as refinement and the default definee are
@@ -48,18 +47,6 @@ public void unsafeSetLiveModule(RubyModule liveModule) {
         this.liveModule = liveModule;
     }
 
-    @TruffleBoundary
-    public static RubyModule resolveTargetModuleForClassVariables(LexicalScope lexicalScope) {
-        LexicalScope scope = lexicalScope;
-
-        // MRI logic: ignore lexical scopes (cref) referring to singleton classes
-        while (RubyGuards.isSingletonClass(scope.liveModule)) {
-            scope = scope.parent;
-        }
-
-        return scope.liveModule;
-    }
-
     @Override
     public String toString() {
         return " :: " + liveModule + (parent == null ? "" : parent.toString());