Rework framing for super and __callee__

core/src/main/java/org/jruby/RubyModule.java

@@ -89,7 +89,18 @@
 import org.jruby.internal.runtime.methods.UndefinedMethod;
 import org.jruby.ir.IRClosure;
 import org.jruby.ir.IRMethod;
+import org.jruby.ir.instructions.CallBase;
+import org.jruby.ir.instructions.ClassSuperInstr;
+import org.jruby.ir.instructions.InstanceSuperInstr;
+import org.jruby.ir.instructions.Instr;
+import org.jruby.ir.instructions.ModuleSuperInstr;
+import org.jruby.ir.instructions.SuperInstr;
+import org.jruby.ir.instructions.UnresolvedSuperInstr;
+import org.jruby.ir.interpreter.InterpreterContext;
+import org.jruby.ir.operands.WrappedIRClosure;
+import org.jruby.ir.representations.BasicBlock;
 import org.jruby.ir.targets.indy.Bootstrap;
+import org.jruby.ir.transformations.inlining.SimpleCloneInfo;
 import org.jruby.javasupport.JavaClass;
 import org.jruby.javasupport.binding.MethodGatherer;
 import org.jruby.parser.StaticScope;
@@ -102,6 +113,7 @@
 import org.jruby.runtime.IRBlockBody;
 import org.jruby.runtime.MethodFactory;
 import org.jruby.runtime.MethodIndex;
+import org.jruby.runtime.MixedModeIRBlockBody;
 import org.jruby.runtime.ObjectAllocator;
 import org.jruby.runtime.ThreadContext;
 import org.jruby.runtime.Visibility;
