jruby · Mar 9, 2015 · Mar 9, 2015
diff --git a/core/src/main/java/org/jruby/RubyBasicObject.java b/core/src/main/java/org/jruby/RubyBasicObject.java
@@ -163,6 +163,10 @@ public class RubyBasicObject implements Cloneable, IRubyObject, Serializable, Co
     public static final int USER6_F = (1<<(FL_USHIFT+6));
     public static final int USER7_F = (1<<(FL_USHIFT+7));
     public static final int USER8_F = (1<<(FL_USHIFT+8));
+    public static final int USER9_F = (1<<(FL_USHIFT+9));
+    public static final int USERA_F = (1<<(FL_USHIFT+10));
+    public static final int REFINED_MODULE_F = USER9_F;
+    public static final int IS_OVERLAID_F = USERA_F;
 
     public static final int COMPARE_BY_IDENTITY_F = USER8_F;
 

diff --git a/core/src/main/java/org/jruby/RubyModule.java b/core/src/main/java/org/jruby/RubyModule.java
@@ -62,6 +62,7 @@
 import org.jruby.internal.runtime.methods.SynchronizedDynamicMethod;
 import org.jruby.internal.runtime.methods.UndefinedMethod;
 import org.jruby.internal.runtime.methods.WrapperMethod;
+import org.jruby.ir.IRMethod;
 import org.jruby.parser.StaticScope;
 import org.jruby.runtime.Arity;
 import org.jruby.runtime.Block;
@@ -545,6 +546,170 @@ private String calculateAnonymousName() {
         return anonymousName;
     }
 
