gistfile1.diff

diff --git a/core/src/main/java/org/jruby/ir/IRBuilder.java b/core/src/main/java/org/jruby/ir/IRBuilder.java
index 1f0cb3f..3c6e1a0 100644
--- a/core/src/main/java/org/jruby/ir/IRBuilder.java
+++ b/core/src/main/java/org/jruby/ir/IRBuilder.java
@@ -2214,8 +2214,14 @@ public class IRBuilder {
           L3:
 
      * ****************************************************************/
-    public Operand buildEnsureNode(EnsureNode ensureNode) {
-        Node bodyNode = ensureNode.getBodyNode();
+    public Operand buildEnsureNode(final EnsureNode ensureNode) {
+        return buildEnsureInternal(ensureNode.getBodyNode(), ensureNode.getEnsureNode());
+    }
+
+    public Operand buildEnsureInternal(Node ensureBodyNode, Node ensurerNode) {
+        // Save $!
+        final Variable savedGlobalException = createTemporaryVariable();
+        addInstr(new GetGlobalVariableInstr(savedGlobalException, "$!"));
 
         // ------------ Build the body of the ensure block ------------
         //
@@ -2225,12 +2231,14 @@ public class IRBuilder {
         // Push a new ensure block node onto the stack of ensure bodies being built
         // The body's instructions are stashed and emitted later.
         EnsureBlockInfo ebi = new EnsureBlockInfo(scope,
-            (bodyNode instanceof RescueNode) ? (RescueNode)bodyNode : null,
+            (ensureBodyNode instanceof RescueNode) ? (RescueNode)ensureBodyNode : null,
             getCurrentLoop(),
             activeRescuers.peek());
 
         ensureBodyBuildStack.push(ebi);
-        Operand ensureRetVal = (ensureNode.getEnsureNode() == null) ? manager.getNil() : build(ensureNode.getEnsureNode());
+        Operand ensureRetVal = ensurerNode == null ? manager.getNil() : build(ensurerNode);
+        // Restore $!
+        addInstr(new PutGlobalVarInstr("$!", savedGlobalException));
         ensureBodyBuildStack.pop();
 
         // ------------ Build the protected region ------------
@@ -2242,7 +2250,7 @@ public class IRBuilder {
         activeRescuers.push(ebi.dummyRescueBlockLabel);
 
         // Generate IR for code being protected
-        Operand rv = bodyNode instanceof RescueNode ? buildRescueInternal((RescueNode) bodyNode, ebi) : build(bodyNode);
+        Operand rv = ensureBodyNode instanceof RescueNode ? buildRescueInternal((RescueNode) ensureBodyNode, ebi) : build(ensureBodyNode);
 
         // End of protected region
         addInstr(new ExceptionRegionEndMarkerInstr());
@@ -2250,7 +2258,7 @@ public class IRBuilder {
 
         // Clone the ensure body and jump to the end.
         // Don't bother if the protected body ended in a return.
-        if (rv != U_NIL && !(bodyNode instanceof RescueNode)) {
+        if (rv != U_NIL && !(ensureBodyNode instanceof RescueNode)) {
             ebi.cloneIntoHostScope(this);
             addInstr(new JumpInstr(ebi.end));
         }
@@ -3023,7 +3031,7 @@ public class IRBuilder {
     }
 
     public Operand buildRescue(RescueNode node) {
-        return buildRescueInternal(node, null);
+        return buildEnsureInternal(node, null);
     }
 
     private Operand buildRescueInternal(RescueNode rescueNode, EnsureBlockInfo ensure) {
@@ -3035,7 +3043,7 @@ public class IRBuilder {
         // Save $! in a temp var so it can be restored when the exception gets handled.
         Variable savedGlobalException = createTemporaryVariable();
         addInstr(new GetGlobalVariableInstr(savedGlobalException, "$!"));
-        if (ensure != null) ensure.savedGlobalException = savedGlobalException;
+        ensure.savedGlobalException = savedGlobalException;
 
         addInstr(new LabelInstr(rBeginLabel));
 
@@ -3085,9 +3093,7 @@ public class IRBuilder {
             // No explicit return from the protected body
             // - If we dont have any ensure blocks, simply jump to the end of the rescue block
             // - If we do, execute the ensure code.
-            if (ensure != null) {
-                ensure.cloneIntoHostScope(this);
-            }
+            ensure.cloneIntoHostScope(this);
             addInstr(new JumpInstr(rEndLabel));
         }   //else {
             // If the body had an explicit return, the return instruction IR build takes care of setting
@@ -3108,9 +3114,6 @@ public class IRBuilder {
         // Build the actual rescue block(s)
         buildRescueBodyInternal(rescueNode.getRescueNode(), rv, exc, rEndLabel);
 
-        // End label -- only if there is no ensure block!  With an ensure block, you end at ensureEndLabel.
-        if (ensure == null) addInstr(new LabelInstr(rEndLabel));
-
         activeRescueBlockStack.pop();
         return rv;
     }