@@ -1616,7 +1628,7 @@ private CacheEntry refinedMethodOriginalMethodEntry(Map<RubyModule, RubyModule>
      * failed to return a result. Cache superclass definitions in this class.
      *
      * MRI: method_entry_get_without_cache
-     * 
+     *
      * @param id The name of the method to search for
      * @param cacheUndef Flag for caching UndefinedMethod. This should normally be true.
      * @return The method, or UndefinedMethod if not found
@@ -2386,19 +2398,25 @@ public IRubyObject defineMethodFromBlock(ThreadContext context, IRubyObject arg0
 
         // If we know it comes from IR we can convert this directly to a method and
         // avoid overhead of invoking it as a block
-        if (block.getBody() instanceof IRBlockBody &&
-                runtime.getInstanceConfig().getCompileMode().shouldJIT()) { // FIXME: Once Interp and Mixed Methods are one class we can fix this to work in interp mode too.
+        if (block.getBody() instanceof IRBlockBody) {
             IRBlockBody body = (IRBlockBody) block.getBody();
-            IRClosure closure = body.getScope();
-
-            // closure may be null from AOT scripts
-            if (closure != null) {
-                // Ask closure to give us a method equivalent.
-                IRMethod method = closure.convertToMethod(name.getBytes());
-                if (method != null) {
-                    newMethod = new DefineMethodMethod(method, visibility, this, context.getFrameBlock());
-                    Helpers.addInstanceMethod(this, name, newMethod, visibility, context, runtime);
-                    return name;
+
+            IRClosure closure = rewriteClosure(name, body.getScope());
+
+            body = new MixedModeIRBlockBody(closure, body.getSignature());
+            block = new Block(body, block.getBinding(), block.type);
+
+            if (runtime.getInstanceConfig().getCompileMode().shouldJIT()) { // FIXME: Once Interp and Mixed Methods are one class we can fix this to work in interp mode too.
+
+                // closure may be null from AOT scripts
+                if (closure != null) {
+                    // Ask closure to give us a method equivalent.
+                    IRMethod method = closure.convertToMethod(name.getBytes());
+                    if (method != null) {
+                        newMethod = new DefineMethodMethod(method, visibility, this, context.getFrameBlock());
+                        Helpers.addInstanceMethod(this, name, newMethod, visibility, context, runtime);
+                        return name;
+                    }
                 }
             }
         }
@@ -2409,6 +2427,77 @@ public IRubyObject defineMethodFromBlock(ThreadContext context, IRubyObject arg0
         return name;
     }
 
+    private IRClosure rewriteClosure(RubySymbol name, IRClosure closure) {
+        // FIXME: This should be somewhere within IRClosure or at least within .ir package space.
+        // If we have a super containing closure but it has been built it means we have already been here.
+        if ((closure.usesSuper() || closure.childUsesSuper()) && closure.getFullInterpreterContext() == null) {
+            // We depend on using startup because ACP has not been called and the method rewriting opto below depends
+            // on us still having AST available which would not be the case by the time we made it to full.
+            closure = closure.cloneStartup();
+
+            InterpreterContext ic = closure.getInterpreterContext();
+            Instr[] list = ic.getInstructions();
+
+            List<Instr> instrs = new ArrayList<>(Arrays.asList(list));
+
+            for (int i = 0; i < instrs.size(); i++) {
+                Instr instr = instrs.get(i);
+
+                if (instr instanceof CallBase && ((CallBase) instr).hasLiteralClosure()) {
+                    CallBase callBase = (CallBase) instr;
+
+                    WrappedIRClosure closureArg = (WrappedIRClosure) callBase.getClosureArg();
+
+                    IRClosure wrapped = closureArg.getClosure();
+                    IRClosure rewritten = rewriteClosure(name, wrapped);
+
+                    closure.removeClosure(wrapped);
+                    closure.addClosure(rewritten);
+
+                    callBase.setClosureArg(new WrappedIRClosure(closureArg.getSelf(), rewritten));
+                }
+
+                if (instr instanceof SuperInstr) {
+                    SuperInstr superInstr = (SuperInstr) instr;
+                    instrs.remove(i);
+
+                    if (this.isSingleton()) {
+                        instrs.add(i, new ClassSuperInstr(
+                                closure,
+                                superInstr.getResult(),
+                                superInstr.getDefiningModule(),
+                                name,
+                                superInstr.getCallArgs(),
+                                superInstr.getClosureArg(),
+                                superInstr.isPotentiallyRefined()));
+                    } else if (this.isModule()) {
+                        instrs.add(i, new ModuleSuperInstr(
+                                closure,
+                                superInstr.getResult(),
+                                name,
+                                superInstr.getReceiver(),
+                                superInstr.getCallArgs(),
+                                superInstr.getClosureArg(),
+                                superInstr.isPotentiallyRefined()));
+                    } else {
+                        instrs.add(i, new InstanceSuperInstr(
+                                closure,
+                                superInstr.getResult(),
+                                superInstr.getDefiningModule(),
+                                name,
+                                superInstr.getCallArgs(),
+                                superInstr.getClosureArg(),
+                                superInstr.isPotentiallyRefined()));
+                    }
+                }
+            }
+
+            ic.setInstructionsRaw(instrs);
+        }
+
+        return closure;
+    }
+
     @JRubyMethod(name = "define_method", reads = VISIBILITY)
     public IRubyObject define_method(ThreadContext context, IRubyObject arg0, IRubyObject arg1, Block block) {
         Visibility visibility = getCurrentVisibilityForDefineMethod(context);
@@ -2470,6 +2559,7 @@ private DynamicMethod createProcMethod(Ruby runtime, String name, Visibility vis
         // various instructions can tell this scope is not an ordinary block but a block representing
         // a method definition.
         block.getBody().getStaticScope().makeArgumentScope();
+        block.getBody().getStaticScope().setModule(this);
 
         return new ProcMethod(this, proc, visibility, name);
     }

core/src/main/java/org/jruby/ir/IRBuilder.java

@@ -4077,19 +4077,37 @@ public Operand buildStrRaw(StrNode strNode) {
     private Operand buildSuperInstr(Operand block, Operand[] args) {
         CallInstr superInstr;
         Variable ret = createTemporaryVariable();
-        if (scope instanceof IRMethod && scope.getLexicalParent() instanceof IRClassBody) {
-            if (((IRMethod) scope).isInstanceMethod) {
-                superInstr = new InstanceSuperInstr(scope, ret, getCurrentModuleVariable(), getName(), args, block, scope.maybeUsingRefinements());
+        IRMethod nearestMethod = scope.getNearestMethod();
+
+        if (nearestMethod != null) {
+            IRScope lexicalParent;
+            RubySymbol superName;
+
+            if (scope instanceof IRMethod) {
+                lexicalParent = scope.getLexicalParent();
+                superName = getName();
             } else {
-                superInstr = new ClassSuperInstr(scope, ret, getCurrentModuleVariable(), getName(), args, block, scope.maybeUsingRefinements());
+                lexicalParent = nearestMethod.getLexicalParent();
+                superName = nearestMethod.getName();
+            }
+
+            if (lexicalParent instanceof IRClassBody) {
+                if (nearestMethod.isInstanceMethod) {
+                    superInstr = new InstanceSuperInstr(scope, ret, getCurrentModuleVariable(), superName, args, block, scope.maybeUsingRefinements());
+                } else {
+                    superInstr = new ClassSuperInstr(scope, ret, getCurrentModuleVariable(), superName, args, block, scope.maybeUsingRefinements());
+                }
+            } else {
+                superInstr = new ModuleSuperInstr(scope, ret, superName, buildSelf(), args, block, scope.maybeUsingRefinements());
             }
         } else {
             // We dont always know the method name we are going to be invoking if the super occurs in a closure.
             // This is because the super can be part of a block that will be used by 'define_method' to define
             // a new method.  In that case, the method called by super will be determined by the 'name' argument
             // to 'define_method'.
-            superInstr = new UnresolvedSuperInstr(scope, ret, buildSelf(), args, block, scope.maybeUsingRefinements());
+            superInstr = new UnresolvedSuperInstr(scope, ret, getCurrentModuleVariable(), buildSelf(), args, block, scope.maybeUsingRefinements());
         }
+
         receiveBreakException(block, superInstr);
         return ret;
     }
@@ -4252,7 +4270,7 @@ private Operand buildZSuperIfNest(final Operand block) {
         Variable zsuperResult = createTemporaryVariable();
         if (superScope instanceof IRMethod && !defineMethod) {
             Operand[] args = adjustVariableDepth(getCallArgs(superScope, superBuilder), depthFromSuper);
-            addInstr(new ZSuperInstr(scope, zsuperResult, buildSelf(), args, block, scope.maybeUsingRefinements()));
+            addInstr(new ZSuperInstr(scope, zsuperResult, getCurrentModuleVariable(), buildSelf(), args, block, scope.maybeUsingRefinements()));
         } else {
             // We will not have a zsuper show up since we won't emit it but we still need to toggle it.
             // define_method optimization will try and create a method from a closure but it should not in this case.

core/src/main/java/org/jruby/ir/IRClosure.java

@@ -323,6 +323,21 @@ protected IRClosure cloneForInlining(CloneInfo ii, IRClosure clone) {
     private static final ByteList CLOSURE_CLONE =
             new ByteList(new byte[] {'_', 'C', 'L', 'O', 'S', 'U', 'R', 'E', '_', 'C', 'L', 'O', 'N', 'E', '_'}, false);
 
+    // FIXME: We should be cloning children scopes
+    // FIXME: IRScope should be keeping booleans
+    /**
+     * This is for cloning when you know you are still within startup interpreter
+     * @return a new copy of this scope
+     */
+    public IRClosure cloneStartup() {
+        assert this.getFullInterpreterContext() == null : "Trying to clone startup and full has already been made";
+        IRClosure clonedClosure = new IRClosure(this, this.getLexicalParent(), getLexicalParent().getNextClosureId(), getByteName());
+
+        clonedClosure.interpreterContext = new InterpreterContext(this);
+
+        return clonedClosure;
+    }
+
     public IRClosure cloneForInlining(CloneInfo ii) {
         IRClosure clonedClosure;
         IRScope lexicalParent = ii.getScope();

core/src/main/java/org/jruby/ir/IRScope.java

@@ -517,6 +517,13 @@ public boolean usesSuper() {
         return usesSuper;
     }
 
+    public boolean childUsesSuper() {
+        for (IRClosure closure : getClosures()) {
+            if (closure.usesSuper() || closure.childUsesSuper()) return true;
+        }
+        return false;
+    }
+
     public boolean usesZSuper() {
         return usesZSuper;
     }

core/src/main/java/org/jruby/ir/IRVisitor.java

@@ -78,6 +78,7 @@ private void error(Object object) {
     public void LoadFrameClosure(LoadFrameClosureInstr loadframeclosureinstr) { error(loadframeclosureinstr); }
     public void LoadBlockImplicitClosure(LoadBlockImplicitClosureInstr loadblockimplicitclosureinstr) { error(loadblockimplicitclosureinstr); }
     public void MatchInstr(MatchInstr matchInstr) { error(matchInstr); }
+    public void ModuleSuperInstr(ModuleSuperInstr modulesuperinstr) { error(modulesuperinstr); }
     public void ModuleVersionGuardInstr(ModuleVersionGuardInstr moduleversionguardinstr) { error(moduleversionguardinstr); }
     public void NonlocalReturnInstr(NonlocalReturnInstr nonlocalreturninstr) { error(nonlocalreturninstr); }
     public void NopInstr(NopInstr nopinstr) { error(nopinstr); }

core/src/main/java/org/jruby/ir/Operation.java

@@ -70,6 +70,7 @@ public enum Operation {
     AS_STRING(OpFlags.f_is_call | OpFlags.f_has_side_effect | OpFlags.f_can_raise_exception),
     CLASS_SUPER(OpFlags.f_has_side_effect | OpFlags.f_is_call | OpFlags.f_can_raise_exception),
     INSTANCE_SUPER(OpFlags.f_has_side_effect | OpFlags.f_is_call | OpFlags.f_can_raise_exception),
+    MODULE_SUPER(OpFlags.f_has_side_effect | OpFlags.f_is_call | OpFlags.f_can_raise_exception),
     UNRESOLVED_SUPER(OpFlags.f_has_side_effect | OpFlags.f_is_call | OpFlags.f_can_raise_exception),
     ZSUPER(OpFlags.f_has_side_effect | OpFlags.f_is_call | OpFlags.f_can_raise_exception),
 

core/src/main/java/org/jruby/ir/instructions/CallBase.java

@@ -30,7 +30,7 @@ public abstract class CallBase extends NOperandInstr implements ClosureAccepting
 
     public transient long callSiteId;
     private final CallType callType;
-    protected final RubySymbol name;
+    protected RubySymbol name;
     protected final transient CallSite callSite;
     protected final transient int argsCount;
     protected final transient boolean hasClosure;
@@ -104,6 +104,10 @@ public String getId() {
         return name.idString();
     }
 
+    public void setName(RubySymbol name) {
+        this.name = name;
+    }
+
     public long getCallSiteId() {
         return callSiteId;
     }
@@ -121,6 +125,10 @@ public Operand getClosureArg() {
         return hasClosure ? operands[argsCount + 1] : null;
     }
 
+    public void setClosureArg(WrappedIRClosure closureArg) {
+        operands[argsCount + 1] = closureArg;
+    }
+
     public Operand getClosureArg(Operand ifUnspecified) {
         return hasClosure ? getClosureArg() : ifUnspecified;
     }

core/src/main/java/org/jruby/ir/instructions/ClassSuperInstr.java

@@ -18,21 +18,21 @@
 
 import java.util.EnumSet;
 
-public class ClassSuperInstr extends CallInstr {
+public class ClassSuperInstr extends SuperInstr {
     private final boolean isLiteralBlock;
 
     // clone constructor
     protected ClassSuperInstr(IRScope scope, Variable result, Operand receiver, RubySymbol name, Operand[] args,
                               Operand closure, boolean potentiallyRefined, CallSite callSite, long callSiteId) {
-        super(scope, Operation.CLASS_SUPER, CallType.SUPER, result, name, receiver, args, closure, potentiallyRefined, callSite, callSiteId);
+        super(scope, Operation.CLASS_SUPER, result, receiver, name, args, closure, potentiallyRefined, callSite, callSiteId);
 
         isLiteralBlock = closure instanceof WrappedIRClosure;
     }
 
     // normal constructor
     public ClassSuperInstr(IRScope scope, Variable result, Operand definingModule, RubySymbol name, Operand[] args, Operand closure,
                            boolean isPotentiallyRefined) {
-        super(scope, Operation.CLASS_SUPER, CallType.SUPER, result, name, definingModule, args, closure, isPotentiallyRefined);
+        super(scope, Operation.CLASS_SUPER, result, definingModule, name, args, closure, isPotentiallyRefined);
 
         isLiteralBlock = closure instanceof WrappedIRClosure;
     }

core/src/main/java/org/jruby/ir/instructions/InstanceSuperInstr.java

@@ -49,13 +49,13 @@
 
 import java.util.EnumSet;
 
-public class InstanceSuperInstr extends CallInstr {
+public class InstanceSuperInstr extends SuperInstr {
     private final boolean isLiteralBlock;
 
     // clone constructor
     protected InstanceSuperInstr(IRScope scope, Variable result, Operand definingModule, RubySymbol name, Operand[] args,
                                  Operand closure, boolean isPotentiallyRefined, CallSite callSite, long callSiteId) {
-        super(scope, Operation.INSTANCE_SUPER, CallType.SUPER, result, name, definingModule, args, closure,
+        super(scope, Operation.INSTANCE_SUPER, result, definingModule, name, args, closure,
                 isPotentiallyRefined, callSite, callSiteId);
 
         isLiteralBlock = closure instanceof WrappedIRClosure;
@@ -64,7 +64,7 @@ protected InstanceSuperInstr(IRScope scope, Variable result, Operand definingMod
     // normal constructor
     public InstanceSuperInstr(IRScope scope, Variable result, Operand definingModule, RubySymbol name, Operand[] args,
                               Operand closure, boolean isPotentiallyRefined) {
-        super(scope, Operation.INSTANCE_SUPER, CallType.SUPER, result, name, definingModule, args, closure, isPotentiallyRefined);
+        super(scope, Operation.INSTANCE_SUPER, result, definingModule, name, args, closure, isPotentiallyRefined);
 
         isLiteralBlock = closure instanceof WrappedIRClosure;
     }