Unwind the current workInProgress if it's suspended (#25247)

Usually we complete workInProgress before yielding but if that's the
currently suspended one, we don't yet complete it in case we can
immediately unblock it.

If we get interrupted, however, we must unwind it. Where as we usually
assume that we've already completed it.

This shows up when the current work in progress was a Context that pushed
and then it suspends in its immediate children. If we don't unwind,
it won't pop and so we get an imbalance.

packages/react-reconciler/src/ReactFiberWorkLoop.new.js

@@ -1651,7 +1651,9 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   }
 
   if (workInProgress !== null) {
-    let interruptedWork = workInProgress.return;
+    let interruptedWork = workInProgressIsSuspended
+      ? workInProgress
+      : workInProgress.return;
     while (interruptedWork !== null) {
       const current = interruptedWork.alternate;
       unwindInterruptedWork(

packages/react-reconciler/src/ReactFiberWorkLoop.old.js

@@ -1651,7 +1651,9 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   }
 
   if (workInProgress !== null) {
-    let interruptedWork = workInProgress.return;
+    let interruptedWork = workInProgressIsSuspended
+      ? workInProgress
+      : workInProgress.return;
     while (interruptedWork !== null) {
       const current = interruptedWork.alternate;
       unwindInterruptedWork(

packages/react-reconciler/src/__tests__/ReactWakeable-test.js

@@ -339,4 +339,51 @@ describe('ReactWakeable', () => {
     expect(Scheduler).toFlushWithoutYielding();
     expect(root).toMatchRenderedOutput('AB');
   });
+
+  // @gate enableUseHook
+  test('interrupting while yielded should reset contexts', async () => {
+    let resolve;
+    const promise = new Promise(r => {
+      resolve = r;
+    });
+
+    const Context = React.createContext();
+
+    const lazy = React.lazy(() => {
+      return promise;
+    });
+
+    function ContextText() {
+      return <Text text={use(Context)} />;
+    }
+
+    function App({text}) {
+      return (
+        <div>
+          <Context.Provider value={text}>
+            {lazy}
+            <ContextText />
+          </Context.Provider>
+        </div>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    startTransition(() => {
+      root.render(<App text="world" />);
+    });
+    expect(Scheduler).toFlushUntilNextPaint([]);
+    expect(root).toMatchRenderedOutput(null);
+
+    await resolve({default: <Text key="hi" text="Hello " />});
+
+    // Higher priority update that interrupts the first render
+    ReactNoop.flushSync(() => {
+      root.render(<App text="world!" />);
+    });
+
+    expect(Scheduler).toHaveYielded(['Hello ', 'world!']);
+
+    expect(root).toMatchRenderedOutput(<div>Hello world!</div>);
+  });
 });