Fix reenter after refactoring

Author: msimacek

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java

@@ -2172,6 +2172,7 @@ private Object bytecodeLoop(VirtualFrame virtualFrame, Frame localFrame, Bytecod
                     // Need to handle instrumentation frame unwind
                     Object result = notifyException(virtualFrame, instrumentation, returnCalled, bci, e);
                     if (result == ProbeNode.UNWIND_ACTION_REENTER) {
+                        copyArgs(virtualFrame.getArguments(), virtualFrame);
                         bci = 0;
                         stackTop = getInitialStackTop();
                         oparg = 0;
@@ -2269,21 +2270,29 @@ private int handleException(VirtualFrame virtualFrame, Frame localFrame, boolean
 
     @InliningCutoff
     private Object notifyException(VirtualFrame virtualFrame, InstrumentationSupport instrumentation, boolean returnCalled, int bci, Throwable e) {
-        instrumentation.notifyException(virtualFrame, bciToLine(bci), e);
+        try {
+            instrumentation.notifyException(virtualFrame, bciToLine(bci), e);
+        } catch (Throwable t) {
+            e = t;
+        }
         if (instrumentationRoot instanceof WrapperNode) {
             WrapperNode wrapper = (WrapperNode) instrumentationRoot;
             Object result = wrapper.getProbeNode().onReturnExceptionalOrUnwind(virtualFrame, e, returnCalled);
-            if (result != null) {
-                if (co.isGeneratorOrCoroutine()) {
-                    CompilerDirectives.transferToInterpreterAndInvalidate();
-                    throw new IllegalStateException(result == ProbeNode.UNWIND_ACTION_REENTER ? "Frame restarting is not possible in generators" : "Cannot replace return value of generators");
-                }
-            }
+            checkOnReturnExceptionalOrUnwindResult(result);
             return result;
         }
         return null;
     }
 
+    private void checkOnReturnExceptionalOrUnwindResult(Object result) {
+        if (result != null) {
+            if (co.isGeneratorOrCoroutine()) {
+                CompilerDirectives.transferToInterpreterAndInvalidate();
+                throw new IllegalStateException(result == ProbeNode.UNWIND_ACTION_REENTER ? "Frame restarting is not possible in generators" : "Cannot replace return value of generators");
+            }
+        }
+    }
+
     @InliningCutoff
     private void notifyRootReturn(VirtualFrame virtualFrame, Object value) {
         WrapperNode wrapper = (WrapperNode) instrumentationRoot;
@@ -2307,13 +2316,12 @@ private Object enterRoot(VirtualFrame virtualFrame) {
                 wrapper.getProbeNode().onEnter(virtualFrame);
             } catch (Throwable t) {
                 Object result = wrapper.getProbeNode().onReturnExceptionalOrUnwind(virtualFrame, t, false);
+                checkOnReturnExceptionalOrUnwindResult(result);
                 if (result == ProbeNode.UNWIND_ACTION_REENTER) {
-                    // We're at the beginning, reenter means just continue
+                    // We're at the beginning, reenter means just restore args and continue
+                    copyArgs(virtualFrame.getArguments(), virtualFrame);
                     return null;
                 } else if (result != null) {
-                    if (co.isGeneratorOrCoroutine()) {
-                        throw CompilerDirectives.shouldNotReachHere("Cannot replace return value of generators");
-                    }
                     return result;
                 }
             }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/instrumentation/InstrumentationSupport.java

@@ -134,7 +134,7 @@ private static void handleException(VirtualFrame frame, InstrumentableNode.Wrapp
     public void notifyException(VirtualFrame frame, int line, Throwable exception) {
         InstrumentableNode.WrapperNode wrapper = getWrapperAtLine(line);
         if (wrapper != null) {
-            wrapper.getProbeNode().onReturnExceptionalOrUnwind(frame, exception, false);
+            handleException(frame, wrapper, exception, false);
         }
     }