jruby · Jul 2, 2015 · Jul 2, 2015 · Jul 2, 2015
diff --git a/spec/truffle/tags/core/module/class_variable_defined_tags.txt b/spec/truffle/tags/core/module/class_variable_defined_tags.txt
diff --git a/spec/truffle/tags/core/module/class_variable_get_tags.txt b/spec/truffle/tags/core/module/class_variable_get_tags.txt
diff --git a/spec/truffle/tags/core/module/class_variables_tags.txt b/spec/truffle/tags/core/module/class_variables_tags.txt
diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/core/ModuleNodes.java b/truffle/src/main/java/org/jruby/truffle/nodes/core/ModuleNodes.java
@@ -770,7 +770,8 @@ public RubyNode coerceToString(RubyNode name) {
         public Object setClassVariable(RubyModule module, String name, Object value) {
             RubyContext.checkClassVariableName(getContext(), name, this);
 
-            module.getClassVariables().put(name, value);
+            ModuleOperations.setClassVariable(module, name, value, this);
+
             return value;
         }
 

diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/objects/WriteClassVariableNode.java b/truffle/src/main/java/org/jruby/truffle/nodes/objects/WriteClassVariableNode.java
@@ -12,10 +12,12 @@
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.source.SourceSection;
+
 import org.jruby.truffle.nodes.RubyNode;
 import org.jruby.truffle.runtime.LexicalScope;
 import org.jruby.truffle.runtime.ModuleOperations;
 import org.jruby.truffle.runtime.RubyContext;
+import org.jruby.truffle.runtime.core.RubyClass;
 import org.jruby.truffle.runtime.core.RubyModule;
 
 public class WriteClassVariableNode extends RubyNode {
@@ -35,11 +37,16 @@ public WriteClassVariableNode(RubyContext context, SourceSection sourceSection,
     public Object execute(VirtualFrame frame) {
         CompilerDirectives.transferToInterpreter();
 
-        final RubyModule moduleObject = lexicalScope.getLiveModule();
-
         final Object rhsValue = rhs.execute(frame);
 
-        ModuleOperations.setClassVariable(moduleObject, name, rhsValue, this);
+        // MRI logic: ignore lexical scopes (cref) referring to singleton classes
+        LexicalScope scope = this.lexicalScope;
+        while ((scope.getLiveModule() instanceof RubyClass) && ((RubyClass) scope.getLiveModule()).isSingleton()) {
+            scope = scope.getParent();
+        }
+        final RubyModule module = scope.getLiveModule();
+
+        ModuleOperations.setClassVariable(module, name, rhsValue, this);
 
         return rhsValue;
     }

diff --git a/truffle/src/main/java/org/jruby/truffle/runtime/ModuleOperations.java b/truffle/src/main/java/org/jruby/truffle/runtime/ModuleOperations.java
@@ -13,10 +13,12 @@
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.nodes.Node;
+
 import org.jruby.truffle.runtime.control.RaiseException;
 import org.jruby.truffle.runtime.core.RubyClass;
 import org.jruby.truffle.runtime.core.RubyModule;
 import org.jruby.truffle.runtime.methods.InternalMethod;
+import org.jruby.truffle.runtime.util.Function;
 import org.jruby.util.IdUtil;
 
 import java.util.HashMap;
@@ -303,29 +305,54 @@ public static Map<String, Object> getAllClassVariables(RubyModule module) {
 
         final Map<String, Object> classVariables = new HashMap<>();
 
-        // Look in the current module
-        classVariables.putAll(module.getClassVariables());
-
-        // Look in ancestors
-        for (RubyModule ancestor : module.parentAncestors()) {
-            for (Map.Entry<String, Object> classVariable : ancestor.getClassVariables().entrySet()) {
-                if (!classVariables.containsKey(classVariable.getKey())) {
-                    classVariables.put(classVariable.getKey(), classVariable.getValue());
-                }
+        classVariableLookup(module, new Function<RubyModule, Object>() {
+            @Override
+            public Object apply(RubyModule module) {
+                classVariables.putAll(module.getClassVariables());
+                return null;
             }
-        }
+        });
 
         return classVariables;
     }
 
     @TruffleBoundary
-    public static Object lookupClassVariable(RubyModule module, String name) {
+    public static Object lookupClassVariable(RubyModule module, final String name) {
+        return classVariableLookup(module, new Function<RubyModule, Object>() {
+            @Override
+            public Object apply(RubyModule module) {
+                return module.getClassVariables().get(name);
+            }
+        });
+    }
+
+    @TruffleBoundary
+    public static void setClassVariable(RubyModule module, final String name, final Object value, final Node currentNode) {
+        Boolean found = classVariableLookup(module, new Function<RubyModule, Boolean>() {
+            @Override
+            public Boolean apply(RubyModule module) {
+                if (module.getClassVariables().containsKey(name)) {
+                    module.setClassVariable(currentNode, name, value);
+                    return true;
+                } else {
+                    return null;
+                }
+            }
+        });
+
+        if (found == null) {
+            // Not existing class variable - set in the current module
+            module.setClassVariable(currentNode, name, value);
+        }
+    }
+
+    private static <R> R classVariableLookup(RubyModule module, Function<RubyModule, R> action) {
         CompilerAsserts.neverPartOfCompilation();
 
         // Look in the current module
-        Object value = module.getClassVariables().get(name);
-        if (value != null) {
-            return value;
+        R result = action.apply(module);
+        if (result != null) {
+            return result;
         }
 
         // If singleton class, check attached module.
@@ -334,50 +361,24 @@ public static Object lookupClassVariable(RubyModule module, String name) {
             if (klass.isSingleton() && klass.getAttached() != null) {
                 module = klass.getAttached();
 
-                value = module.getClassVariables().get(name);
-                if (value != null) {
-                    return value;
+                result = action.apply(module);
+                if (result != null) {
+                    return result;
                 }
             }
         }
 
         // Look in ancestors
         for (RubyModule ancestor : module.parentAncestors()) {
-            value = ancestor.getClassVariables().get(name);
-            if (value != null) {
-                return value;
+            result = action.apply(ancestor);
+            if (result != null) {
+                return result;
             }
         }
 
-        // Nothing found
         return null;
     }
 
-    @TruffleBoundary
-    public static void setClassVariable(RubyModule module, String name, Object value, Node currentNode) {
-        CompilerAsserts.neverPartOfCompilation();
-
-        // Look in the current module
-
-        if (module.getClassVariables().containsKey(name)) {
-            module.setClassVariable(currentNode, name, value);
-            return;
-        }
-
-        // Look in ancestors
-
-        for (RubyModule ancestor : module.parentAncestors()) {
-            if (ancestor.getClassVariables().containsKey(name)) {
-                ancestor.setClassVariable(currentNode, name, value);
-                return;
-            }
-        }
-
-        // Not existing class variable - set in the current module
-
-        module.setClassVariable(currentNode, name, value);
-    }
-
     public static boolean isMethodPrivateFromName(String name) {
         CompilerAsserts.neverPartOfCompilation();
 

diff --git a/truffle/src/main/java/org/jruby/truffle/runtime/util/Function.java b/truffle/src/main/java/org/jruby/truffle/runtime/util/Function.java
@@ -0,0 +1,16 @@
+/*
+ * Copyright (c) 2015 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 1.0
+ * GNU General Public License version 2
+ * GNU Lesser General Public License version 2.1
+ */
+package org.jruby.truffle.runtime.util;
+
+public interface Function<T, R> {
+
+    R apply(T value);
+
+}