+
+    @JRubyMethod(name = "refine", required = 1)
+    public IRubyObject refine(ThreadContext context, IRubyObject classArg, Block block) {
+        if (!block.isGiven()) throw context.runtime.newArgumentError("no block given");
+        if (block.isEscaped()) throw context.runtime.newArgumentError("can't pass a Proc as a block to Module#refine");
+        if (!(classArg instanceof RubyClass)) throw context.runtime.newTypeError(classArg, context.runtime.getClassClass());
+        if (refinements == null) refinements = new HashMap<>();
+        if (activatedRefinements == null) activatedRefinements = new HashMap<>();
+
+        RubyClass classWeAreRefining = (RubyClass) classArg;
+        RubyModule refinement = refinements.get(classWeAreRefining);
+        if (refinement == null) {
+            refinement = createNewRefinedModule(context, classWeAreRefining);
+
+            // Add it to the activated chain of other refinements already added to this class we are refining
+            addActivatedRefinement(context, classWeAreRefining, refinement);
+        }
+
+        // Executes the block supplied with the defined method definitions using the refinment as it's module.
+        yieldRefineBlock(context, refinement, block);
+
+        return refinement;
+    }
+
+    private RubyModule createNewRefinedModule(ThreadContext context, RubyClass classWeAreRefining) {
+        RubyModule newRefinement = new RubyModule(context.runtime, classWeAreRefining);
+        newRefinement.setFlag(REFINED_MODULE_F, true);
+        newRefinement.refinedClass = classWeAreRefining;
+        newRefinement.definedAt = this;
+        refinements.put(classWeAreRefining, newRefinement);
+
+        return newRefinement;
+    }
+
+    private void yieldRefineBlock(ThreadContext context, RubyModule refinement, Block block) {
+        block.setEvalType(EvalType.MODULE_EVAL);
+        block.getBinding().setSelf(refinement);
+        block.yieldSpecific(context);
+    }
+
+    // This has three cases:
+    // 1. class being refined has never had any refines happen to it yet: return itself
+    // 2. class has been refined: return already existing refinementwrapper (chain of modules to call against)
+    // 3. refinement is already in the refinementwrapper so we do not need to add it to the wrapper again: return null
+    private RubyClass getAlreadyActivatedRefinementWrapper(RubyClass classWeAreRefining, RubyModule refinement) {
+        // We have already encountered at least one refine on this class.  Return that wrapper.
+        RubyClass moduleWrapperForRefinment = activatedRefinements.get(classWeAreRefining);
+        if (moduleWrapperForRefinment == null) return classWeAreRefining;
+
+        for (RubyModule c = moduleWrapperForRefinment; c != null && c.isIncluded(); c = c.getSuperClass()) {
+            if (c.getNonIncludedClass() == refinement) return null;
+        }
+
+        return moduleWrapperForRefinment;
+    }
+
+    /*
+     * We will find whether we have already refined once and get that set of includedmodules or we will start to create
+     * one.  The new refinement will be added as a new included module on the front.  It will also add all superclasses
+     * of the refinement into this call chain.
+     */
+    // MRI: add_activated_refinement
+    private void addActivatedRefinement(ThreadContext context, RubyClass classWeAreRefining, RubyModule refinement) {
+        RubyClass superClass = getAlreadyActivatedRefinementWrapper(classWeAreRefining, refinement);
+        if (superClass == null) return; // already been refined and added to refinementwrapper
+
+        refinement.setFlag(IS_OVERLAID_F, true);
+        IncludedModuleWrapper iclass = new IncludedModuleWrapper(context.runtime, superClass, refinement);
+        RubyClass c = iclass;
+        c.refinedClass = classWeAreRefining;
+        for (refinement = refinement.getSuperClass(); refinement != null; refinement = refinement.getSuperClass()) {
+            refinement.setFlag(IS_OVERLAID_F, true);
+            RubyClass superClazz = c.getSuperClass();
+            c.setModuleSuperClass(new IncludedModuleWrapper(context.runtime, c.getSuperClass(), refinement));
+            c.refinedClass = classWeAreRefining;
+            c = superClazz;
+        }
+        activatedRefinements.put(classWeAreRefining, iclass);
+    }
+
+    @JRubyMethod(name = "using", required = 1, frame = true)
+    public IRubyObject using(ThreadContext context, IRubyObject refinedModule) {
+        if (context.getFrameSelf() != this) throw context.runtime.newRuntimeError("Module#using is not called on self");
+        // FIXME: This is a lame test and I am unsure it works with JIT'd bodies...
+        if (context.getCurrentScope().getStaticScope().getIRScope() instanceof IRMethod) {
+            throw context.runtime.newRuntimeError("Module#using is not permitted in methods");
+        }
+
+        // I pass the cref even though I don't need to so that the concept is simpler to read
+        usingModule(context, this, refinedModule);
+
+        return this;
+    }
+
+    // mri: rb_using_module
+    public void usingModule(ThreadContext context, RubyModule cref, IRubyObject refinedModule) {
+        if (!(refinedModule instanceof RubyModule))throw context.runtime.newTypeError(refinedModule, context.runtime.getModule());
+
+        usingModuleRecursive(cref, (RubyModule) refinedModule);
+    }
+
+    // mri: using_module_recursive
+    private void usingModuleRecursive(RubyModule cref, RubyModule refinedModule) {
+        RubyClass superClass = refinedModule.getSuperClass();
+
+        // For each superClass of the refined module also use their refinements for the given cref
+        if (superClass != null) usingModuleRecursive(cref, superClass);
+
+        //RubyModule realRefinedModule = refinedModule instanceof IncludedModule ?
+        //                ((IncludedModule) refinedModule).getRealClass() : refinedModule;
+
+        Map<RubyClass, RubyModule> refinements = refinedModule.refinements;
+        if (refinements == null) return; // No refinements registered for this module
+
+        for (Map.Entry<RubyClass, RubyModule> entry: refinements.entrySet()) {
+            usingRefinement(cref, entry.getKey(), entry.getValue());
+        }
+    }
+
+    // This is nearly identical to getAlreadyActivatedRefinementWrapper but thw maps they work against are different.
+    // This has three cases:
+    // 1. class being refined has never had any refines happen to it yet: return itself
+    // 2. class has been refined: return already existing refinementwrapper (chain of modules to call against)
+    // 3. refinement is already in the refinementwrapper so we do not need to add it to the wrapper again: return null
+    private RubyModule getAlreadyRefinementWrapper(RubyModule cref, RubyClass classWeAreRefining, RubyModule refinement) {
+        // We have already encountered at least one refine on this class.  Return that wrapper.
+        RubyModule moduleWrapperForRefinment = cref.refinements.get(classWeAreRefining);
+        if (moduleWrapperForRefinment == null) return classWeAreRefining;
+
+        for (RubyModule c = moduleWrapperForRefinment; c != null && c.isIncluded(); c = c.getSuperClass()) {
+            if (c.getNonIncludedClass() == refinement) return null;
+        }
+
+        return moduleWrapperForRefinment;
+    }
+
+    /*
+     * Within the context of this cref any references to the class we are refining will try and find
+     * that definition from the refinement instead.  At one point I was confused how this would not
+     * conflict if the same module was used in two places but the cref must be a lexically containing
+     * module so it cannot live in two files.
+     */
+    private void usingRefinement(RubyModule cref, RubyClass classWeAreRefining, RubyModule refinement) {
+        // Our storage cubby in cref for all known refinements
+        if (cref.refinements == null) cref.refinements = new HashMap<>();
+
+        RubyModule superClass = getAlreadyRefinementWrapper(cref, classWeAreRefining, refinement);
+        if (superClass == null) return; // already been refined and added to refinementwrapper
+
+        refinement.setFlag(IS_OVERLAID_F, true);
+        RubyModule lookup = new IncludedModuleWrapper(getRuntime(), (RubyClass) superClass, refinement);
+        RubyModule iclass = lookup;
+        lookup.refinedClass = classWeAreRefining;
+
+        for (refinement = refinement.getSuperClass(); refinement != null && refinement != classWeAreRefining; refinement = refinement.getSuperClass()) {
+            refinement.setFlag(IS_OVERLAID_F, true);
+            RubyClass newInclude = new IncludedModuleWrapper(getRuntime(), lookup.getSuperClass(), refinement);
+            lookup.setSuperClass(newInclude);
+            lookup = newInclude;
+            lookup.refinedClass = classWeAreRefining;
+        }
+        cref.refinements.put(classWeAreRefining, iclass);
+    }
+
     /**
      * Create a wrapper to use for including the specified module into this one.
      * 
@@ -4155,7 +4320,15 @@ public boolean isMethodBuiltin(String methodName) {
 
         return method != null && method.isBuiltin();
     }
-
+
+    public Map<RubyClass, RubyModule> getRefinements() {
+        return refinements;
+    }
+
+    public void setRefinements(Map<RubyClass, RubyModule> refinements) {
+        this.refinements = refinements;
+    }
+
     private volatile Map<String, Autoload> autoloads = Collections.EMPTY_MAP;
     protected volatile Map<String, DynamicMethod> methods = Collections.EMPTY_MAP;
     protected Map<String, CacheEntry> cachedMethods = Collections.EMPTY_MAP;
@@ -4199,6 +4372,18 @@ public boolean isMethodBuiltin(String methodName) {
 
     private volatile Map<String, IRubyObject> classVariables = Collections.EMPTY_MAP;
 
+    /** Refinements added to this module are stored here **/
+    private volatile Map<RubyClass, RubyModule> refinements = null;
+
+    /** A list of refinement hosts for this refinement */
+    private volatile Map<RubyClass, IncludedModuleWrapper> activatedRefinements = null;
+
+    /** The class this refinement refines */
+    volatile RubyClass refinedClass = null;
+
+    /** The moduel where this refinement was defined */
+    private volatile RubyModule definedAt = null;
+
     private static final AtomicReferenceFieldUpdater CLASSVARS_UPDATER;
 
     static {

diff --git a/core/src/main/java/org/jruby/ir/IRBuilder.java b/core/src/main/java/org/jruby/ir/IRBuilder.java
@@ -85,6 +85,8 @@ public class IRBuilder {
     static final Operand[] NO_ARGS = new Operand[]{};
     static final UnexecutableNil U_NIL = UnexecutableNil.U_NIL;
 
+    public static final String USING_METHOD = "using";
+
     public static Node buildAST(boolean isCommandLineScript, String arg) {
         Ruby ruby = Ruby.getGlobalRuntime();
 
@@ -1015,7 +1017,7 @@ public Operand buildCall(CallNode callNode) {
         Operand[] args       = setupCallArgs(callArgsNode);
         Operand   block      = setupCallClosure(callNode.getIterNode());
         Variable  callResult = createTemporaryVariable();
-        CallInstr callInstr  = CallInstr.create(callResult, callNode.getName(), receiver, args, block);
+        CallInstr callInstr  = CallInstr.create(scope, callResult, callNode.getName(), receiver, args, block);
 
         // This is to support the ugly Proc.new with no block, which must see caller's frame
         if ( callNode.getName().equals("new") &&
@@ -2283,7 +2285,10 @@ public Operand buildFCall(FCallNode fcallNode) {
         Operand[] args         = setupCallArgs(callArgsNode);
         Operand   block        = setupCallClosure(fcallNode.getIterNode());
         Variable  callResult   = createTemporaryVariable();
-        CallInstr callInstr    = CallInstr.create(CallType.FUNCTIONAL, callResult, fcallNode.getName(), buildSelf(), args, block);
+
+        determineIfMaybeUsingMethod(fcallNode.getName(), args);
+
+        CallInstr callInstr    = CallInstr.create(scope, CallType.FUNCTIONAL, callResult, fcallNode.getName(), buildSelf(), args, block);
         receiveBreakException(block, callInstr);
         return callResult;
     }
@@ -2301,6 +2306,16 @@ private Operand setupCallClosure(Node node) {
         }
     }
 
+    // FIXME: This needs to be called on super/zsuper too
+    private void determineIfMaybeUsingMethod(String methodName, Operand[] args) {
+        IRScope outerScope = scope.getNearestTopLocalVariableScope();
+
+        // 'using single_mod_arg' possible nearly everywhere but method scopes.
+        if (USING_METHOD.equals(methodName) && !(outerScope instanceof IRMethod) && args.length == 1) {
+            scope.setIsMaybeUsingRefinements();
+        }
+    }
+
     public Operand buildFixnum(FixnumNode node) {
         return new Fixnum(node.getValue());
     }
@@ -2343,7 +2358,7 @@ public Operand buildFlip(FlipNode flipNode) {
         if (nearestNonClosureBuilder == null) {
             Variable excType = createTemporaryVariable();
             addInstr(new InheritanceSearchConstInstr(excType, new ObjectClass(), "NotImplementedError", false));
-            Variable exc = addResultInstr(CallInstr.create(createTemporaryVariable(), "new", excType, new Operand[] {new FrozenString("Flip support currently broken")}, null));
+            Variable exc = addResultInstr(CallInstr.create(scope, createTemporaryVariable(), "new", excType, new Operand[] {new FrozenString("Flip support currently broken")}, null));
             addInstr(new ThrowExceptionInstr(exc));
             return buildNil();
         }
@@ -2751,7 +2766,7 @@ public Operand buildOpAsgn(OpAsgnNode opAsgnNode) {
 
         // get attr
         Operand  v1 = build(opAsgnNode.getReceiverNode());
-        addInstr(CallInstr.create(readerValue, opAsgnNode.getVariableName(), v1, NO_ARGS, null));
+        addInstr(CallInstr.create(scope, readerValue, opAsgnNode.getVariableName(), v1, NO_ARGS, null));
 
         // Ex: e.val ||= n
         //     e.val &&= n
@@ -2762,7 +2777,7 @@ public Operand buildOpAsgn(OpAsgnNode opAsgnNode) {
 
             // compute value and set it
             Operand  v2 = build(opAsgnNode.getValueNode());
-            addInstr(CallInstr.create(writerValue, opAsgnNode.getVariableNameAsgn(), v1, new Operand[] {v2}, null));
+            addInstr(CallInstr.create(scope, writerValue, opAsgnNode.getVariableNameAsgn(), v1, new Operand[] {v2}, null));
             // It is readerValue = v2.
             // readerValue = writerValue is incorrect because the assignment method
             // might return something else other than the value being set!
@@ -2776,10 +2791,10 @@ public Operand buildOpAsgn(OpAsgnNode opAsgnNode) {
             // call operator
             Operand  v2 = build(opAsgnNode.getValueNode());
             Variable setValue = createTemporaryVariable();
-            addInstr(CallInstr.create(setValue, opAsgnNode.getOperatorName(), readerValue, new Operand[]{v2}, null));
+            addInstr(CallInstr.create(scope, setValue, opAsgnNode.getOperatorName(), readerValue, new Operand[]{v2}, null));
 
             // set attr
-            addInstr(CallInstr.create(writerValue, opAsgnNode.getVariableNameAsgn(), v1, new Operand[] {setValue}, null));
+            addInstr(CallInstr.create(scope, writerValue, opAsgnNode.getVariableNameAsgn(), v1, new Operand[] {setValue}, null));
             // Returning writerValue is incorrect becuase the assignment method
             // might return something else other than the value being set!
             return setValue;
@@ -2856,12 +2871,12 @@ private Operand buildOpElementAsgnWith(OpElementAsgnNode opElementAsgnNode, Bool
         Label endLabel = getNewLabel();
         Variable elt = createTemporaryVariable();
         Operand[] argList = setupCallArgs(opElementAsgnNode.getArgsNode());
-        addInstr(CallInstr.create(elt, "[]", array, argList, null));
+        addInstr(CallInstr.create(scope, elt, "[]", array, argList, null));
         addInstr(BEQInstr.create(elt, truthy, endLabel));
         Operand value = build(opElementAsgnNode.getValueNode());
 
         argList = addArg(argList, value);
-        addInstr(CallInstr.create(elt, "[]=", array, argList, null));
+        addInstr(CallInstr.create(scope, elt, "[]=", array, argList, null));
         addInstr(new CopyInstr(elt, value));
 
         addInstr(new LabelInstr(endLabel));
@@ -2873,15 +2888,15 @@ public Operand buildOpElementAsgnWithMethod(OpElementAsgnNode opElementAsgnNode)
         Operand array = buildWithOrder(opElementAsgnNode.getReceiverNode(), opElementAsgnNode.containsVariableAssignment());
         Operand[] argList = setupCallArgs(opElementAsgnNode.getArgsNode());
         Variable elt = createTemporaryVariable();
-        addInstr(CallInstr.create(elt, "[]", array, argList, null)); // elt = a[args]
+        addInstr(CallInstr.create(scope, elt, "[]", array, argList, null)); // elt = a[args]
         Operand value = build(opElementAsgnNode.getValueNode());                                       // Load 'value'
         String  operation = opElementAsgnNode.getOperatorName();
-        addInstr(CallInstr.create(elt, operation, elt, new Operand[] { value }, null)); // elt = elt.OPERATION(value)
+        addInstr(CallInstr.create(scope, elt, operation, elt, new Operand[] { value }, null)); // elt = elt.OPERATION(value)
         // SSS: do not load the call result into 'elt' to eliminate the RAW dependency on the call
         // We already know what the result is going be .. we are just storing it back into the array
         Variable tmp = createTemporaryVariable();
         argList = addArg(argList, elt);
-        addInstr(CallInstr.create(tmp, "[]=", array, argList, null));   // a[args] = elt
+        addInstr(CallInstr.create(scope, tmp, "[]=", array, argList, null));   // a[args] = elt
         return elt;
     }
 
@@ -3374,7 +3389,7 @@ public Operand buildVAlias(VAliasNode valiasNode) {
     }
 
     public Operand buildVCall(VCallNode node) {
-        return addResultInstr(CallInstr.create(CallType.VARIABLE, createTemporaryVariable(),
+        return addResultInstr(CallInstr.create(scope, CallType.VARIABLE, createTemporaryVariable(),
                 node.getName(), buildSelf(), NO_ARGS, null));
     }
 

diff --git a/core/src/main/java/org/jruby/ir/IRFlags.java b/core/src/main/java/org/jruby/ir/IRFlags.java
@@ -44,6 +44,7 @@ public enum IRFlags {
     HAS_EXPLICIT_CALL_PROTOCOL,   // contains call protocol instrs => we don't need to manage bindings frame implicitly
     HAS_LOOPS,                    // has a loop
     HAS_NONLOCAL_RETURNS,         // has a non-local return
+    MAYBE_USING_REFINEMENTS,      // a call to 'using' discovered...is it "the" using...maybe?
     RECEIVES_CLOSURE_ARG,         // This scope (or parent receives a closure
     RECEIVES_KEYWORD_ARGS,        // receives keyword args
     REQUIRES_DYNSCOPE,            // does this scope require a dynamic scope?