Prerendering support for useDeferredValue (facebook#27512)

### Based on facebook#27509 

Revealing a prerendered tree (hidden -> visible) is considered the same
as mounting a brand new tree. So, when an initialValue argument is
passed to useDeferredValue, and it's prerendered inside a hidden tree,
we should first prerender the initial value.

After the initial value has been prerendered, we switch to prerendering
the final one. This is the same sequence that we use when mounting new
visible tree. Depending on how much prerendering work has been finished
by the time the tree is revealed, we may or may not be able to skip all
the way to the final value.

This means we get the benefits of both prerendering and preview states:
if we have enough resources to prerender the whole thing, we do that. If
we don't, we have a preview state to show for immediate feedback.

Author: acdlite

packages/react-reconciler/src/ReactFiberHooks.js

@@ -150,6 +150,7 @@ import {
 } from './ReactFiberAsyncAction';
 import {HostTransitionContext} from './ReactFiberHostContext';
 import {requestTransitionLane} from './ReactFiberRootScheduler';
+import {isCurrentTreeHidden} from './ReactFiberHiddenContext';
 
 const {ReactCurrentDispatcher, ReactCurrentBatchConfig} = ReactSharedInternals;
 
@@ -2688,12 +2689,6 @@ function mountDeferredValueImpl<T>(hook: Hook, value: T, initialValue?: T): T {
     );
     markSkippedUpdateLanes(deferredLane);
 
-    // Set this to true to indicate that the rendered value is inconsistent
-    // from the latest value. The name "baseState" doesn't really match how we
-    // use it because we're reusing a state hook field instead of creating a
-    // new one.
-    hook.baseState = true;
-
     return initialValue;
   } else {
     hook.memoizedState = value;
@@ -2705,17 +2700,33 @@ function updateDeferredValueImpl<T>(
   hook: Hook,
   prevValue: T,
   value: T,
-  initialValue: ?T,
+  initialValue?: T,
 ): T {
-  // TODO: We should also check if this component is going from
-  // hidden -> visible. If so, it should use the initialValue arg.
+  if (is(value, prevValue)) {
+    // The incoming value is referentially identical to the currently rendered
+    // value, so we can bail out quickly.
+    return value;
+  } else {
+    // Received a new value that's different from the current value.
+
+    // Check if we're inside a hidden tree
+    if (isCurrentTreeHidden()) {
+      // Revealing a prerendered tree is considered the same as mounting new
+      // one, so we reuse the "mount" path in this case.
+      const resultValue = mountDeferredValueImpl(hook, value, initialValue);
+      // Unlike during an actual mount, we need to mark this as an update if
+      // the value changed.
+      if (!is(resultValue, prevValue)) {
+        markWorkInProgressReceivedUpdate();
+      }
+      return resultValue;
+    }
 
-  const shouldDeferValue = !includesOnlyNonUrgentLanes(renderLanes);
-  if (shouldDeferValue) {
-    // This is an urgent update. If the value has changed, keep using the
-    // previous value and spawn a deferred render to update it later.
+    const shouldDeferValue = !includesOnlyNonUrgentLanes(renderLanes);
+    if (shouldDeferValue) {
+      // This is an urgent update. Since the value has changed, keep using the
+      // previous value and spawn a deferred render to update it later.
 
-    if (!is(value, prevValue)) {
       // Schedule a deferred render
       const deferredLane = requestDeferredLane();
       currentlyRenderingFiber.lanes = mergeLanes(
@@ -2724,33 +2735,18 @@ function updateDeferredValueImpl<T>(
       );
       markSkippedUpdateLanes(deferredLane);
 
-      // Set this to true to indicate that the rendered value is inconsistent
-      // from the latest value. The name "baseState" doesn't really match how we
-      // use it because we're reusing a state hook field instead of creating a
-      // new one.
-      hook.baseState = true;
-    }
-
-    // Reuse the previous value
-    return prevValue;
-  } else {
-    // This is not an urgent update, so we can use the latest value regardless
-    // of what it is. No need to defer it.
+      // Reuse the previous value. We do not need to mark this as an update,
+      // because we did not render a new value.
+      return prevValue;
+    } else {
+      // This is not an urgent update, so we can use the latest value regardless
+      // of what it is. No need to defer it.
 
-    // However, if we're currently inside a spawned render, then we need to mark
-    // this as an update to prevent the fiber from bailing out.
-    //
-    // `baseState` is true when the current value is different from the rendered
-    // value. The name doesn't really match how we use it because we're reusing
-    // a state hook field instead of creating a new one.
-    if (hook.baseState) {
-      // Flip this back to false.
-      hook.baseState = false;
+      // Mark this as an update to prevent the fiber from bailing out.
       markWorkInProgressReceivedUpdate();
+      hook.memoizedState = value;
+      return value;
     }
-
-    hook.memoizedState = value;
-    return value;
   }
 }
 

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

@@ -616,4 +616,190 @@ describe('ReactDeferredValue', () => {
     assertLog([]);
     expect(root).toMatchRenderedOutput(<div>Final</div>);
   });
+
+  // @gate enableUseDeferredValueInitialArg
+  // @gate enableOffscreen
+  it('useDeferredValue can prerender the initial value inside a hidden tree', async () => {
+    function App({text}) {
+      const renderedText = useDeferredValue(text, `Preview [${text}]`);
+      return (
+        <div>
+          <Text text={renderedText} />
+        </div>
+      );
+    }
+
+    let revealContent;
+    function Container({children}) {
+      const [shouldShow, setState] = useState(false);
+      revealContent = () => setState(true);
+      return (
+        <Offscreen mode={shouldShow ? 'visible' : 'hidden'}>
+          {children}
+        </Offscreen>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+
+    // Prerender some content
+    await act(() => {
+      root.render(
+        <Container>
+          <App text="A" />
+        </Container>,
+      );
+    });
+    assertLog(['Preview [A]', 'A']);
+    expect(root).toMatchRenderedOutput(<div hidden={true}>A</div>);
+
+    await act(async () => {
+      // While the tree is still hidden, update the pre-rendered tree.
+      root.render(
+        <Container>
+          <App text="B" />
+        </Container>,
+      );
+      // We should switch to pre-rendering the new preview.
+      await waitForPaint(['Preview [B]']);
+      expect(root).toMatchRenderedOutput(<div hidden={true}>Preview [B]</div>);
+
+      // Before the prerender is complete, reveal the hidden tree. Because we
+      // consider revealing a hidden tree to be the same as mounting a new one,
+      // we should not skip the preview state.
+      revealContent();
+      // Because the preview state was already prerendered, we can reveal it
+      // without any addditional work.
+      await waitForPaint([]);
+      expect(root).toMatchRenderedOutput(<div>Preview [B]</div>);
+    });
+    // Finally, finish rendering the final value.
+    assertLog(['B']);
+    expect(root).toMatchRenderedOutput(<div>B</div>);
+  });
+
+  // @gate enableUseDeferredValueInitialArg
+  // @gate enableOffscreen
+  it(
+    'useDeferredValue skips the preview state when revealing a hidden tree ' +
+      'if the final value is referentially identical',
+    async () => {
+      function App({text}) {
+        const renderedText = useDeferredValue(text, `Preview [${text}]`);
+        return (
+          <div>
+            <Text text={renderedText} />
+          </div>
+        );
+      }
+
+      function Container({text, shouldShow}) {
+        return (
+          <Offscreen mode={shouldShow ? 'visible' : 'hidden'}>
+            <App text={text} />
+          </Offscreen>
+        );
+      }
+
+      const root = ReactNoop.createRoot();
+
+      // Prerender some content
+      await act(() => root.render(<Container text="A" shouldShow={false} />));
+      assertLog(['Preview [A]', 'A']);
+      expect(root).toMatchRenderedOutput(<div hidden={true}>A</div>);
+
+      // Reveal the prerendered tree. Because the final value is referentially
+      // equal to what was already prerendered, we can skip the preview state
+      // and go straight to the final one. The practical upshot of this is
+      // that we can completely prerender the final value without having to
+      // do additional rendering work when the tree is revealed.
+      await act(() => root.render(<Container text="A" shouldShow={true} />));
+      assertLog(['A']);
+      expect(root).toMatchRenderedOutput(<div>A</div>);
+    },
+  );
+
+  // @gate enableUseDeferredValueInitialArg
+  // @gate enableOffscreen
+  it(
+    'useDeferredValue does not skip the preview state when revealing a ' +
+      'hidden tree if the final value is different from the currently rendered one',
+    async () => {
+      function App({text}) {
+        const renderedText = useDeferredValue(text, `Preview [${text}]`);
+        return (
+          <div>
+            <Text text={renderedText} />
+          </div>
+        );
+      }
+
+      function Container({text, shouldShow}) {
+        return (
+          <Offscreen mode={shouldShow ? 'visible' : 'hidden'}>
+            <App text={text} />
+          </Offscreen>
+        );
+      }
+
+      const root = ReactNoop.createRoot();
+
+      // Prerender some content
+      await act(() => root.render(<Container text="A" shouldShow={false} />));
+      assertLog(['Preview [A]', 'A']);
+      expect(root).toMatchRenderedOutput(<div hidden={true}>A</div>);
+
+      // Reveal the prerendered tree. Because the final value is different from
+      // what was already prerendered, we can't bail out. Since we treat
+      // revealing a hidden tree the same as a new mount, show the preview state
+      // before switching to the final one.
+      await act(async () => {
+        root.render(<Container text="B" shouldShow={true} />);
+        // First commit the preview state
+        await waitForPaint(['Preview [B]']);
+        expect(root).toMatchRenderedOutput(<div>Preview [B]</div>);
+      });
+      // Then switch to the final state
+      assertLog(['B']);
+      expect(root).toMatchRenderedOutput(<div>B</div>);
+    },
+  );
+
+  // @gate enableOffscreen
+  it(
+    'useDeferredValue does not show "previous" value when revealing a hidden ' +
+      'tree (no initial value)',
+    async () => {
+      function App({text}) {
+        const renderedText = useDeferredValue(text);
+        return (
+          <div>
+            <Text text={renderedText} />
+          </div>
+        );
+      }
+
+      function Container({text, shouldShow}) {
+        return (
+          <Offscreen mode={shouldShow ? 'visible' : 'hidden'}>
+            <App text={text} />
+          </Offscreen>
+        );
+      }
+
+      const root = ReactNoop.createRoot();
+
+      // Prerender some content
+      await act(() => root.render(<Container text="A" shouldShow={false} />));
+      assertLog(['A']);
+      expect(root).toMatchRenderedOutput(<div hidden={true}>A</div>);
+
+      // Update the prerendered tree and reveal it at the same time. Even though
+      // this is a sync update, we should update B immediately rather than stay
+      // on the old value (A), because conceptually this is a new tree.
+      await act(() => root.render(<Container text="B" shouldShow={true} />));
+      assertLog(['B']);
+      expect(root).toMatchRenderedOutput(<div>B</div>);
+    },
+  );
 });