Implement useSyncExternalStore in Fiber (#22239)

This adds an initial implementation of useSyncExternalStore to the
fiber reconciler. It's mostly a copy-paste of the userspace
implementation, which is not ideal but is a good enough starting place.

The main change we'll want to make to this native implementation is to
move the tearing checks from the layout phase to an earlier, pre-commit
phase so that code that runs in the commit phase always observes a
consistent tree.

Follow-ups:

- Implement in Fizz
- Implement in old SSR renderer
- Implement in react-debug-hooks

Author: acdlite

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

@@ -51,6 +51,7 @@ import {
   markRootMutableRead,
 } from './ReactFiberLane.new';
 import {
+  DiscreteEventPriority,
   ContinuousEventPriority,
   getCurrentUpdatePriority,
   setCurrentUpdatePriority,
@@ -136,6 +137,7 @@ export type UpdateQueue<S, A> = {|
 
 let didWarnAboutMismatchedHooksForComponent;
 let didWarnAboutUseOpaqueIdentifier;
+let didWarnUncachedGetSnapshot;
 if (__DEV__) {
   didWarnAboutUseOpaqueIdentifier = {};
   didWarnAboutMismatchedHooksForComponent = new Set();
@@ -1246,14 +1248,127 @@ function mountSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
 ): T {
-  throw new Error('Not yet implemented');
+  const hook = mountWorkInProgressHook();
+  return useSyncExternalStore(hook, subscribe, getSnapshot);
 }
 
 function updateSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
 ): T {
-  throw new Error('Not yet implemented');
+  const hook = updateWorkInProgressHook();
+  return useSyncExternalStore(hook, subscribe, getSnapshot);
+}
+
+function useSyncExternalStore<T>(
+  hook: Hook,
+  subscribe: (() => void) => () => void,
+  getSnapshot: () => T,
+): T {
+  // TODO: This is a copy-paste of the userspace shim. We can improve the
+  // built-in implementation using lower-level APIs. We also intend to move
+  // the tearing checks to an earlier, pre-commit phase so that the layout
+  // effects always observe a consistent tree.
+
+  const dispatcher = ReactCurrentDispatcher.current;
+
+  // Read the current snapshot from the store on every render. Again, this
+  // breaks the rules of React, and only works here because of specific
+  // implementation details, most importantly that updates are
+  // always synchronous.
+  const value = getSnapshot();
+  if (__DEV__) {
+    if (!didWarnUncachedGetSnapshot) {
+      if (value !== getSnapshot()) {
+        console.error(
+          'The result of getSnapshot should be cached to avoid an infinite loop',
+        );
+        didWarnUncachedGetSnapshot = true;
+      }
+    }
+  }
+
+  // Because updates are synchronous, we don't queue them. Instead we force a
+  // re-render whenever the subscribed state changes by updating an some
+  // arbitrary useState hook. Then, during render, we call getSnapshot to read
+  // the current value.
+  //
+  // Because we don't actually use the state returned by the useState hook, we
+  // can save a bit of memory by storing other stuff in that slot.
+  //
+  // To implement the early bailout, we need to track some things on a mutable
+  // object. Usually, we would put that in a useRef hook, but we can stash it in
+  // our useState hook instead.
+  //
+  // To force a re-render, we call forceUpdate({inst}). That works because the
+  // new object always fails an equality check.
+  const [{inst}, forceUpdate] = dispatcher.useState({
+    inst: {value, getSnapshot},
+  });
+
+  // Track the latest getSnapshot function with a ref. This needs to be updated
+  // in the layout phase so we can access it during the tearing check that
+  // happens on subscribe.
+  // TODO: Circumvent SSR warning
+  dispatcher.useLayoutEffect(() => {
+    inst.value = value;
+    inst.getSnapshot = getSnapshot;
+
+    // Whenever getSnapshot or subscribe changes, we need to check in the
+    // commit phase if there was an interleaved mutation. In concurrent mode
+    // this can happen all the time, but even in synchronous mode, an earlier
+    // effect may have mutated the store.
+    if (checkIfSnapshotChanged(inst)) {
+      // Force a re-render.
+      const prevTransition = ReactCurrentBatchConfig.transition;
+      const prevPriority = getCurrentUpdatePriority();
+      ReactCurrentBatchConfig.transition = 0;
+      setCurrentUpdatePriority(DiscreteEventPriority);
+      forceUpdate({inst});
+      setCurrentUpdatePriority(prevPriority);
+      ReactCurrentBatchConfig.transition = prevTransition;
+    }
+  }, [subscribe, value, getSnapshot]);
+
+  dispatcher.useEffect(() => {
+    const handleStoreChange = () => {
+      // TODO: Because there is no cross-renderer API for batching updates, it's
+      // up to the consumer of this library to wrap their subscription event
+      // with unstable_batchedUpdates. Should we try to detect when this isn't
+      // the case and print a warning in development?
+
+      // The store changed. Check if the snapshot changed since the last time we
+      // read from the store.
+      if (checkIfSnapshotChanged(inst)) {
+        // Force a re-render.
+        const prevTransition = ReactCurrentBatchConfig.transition;
+        const prevPriority = getCurrentUpdatePriority();
+        ReactCurrentBatchConfig.transition = 0;
+        setCurrentUpdatePriority(DiscreteEventPriority);
+        forceUpdate({inst});
+        setCurrentUpdatePriority(prevPriority);
+        ReactCurrentBatchConfig.transition = prevTransition;
+      }
+    };
+    // Check for changes right before subscribing. Subsequent changes will be
+    // detected in the subscription handler.
+    handleStoreChange();
+    // Subscribe to the store and return a clean-up function.
+    return subscribe(handleStoreChange);
+  }, [subscribe]);
+
+  return value;
+}
+
+function checkIfSnapshotChanged(inst) {
+  const latestGetSnapshot = inst.getSnapshot;
+  const prevValue = inst.value;
+  try {
+    const nextValue = latestGetSnapshot();
+    return !is(prevValue, nextValue);
+  } catch (error) {
+    return true;
+  }
 }
 
 function mountState<S>(

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

@@ -51,6 +51,7 @@ import {
   markRootMutableRead,
 } from './ReactFiberLane.old';
 import {
+  DiscreteEventPriority,
   ContinuousEventPriority,
   getCurrentUpdatePriority,
   setCurrentUpdatePriority,
@@ -136,6 +137,7 @@ export type UpdateQueue<S, A> = {|
 
 let didWarnAboutMismatchedHooksForComponent;
 let didWarnAboutUseOpaqueIdentifier;
+let didWarnUncachedGetSnapshot;
 if (__DEV__) {
   didWarnAboutUseOpaqueIdentifier = {};
   didWarnAboutMismatchedHooksForComponent = new Set();
@@ -1246,14 +1248,127 @@ function mountSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
 ): T {
-  throw new Error('Not yet implemented');
+  const hook = mountWorkInProgressHook();
+  return useSyncExternalStore(hook, subscribe, getSnapshot);
 }
 
 function updateSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
 ): T {
-  throw new Error('Not yet implemented');
+  const hook = updateWorkInProgressHook();
+  return useSyncExternalStore(hook, subscribe, getSnapshot);
+}
+
+function useSyncExternalStore<T>(
+  hook: Hook,
+  subscribe: (() => void) => () => void,
+  getSnapshot: () => T,
+): T {
+  // TODO: This is a copy-paste of the userspace shim. We can improve the
+  // built-in implementation using lower-level APIs. We also intend to move
+  // the tearing checks to an earlier, pre-commit phase so that the layout
+  // effects always observe a consistent tree.
+
+  const dispatcher = ReactCurrentDispatcher.current;
+
+  // Read the current snapshot from the store on every render. Again, this
+  // breaks the rules of React, and only works here because of specific
+  // implementation details, most importantly that updates are
+  // always synchronous.
+  const value = getSnapshot();
+  if (__DEV__) {
+    if (!didWarnUncachedGetSnapshot) {
+      if (value !== getSnapshot()) {
+        console.error(
+          'The result of getSnapshot should be cached to avoid an infinite loop',
+        );
+        didWarnUncachedGetSnapshot = true;
+      }
+    }
+  }
+
+  // Because updates are synchronous, we don't queue them. Instead we force a
+  // re-render whenever the subscribed state changes by updating an some
+  // arbitrary useState hook. Then, during render, we call getSnapshot to read
+  // the current value.
+  //
+  // Because we don't actually use the state returned by the useState hook, we
+  // can save a bit of memory by storing other stuff in that slot.
+  //
+  // To implement the early bailout, we need to track some things on a mutable
+  // object. Usually, we would put that in a useRef hook, but we can stash it in
+  // our useState hook instead.
+  //
+  // To force a re-render, we call forceUpdate({inst}). That works because the
+  // new object always fails an equality check.
+  const [{inst}, forceUpdate] = dispatcher.useState({
+    inst: {value, getSnapshot},
+  });
+
+  // Track the latest getSnapshot function with a ref. This needs to be updated
+  // in the layout phase so we can access it during the tearing check that
+  // happens on subscribe.
+  // TODO: Circumvent SSR warning
+  dispatcher.useLayoutEffect(() => {
+    inst.value = value;
+    inst.getSnapshot = getSnapshot;
+
+    // Whenever getSnapshot or subscribe changes, we need to check in the
+    // commit phase if there was an interleaved mutation. In concurrent mode
+    // this can happen all the time, but even in synchronous mode, an earlier
+    // effect may have mutated the store.
+    if (checkIfSnapshotChanged(inst)) {
+      // Force a re-render.
+      const prevTransition = ReactCurrentBatchConfig.transition;
+      const prevPriority = getCurrentUpdatePriority();
+      ReactCurrentBatchConfig.transition = 0;
+      setCurrentUpdatePriority(DiscreteEventPriority);
+      forceUpdate({inst});
+      setCurrentUpdatePriority(prevPriority);
+      ReactCurrentBatchConfig.transition = prevTransition;
+    }
+  }, [subscribe, value, getSnapshot]);
+
+  dispatcher.useEffect(() => {
+    const handleStoreChange = () => {
+      // TODO: Because there is no cross-renderer API for batching updates, it's
+      // up to the consumer of this library to wrap their subscription event
+      // with unstable_batchedUpdates. Should we try to detect when this isn't
+      // the case and print a warning in development?
+
+      // The store changed. Check if the snapshot changed since the last time we
+      // read from the store.
+      if (checkIfSnapshotChanged(inst)) {
+        // Force a re-render.
+        const prevTransition = ReactCurrentBatchConfig.transition;
+        const prevPriority = getCurrentUpdatePriority();
+        ReactCurrentBatchConfig.transition = 0;
+        setCurrentUpdatePriority(DiscreteEventPriority);
+        forceUpdate({inst});
+        setCurrentUpdatePriority(prevPriority);
+        ReactCurrentBatchConfig.transition = prevTransition;
+      }
+    };
+    // Check for changes right before subscribing. Subsequent changes will be
+    // detected in the subscription handler.
+    handleStoreChange();
+    // Subscribe to the store and return a clean-up function.
+    return subscribe(handleStoreChange);
+  }, [subscribe]);
+
+  return value;
+}
+
+function checkIfSnapshotChanged(inst) {
+  const latestGetSnapshot = inst.getSnapshot;
+  const prevValue = inst.value;
+  try {
+    const nextValue = latestGetSnapshot();
+    return !is(prevValue, nextValue);
+  } catch (error) {
+    return true;
+  }
 }
 
 function mountState<S>(