From 422d102bd58ab16f93b6d84a2a25c759b6a1288f Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Fri, 1 Apr 2022 13:46:55 -0400
Subject: [PATCH 1/2] Fix infinite loop if unmemoized val passed to uDV

The current implementation of useDeferredValue will spawn a new
render any time the input value is different from the previous one. So
if you pass an unmemoized value (like an inline object), it will never
stop spawning new renders.

The fix is to only defer during an urgent render. If we're already
inside a transition, retry, offscreen, or other non-urgen render, then
we can use the latest value.
---
 .../src/ReactFiberHooks.new.js                |  94 +++++---
 .../src/ReactFiberHooks.old.js                |  94 +++++---
 .../src/__tests__/ReactDeferredValue-test.js  | 224 ++++++++++++++++++
 3 files changed, 348 insertions(+), 64 deletions(-)
 create mode 100644 packages/react-reconciler/src/__tests__/ReactDeferredValue-test.js

diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js
index d1c50b42d8cb..1008f835b6e7 100644
--- a/packages/react-reconciler/src/ReactFiberHooks.new.js
+++ b/packages/react-reconciler/src/ReactFiberHooks.new.js
@@ -47,6 +47,8 @@ import {
   NoLanes,
   isSubsetOfLanes,
   includesBlockingLane,
+  includesOnlyNonUrgentLanes,
+  claimNextTransitionLane,
   mergeLanes,
   removeLanes,
   intersectLanes,
@@ -1929,45 +1931,73 @@ function updateMemo<T>(
 }
 
 function mountDeferredValue<T>(value: T): T {
-  const [prevValue, setValue] = mountState(value);
-  mountEffect(() => {
-    const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = {};
-    try {
-      setValue(value);
-    } finally {
-      ReactCurrentBatchConfig.transition = prevTransition;
-    }
-  }, [value]);
-  return prevValue;
+  const hook = mountWorkInProgressHook();
+  hook.memoizedState = value;
+  return value;
 }
 
 function updateDeferredValue<T>(value: T): T {
-  const [prevValue, setValue] = updateState(value);
-  updateEffect(() => {
-    const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = {};
-    try {
-      setValue(value);
-    } finally {
-      ReactCurrentBatchConfig.transition = prevTransition;
-    }
-  }, [value]);
-  return prevValue;
+  const hook = updateWorkInProgressHook();
+  const resolvedCurrentHook: Hook = (currentHook: any);
+  const prevValue: T = resolvedCurrentHook.memoizedState;
+  return updateDeferredValueImpl(hook, prevValue, value);
 }
 
 function rerenderDeferredValue<T>(value: T): T {
-  const [prevValue, setValue] = rerenderState(value);
-  updateEffect(() => {
-    const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = {};
-    try {
-      setValue(value);
-    } finally {
-      ReactCurrentBatchConfig.transition = prevTransition;
+  const hook = updateWorkInProgressHook();
+  if (currentHook === null) {
+    // This is a rerender during a mount.
+    hook.memoizedState = value;
+    return value;
+  } else {
+    // This is a rerender during an update.
+    const prevValue: T = currentHook.memoizedState;
+    return updateDeferredValueImpl(hook, prevValue, value);
+  }
+}
+
+function updateDeferredValueImpl<T>(hook: Hook, prevValue: T, value: T): T {
+  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.
+
+    if (!is(value, prevValue)) {
+      // Schedule a deferred render
+      const deferredLane = claimNextTransitionLane();
+      currentlyRenderingFiber.lanes = mergeLanes(
+        currentlyRenderingFiber.lanes,
+        deferredLane,
+      );
+      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;
     }
-  }, [value]);
-  return prevValue;
+
+    // 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.
+
+    // 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;
+      markWorkInProgressReceivedUpdate();
+    }
+
+    return value;
+  }
 }
 
 function startTransition(setPending, callback, options) {
diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js
index 137ee33cba51..139e86922422 100644
--- a/packages/react-reconciler/src/ReactFiberHooks.old.js
+++ b/packages/react-reconciler/src/ReactFiberHooks.old.js
@@ -47,6 +47,8 @@ import {
   NoLanes,
   isSubsetOfLanes,
   includesBlockingLane,
+  includesOnlyNonUrgentLanes,
+  claimNextTransitionLane,
   mergeLanes,
   removeLanes,
   intersectLanes,
@@ -1929,45 +1931,73 @@ function updateMemo<T>(
 }
 
 function mountDeferredValue<T>(value: T): T {
-  const [prevValue, setValue] = mountState(value);
-  mountEffect(() => {
-    const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = {};
-    try {
-      setValue(value);
-    } finally {
-      ReactCurrentBatchConfig.transition = prevTransition;
-    }
-  }, [value]);
-  return prevValue;
+  const hook = mountWorkInProgressHook();
+  hook.memoizedState = value;
+  return value;
 }
 
 function updateDeferredValue<T>(value: T): T {
-  const [prevValue, setValue] = updateState(value);
-  updateEffect(() => {
-    const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = {};
-    try {
-      setValue(value);
-    } finally {
-      ReactCurrentBatchConfig.transition = prevTransition;
-    }
-  }, [value]);
-  return prevValue;
+  const hook = updateWorkInProgressHook();
+  const resolvedCurrentHook: Hook = (currentHook: any);
+  const prevValue: T = resolvedCurrentHook.memoizedState;
+  return updateDeferredValueImpl(hook, prevValue, value);
 }
 
 function rerenderDeferredValue<T>(value: T): T {
-  const [prevValue, setValue] = rerenderState(value);
-  updateEffect(() => {
-    const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = {};
-    try {
-      setValue(value);
-    } finally {
-      ReactCurrentBatchConfig.transition = prevTransition;
+  const hook = updateWorkInProgressHook();
+  if (currentHook === null) {
+    // This is a rerender during a mount.
+    hook.memoizedState = value;
+    return value;
+  } else {
+    // This is a rerender during an update.
+    const prevValue: T = currentHook.memoizedState;
+    return updateDeferredValueImpl(hook, prevValue, value);
+  }
+}
+
+function updateDeferredValueImpl<T>(hook: Hook, prevValue: T, value: T): T {
+  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.
+
+    if (!is(value, prevValue)) {
+      // Schedule a deferred render
+      const deferredLane = claimNextTransitionLane();
+      currentlyRenderingFiber.lanes = mergeLanes(
+        currentlyRenderingFiber.lanes,
+        deferredLane,
+      );
+      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;
     }
-  }, [value]);
-  return prevValue;
+
+    // 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.
+
+    // 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;
+      markWorkInProgressReceivedUpdate();
+    }
+
+    return value;
+  }
 }
 
 function startTransition(setPending, callback, options) {
diff --git a/packages/react-reconciler/src/__tests__/ReactDeferredValue-test.js b/packages/react-reconciler/src/__tests__/ReactDeferredValue-test.js
new file mode 100644
index 000000000000..9cd07c066e4c
--- /dev/null
+++ b/packages/react-reconciler/src/__tests__/ReactDeferredValue-test.js
@@ -0,0 +1,224 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+'use strict';
+
+let React;
+let ReactNoop;
+let Scheduler;
+let act;
+let startTransition;
+let useDeferredValue;
+let useMemo;
+let useState;
+
+describe('ReactDeferredValue', () => {
+  beforeEach(() => {
+    jest.resetModules();
+
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+    Scheduler = require('scheduler');
+    act = require('jest-react').act;
+    startTransition = React.startTransition;
+    useDeferredValue = React.useDeferredValue;
+    useMemo = React.useMemo;
+    useState = React.useState;
+  });
+
+  function Text({text}) {
+    Scheduler.unstable_yieldValue(text);
+    return text;
+  }
+
+  it('does not cause an infinite defer loop if the original value isn\t memoized', async () => {
+    function App({value}) {
+      // The object passed to useDeferredValue is never the same as the previous
+      // render. A naive implementation would endlessly spawn deferred renders.
+      const {value: deferredValue} = useDeferredValue({value});
+
+      const child = useMemo(() => <Text text={'Original: ' + value} />, [
+        value,
+      ]);
+
+      const deferredChild = useMemo(
+        () => <Text text={'Deferred: ' + deferredValue} />,
+        [deferredValue],
+      );
+
+      return (
+        <div>
+          <div>{child}</div>
+          <div>{deferredChild}</div>
+        </div>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+
+    // Initial render
+    await act(async () => {
+      root.render(<App value={1} />);
+    });
+    expect(Scheduler).toHaveYielded(['Original: 1', 'Deferred: 1']);
+
+    // If it's an urgent update, the value is deferred
+    await act(async () => {
+      root.render(<App value={2} />);
+
+      expect(Scheduler).toFlushUntilNextPaint(['Original: 2']);
+      // The deferred value updates in a separate render
+      expect(Scheduler).toFlushUntilNextPaint(['Deferred: 2']);
+    });
+    expect(root).toMatchRenderedOutput(
+      <div>
+        <div>Original: 2</div>
+        <div>Deferred: 2</div>
+      </div>,
+    );
+
+    // But if it updates during a transition, it doesn't defer
+    await act(async () => {
+      startTransition(() => {
+        root.render(<App value={3} />);
+      });
+      // The deferred value updates in the same render as the original
+      expect(Scheduler).toFlushUntilNextPaint(['Original: 3', 'Deferred: 3']);
+    });
+    expect(root).toMatchRenderedOutput(
+      <div>
+        <div>Original: 3</div>
+        <div>Deferred: 3</div>
+      </div>,
+    );
+  });
+
+  it('does not defer during a transition', async () => {
+    function App({value}) {
+      const deferredValue = useDeferredValue(value);
+
+      const child = useMemo(() => <Text text={'Original: ' + value} />, [
+        value,
+      ]);
+
+      const deferredChild = useMemo(
+        () => <Text text={'Deferred: ' + deferredValue} />,
+        [deferredValue],
+      );
+
+      return (
+        <div>
+          <div>{child}</div>
+          <div>{deferredChild}</div>
+        </div>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+
+    // Initial render
+    await act(async () => {
+      root.render(<App value={1} />);
+    });
+    expect(Scheduler).toHaveYielded(['Original: 1', 'Deferred: 1']);
+
+    // If it's an urgent update, the value is deferred
+    await act(async () => {
+      root.render(<App value={2} />);
+
+      expect(Scheduler).toFlushUntilNextPaint(['Original: 2']);
+      // The deferred value updates in a separate render
+      expect(Scheduler).toFlushUntilNextPaint(['Deferred: 2']);
+    });
+    expect(root).toMatchRenderedOutput(
+      <div>
+        <div>Original: 2</div>
+        <div>Deferred: 2</div>
+      </div>,
+    );
+
+    // But if it updates during a transition, it doesn't defer
+    await act(async () => {
+      startTransition(() => {
+        root.render(<App value={3} />);
+      });
+      // The deferred value updates in the same render as the original
+      expect(Scheduler).toFlushUntilNextPaint(['Original: 3', 'Deferred: 3']);
+    });
+    expect(root).toMatchRenderedOutput(
+      <div>
+        <div>Original: 3</div>
+        <div>Deferred: 3</div>
+      </div>,
+    );
+  });
+
+  it("works if there's a render phase update", async () => {
+    function App({value: propValue}) {
+      const [value, setValue] = useState(null);
+      if (value !== propValue) {
+        setValue(propValue);
+      }
+
+      const deferredValue = useDeferredValue(value);
+
+      const child = useMemo(() => <Text text={'Original: ' + value} />, [
+        value,
+      ]);
+
+      const deferredChild = useMemo(
+        () => <Text text={'Deferred: ' + deferredValue} />,
+        [deferredValue],
+      );
+
+      return (
+        <div>
+          <div>{child}</div>
+          <div>{deferredChild}</div>
+        </div>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+
+    // Initial render
+    await act(async () => {
+      root.render(<App value={1} />);
+    });
+    expect(Scheduler).toHaveYielded(['Original: 1', 'Deferred: 1']);
+
+    // If it's an urgent update, the value is deferred
+    await act(async () => {
+      root.render(<App value={2} />);
+
+      expect(Scheduler).toFlushUntilNextPaint(['Original: 2']);
+      // The deferred value updates in a separate render
+      expect(Scheduler).toFlushUntilNextPaint(['Deferred: 2']);
+    });
+    expect(root).toMatchRenderedOutput(
+      <div>
+        <div>Original: 2</div>
+        <div>Deferred: 2</div>
+      </div>,
+    );
+
+    // But if it updates during a transition, it doesn't defer
+    await act(async () => {
+      startTransition(() => {
+        root.render(<App value={3} />);
+      });
+      // The deferred value updates in the same render as the original
+      expect(Scheduler).toFlushUntilNextPaint(['Original: 3', 'Deferred: 3']);
+    });
+    expect(root).toMatchRenderedOutput(
+      <div>
+        <div>Original: 3</div>
+        <div>Deferred: 3</div>
+      </div>,
+    );
+  });
+});

From e4fd270c7a8a91bcb3f75f732d6e227755a9c160 Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Mon, 4 Apr 2022 12:14:00 -0400
Subject: [PATCH 2/2] Temporarily disable "long nested update" warning

DevTools' timeline profiler warns if an update inside a layout effect
results in an expensive re-render. However, it misattributes renders
that are spawned from a sync render at lower priority. This affects the
new implementation of useDeferredValue but it would also apply to things
like Offscreen.

It's not obvious to me how to fix this given how DevTools models the
idea of a "nested update" so I'm disabling the warning for now to
unblock the bugfix for useDeferredValue.
---
 .../src/__tests__/preprocessData-test.js             | 12 +++++++++---
 .../src/import-worker/preprocessData.js              |  5 ++++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/packages/react-devtools-shared/src/__tests__/preprocessData-test.js b/packages/react-devtools-shared/src/__tests__/preprocessData-test.js
index 72086a684cd4..880607ff1d76 100644
--- a/packages/react-devtools-shared/src/__tests__/preprocessData-test.js
+++ b/packages/react-devtools-shared/src/__tests__/preprocessData-test.js
@@ -1463,7 +1463,9 @@ describe('Timeline profiler', () => {
             expect(event.warning).toBe(null);
           });
 
-          it('should warn about long nested (state) updates during layout effects', async () => {
+          // This is temporarily disabled because the warning doesn't work
+          // with useDeferredValue
+          it.skip('should warn about long nested (state) updates during layout effects', async () => {
             function Component() {
               const [didMount, setDidMount] = React.useState(false);
               Scheduler.unstable_yieldValue(
@@ -1523,7 +1525,9 @@ describe('Timeline profiler', () => {
             );
           });
 
-          it('should warn about long nested (forced) updates during layout effects', async () => {
+          // This is temporarily disabled because the warning doesn't work
+          // with useDeferredValue
+          it.skip('should warn about long nested (forced) updates during layout effects', async () => {
             class Component extends React.Component {
               _didMount: boolean = false;
               componentDidMount() {
@@ -1654,7 +1658,9 @@ describe('Timeline profiler', () => {
             });
           });
 
-          it('should not warn about deferred value updates scheduled during commit phase', async () => {
+          // This is temporarily disabled because the warning doesn't work
+          // with useDeferredValue
+          it.skip('should not warn about deferred value updates scheduled during commit phase', async () => {
             function Component() {
               const [value, setValue] = React.useState(0);
               const deferredValue = React.useDeferredValue(value);
diff --git a/packages/react-devtools-timeline/src/import-worker/preprocessData.js b/packages/react-devtools-timeline/src/import-worker/preprocessData.js
index 258eef931a38..110b4e8923fb 100644
--- a/packages/react-devtools-timeline/src/import-worker/preprocessData.js
+++ b/packages/react-devtools-timeline/src/import-worker/preprocessData.js
@@ -1124,7 +1124,10 @@ export default async function preprocessData(
           lane => profilerData.laneToLabelMap.get(lane) === 'Transition',
         )
       ) {
-        schedulingEvent.warning = WARNING_STRINGS.NESTED_UPDATE;
+        // FIXME: This warning doesn't account for "nested updates" that are
+        // spawned by useDeferredValue. Disabling temporarily until we figure
+        // out the right way to handle this.
+        // schedulingEvent.warning = WARNING_STRINGS.NESTED_UPDATE;
       }
     }
   });