Fix #3104: Ensure that push/pops don't break returns

* return instructions may sometime reference values that rely on the current scope/frame. However, AddCallProtocalInstructionsPass introduces pops of those before returns. To ensure that returns are no longer broken, that pass should copy the value into a temporary variable before popping the scope/frame. * Added a new regression spec that primes the JIT by running the function in a loop 100 times. This ensures that the spec doesn't artificially pass because of relying on the startup interpreter. Verified that the spec fails without the fixes in this patch.
jruby · Jul 7, 2015 · be5e7f1 · be5e7f1
1 parent d0659d1
commit be5e7f1
Showing 3 changed files with 51 additions and 1 deletion.
diff --git a/core/src/main/java/org/jruby/ir/instructions/ReturnBase.java b/core/src/main/java/org/jruby/ir/instructions/ReturnBase.java
@@ -13,4 +13,8 @@ public ReturnBase(Operation op, Operand returnValue) {
     public Operand getReturnValue() {
         return operands[0];
     }
+
+    public void updateReturnValue(Operand val) {
+        operands[0] = val;
+    }
 }
diff --git a/core/src/main/java/org/jruby/ir/passes/AddCallProtocolInstructions.java b/core/src/main/java/org/jruby/ir/passes/AddCallProtocolInstructions.java
@@ -3,7 +3,10 @@
 import org.jruby.ir.*;
 import org.jruby.ir.dataflow.analyses.StoreLocalVarPlacementProblem;
 import org.jruby.ir.instructions.*;
+import org.jruby.ir.operands.ImmutableLiteral;
 import org.jruby.ir.operands.Label;
+import org.jruby.ir.operands.Operand;
+import org.jruby.ir.operands.TemporaryVariable;
 import org.jruby.ir.operands.Variable;
 import org.jruby.ir.representations.BasicBlock;
 import org.jruby.ir.representations.CFG;
@@ -20,6 +23,24 @@ private boolean explicitCallProtocolSupported(IRScope scope) {
         return scope instanceof IRMethod || (scope instanceof IRModuleBody && !(scope instanceof IRMetaClassBody));
     }
 
+    /*
+     * Since the return is now going to be preceded by a pops of bindings/frames,
+     * the return value should continue to be valid after those pops.
+     * If not, introduce a copy into a tmp-var before the pops and use the tmp-var
+     * to return the right value.
+     */
+    private void fixReturn(IRScope scope, ReturnBase i, ListIterator<Instr> instrs) {
+        Operand retVal = i.getReturnValue();
+        if (!(retVal instanceof ImmutableLiteral || retVal instanceof TemporaryVariable)) {
+            TemporaryVariable tmp = scope.createTemporaryVariable();
+            CopyInstr copy = new CopyInstr(tmp, retVal);
+            i.updateReturnValue(tmp);
+            instrs.previous();
+            instrs.add(copy);
+            instrs.next();
+        }
+    }
+
     @Override
     public Object execute(IRScope scope, Object... data) {
         // IRScriptBody do not get explicit call protocol instructions right now.
@@ -78,6 +99,9 @@ public Object execute(IRScope scope, Object... data) {
                         // throw an exception and we don't multiple pops (once before the
                         // return/break, and once when the exception is caught).
                         if (!bb.isExitBB() && i instanceof ReturnBase) {
+                            if (requireBinding || requireFrame) {
+                                fixReturn(scope, (ReturnBase)i, instrs);
+                            }
                             // Add before the break/return
                             instrs.previous();
                             if (requireBinding) instrs.add(new PopBindingInstr());
@@ -88,7 +112,12 @@ public Object execute(IRScope scope, Object... data) {
 
                     if (bb.isExitBB() && !bb.isEmpty()) {
                         // Last instr could be a return -- so, move iterator one position back
-                        if (i != null && i instanceof ReturnBase) instrs.previous();
+                        if (i != null && i instanceof ReturnBase) {
+                            if (requireBinding || requireFrame) {
+                                fixReturn(scope, (ReturnBase)i, instrs);
+                            }
+                            instrs.previous();
+                        }
                         if (requireBinding) instrs.add(new PopBindingInstr());
                         if (requireFrame) instrs.add(new PopFrameInstr());
                     }

diff --git a/spec/regression/GH-3104_fixup_returns.rb b/spec/regression/GH-3104_fixup_returns.rb
@@ -0,0 +1,17 @@
+def foo(string)
+  if string =~ /Prefix (\w)/
+    $1
+  end
+end
+
+# Make sure foo JITs
+(1..100).each do
+  foo('Prefix A')
+end
+
+# https://github.com/jruby/jruby/issues/3104
+describe 'pop scope/frames in AddCallProtocolInstructionsPass' do
+  it 'should not break returns' do
+    expect(foo("Prefix A")).to eq 'A'
+  end
+end