[useSyncExternalStore] Remove extra hook object

Because we already track `getSnapshot` and `value` on the store
instance, we don't need to also track them as effect dependencies. And
because the effect doesn't require any clean-up, we don't need to track
a `destroy` function.

So, we don't need to store any additional state for this effect. We can
call `pushEffect` directly, and only during renders where something
has changed.

This saves some memory, but my main motivation is because I plan to use
this same logic to schedule a pre-commit consistency check. (See the
inline comments for more details.)

Author: acdlite

packages/react-debug-tools/src/ReactDebugHooks.js

@@ -273,7 +273,6 @@ function useSyncExternalStore<T>(
   // Advance the current hook index the same number of times
   // so that subsequent hooks have the right memoized state.
   nextHook(); // SyncExternalStore
-  nextHook(); // LayoutEffect
   nextHook(); // Effect
   const value = getSnapshot();
   hookLog.push({

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

@@ -1256,6 +1256,7 @@ function mountSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
 ): T {
+  const fiber = currentlyRenderingFiber;
   const hook = mountWorkInProgressHook();
   // Read the current snapshot from the store on every render. This breaks the
   // normal rules of React, and only works because store updates are
@@ -1277,13 +1278,37 @@ function mountSyncExternalStore<T>(
     getSnapshot,
   };
   hook.queue = inst;
-  return useSyncExternalStore(hook, inst, subscribe, getSnapshot, nextSnapshot);
+
+  // Schedule an effect to subscribe to the store.
+  mountEffect(subscribeToStore.bind(null, fiber, inst, subscribe), [subscribe]);
+
+  // Schedule an effect to update the mutable instance fields. We will update
+  // this whenever subscribe, getSnapshot, or value changes. Because there's no
+  // clean-up function, and we track the deps correctly, we can call pushEffect
+  // directly, without storing any additional state. For the same reason, we
+  // don't need to set a static flag, either.
+  // TODO: We can move this to the passive phase once we add a pre-commit
+  // consistency check. See the next comment.
+  fiber.flags |= UpdateEffect;
+  pushEffect(
+    HookHasEffect | HookLayout,
+    updateStoreInstance.bind(null, fiber, inst, nextSnapshot, getSnapshot),
+    undefined,
+    null,
+  );
+  // TODO: Unless this is a synchronous render, schedule a consistency check.
+  // Right before committing, we will walk the tree and check if any of the
+  // stores were mutated.
+  // pushConsistencyCheck(inst, getSnapshot, nextSnapshot);
+
+  return nextSnapshot;
 }
 
 function updateSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
 ): T {
+  const fiber = currentlyRenderingFiber;
   const hook = updateWorkInProgressHook();
   // Read the current snapshot from the store on every render. This breaks the
   // normal rules of React, and only works because store updates are
@@ -1300,66 +1325,81 @@ function updateSyncExternalStore<T>(
     }
   }
   const prevSnapshot = hook.memoizedState;
-  if (!is(prevSnapshot, nextSnapshot)) {
+  const snapshotChanged = !is(prevSnapshot, nextSnapshot);
+  if (snapshotChanged) {
     hook.memoizedState = nextSnapshot;
     markWorkInProgressReceivedUpdate();
   }
   const inst = hook.queue;
-  return useSyncExternalStore(hook, inst, subscribe, getSnapshot, nextSnapshot);
+
+  updateEffect(subscribeToStore.bind(null, fiber, inst, subscribe), [
+    subscribe,
+  ]);
+
+  // 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 (
+    inst.getSnapshot !== getSnapshot ||
+    snapshotChanged ||
+    // Check if the susbcribe function changed. We can save some memory by
+    // checking whether we scheduled a subscription effect above.
+    (workInProgressHook !== null &&
+      workInProgressHook.memoizedState.tag & HookHasEffect)
+  ) {
+    // TODO: We can move this to the passive phase once we add a pre-commit
+    // consistency check. See the next comment.
+    fiber.flags |= UpdateEffect;
+    pushEffect(
+      HookHasEffect | HookLayout,
+      updateStoreInstance.bind(null, fiber, inst, nextSnapshot, getSnapshot),
+      undefined,
+      null,
+    );
+
+    // TODO: Unless this is a synchronous render, schedule a consistency check.
+    // Right before committing, we will walk the tree and check if any of the
+    // stores were mutated.
+    // pushConsistencyCheck(inst, getSnapshot, nextSnapshot);
+  }
+
+  return nextSnapshot;
 }
 
-function useSyncExternalStore<T>(
-  hook: Hook,
+function updateStoreInstance<T>(
+  fiber: Fiber,
   inst: StoreInstance<T>,
-  subscribe: (() => void) => () => void,
-  getSnapshot: () => T,
   nextSnapshot: T,
-): T {
-  const fiber = currentlyRenderingFiber;
-  const dispatcher = ReactCurrentDispatcher.current;
+  getSnapshot: () => T,
+) {
+  inst.value = nextSnapshot;
+  inst.getSnapshot = 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 = nextSnapshot;
-    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.
-    // TODO: Move the tearing checks to an earlier, pre-commit phase so that the
-    // layout effects always observe a consistent tree.
+  // TODO: Move the tearing checks to an earlier, pre-commit phase so that the
+  // layout effects always observe a consistent tree.
+  if (checkIfSnapshotChanged(inst)) {
+    // Force a re-render.
+    forceStoreRerender(fiber);
+  }
+}
+
+function subscribeToStore(fiber, inst, subscribe) {
+  const handleStoreChange = () => {
+    // The store changed. Check if the snapshot changed since the last time we
+    // read from the store.
     if (checkIfSnapshotChanged(inst)) {
       // Force a re-render.
       forceStoreRerender(fiber);
     }
-  }, [subscribe, nextSnapshot, 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.
-        forceStoreRerender(fiber);
-      }
-    };
-    // 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 nextSnapshot;
+  };
+  // Check for changes right before subscribing. Subsequent changes will be
+  // detected in the subscription handler.
+  // TODO: Once updateStoreInstance is moved to the passive phase, we can rely
+  // on that check instead of checking again here.
+  handleStoreChange();
+  // Subscribe to the store and return a clean-up function.
+  return subscribe(handleStoreChange);
 }
 
 function checkIfSnapshotChanged(inst) {

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

@@ -1256,6 +1256,7 @@ function mountSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
 ): T {
+  const fiber = currentlyRenderingFiber;
   const hook = mountWorkInProgressHook();
   // Read the current snapshot from the store on every render. This breaks the
   // normal rules of React, and only works because store updates are
@@ -1277,13 +1278,37 @@ function mountSyncExternalStore<T>(
     getSnapshot,
   };
   hook.queue = inst;
-  return useSyncExternalStore(hook, inst, subscribe, getSnapshot, nextSnapshot);
+
+  // Schedule an effect to subscribe to the store.
+  mountEffect(subscribeToStore.bind(null, fiber, inst, subscribe), [subscribe]);
+
+  // Schedule an effect to update the mutable instance fields. We will update
+  // this whenever subscribe, getSnapshot, or value changes. Because there's no
+  // clean-up function, and we track the deps correctly, we can call pushEffect
+  // directly, without storing any additional state. For the same reason, we
+  // don't need to set a static flag, either.
+  // TODO: We can move this to the passive phase once we add a pre-commit
+  // consistency check. See the next comment.
+  fiber.flags |= UpdateEffect;
+  pushEffect(
+    HookHasEffect | HookLayout,
+    updateStoreInstance.bind(null, fiber, inst, nextSnapshot, getSnapshot),
+    undefined,
+    null,
+  );
+  // TODO: Unless this is a synchronous render, schedule a consistency check.
+  // Right before committing, we will walk the tree and check if any of the
+  // stores were mutated.
+  // pushConsistencyCheck(inst, getSnapshot, nextSnapshot);
+
+  return nextSnapshot;
 }
 
 function updateSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
 ): T {
+  const fiber = currentlyRenderingFiber;
   const hook = updateWorkInProgressHook();
   // Read the current snapshot from the store on every render. This breaks the
   // normal rules of React, and only works because store updates are
@@ -1300,66 +1325,81 @@ function updateSyncExternalStore<T>(
     }
   }
   const prevSnapshot = hook.memoizedState;
-  if (!is(prevSnapshot, nextSnapshot)) {
+  const snapshotChanged = !is(prevSnapshot, nextSnapshot);
+  if (snapshotChanged) {
     hook.memoizedState = nextSnapshot;
     markWorkInProgressReceivedUpdate();
   }
   const inst = hook.queue;
-  return useSyncExternalStore(hook, inst, subscribe, getSnapshot, nextSnapshot);
+
+  updateEffect(subscribeToStore.bind(null, fiber, inst, subscribe), [
+    subscribe,
+  ]);
+
+  // 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 (
+    inst.getSnapshot !== getSnapshot ||
+    snapshotChanged ||
+    // Check if the susbcribe function changed. We can save some memory by
+    // checking whether we scheduled a subscription effect above.
+    (workInProgressHook !== null &&
+      workInProgressHook.memoizedState.tag & HookHasEffect)
+  ) {
+    // TODO: We can move this to the passive phase once we add a pre-commit
+    // consistency check. See the next comment.
+    fiber.flags |= UpdateEffect;
+    pushEffect(
+      HookHasEffect | HookLayout,
+      updateStoreInstance.bind(null, fiber, inst, nextSnapshot, getSnapshot),
+      undefined,
+      null,
+    );
+
+    // TODO: Unless this is a synchronous render, schedule a consistency check.
+    // Right before committing, we will walk the tree and check if any of the
+    // stores were mutated.
+    // pushConsistencyCheck(inst, getSnapshot, nextSnapshot);
+  }
+
+  return nextSnapshot;
 }
 
-function useSyncExternalStore<T>(
-  hook: Hook,
+function updateStoreInstance<T>(
+  fiber: Fiber,
   inst: StoreInstance<T>,
-  subscribe: (() => void) => () => void,
-  getSnapshot: () => T,
   nextSnapshot: T,
-): T {
-  const fiber = currentlyRenderingFiber;
-  const dispatcher = ReactCurrentDispatcher.current;
+  getSnapshot: () => T,
+) {
+  inst.value = nextSnapshot;
+  inst.getSnapshot = 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 = nextSnapshot;
-    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.
-    // TODO: Move the tearing checks to an earlier, pre-commit phase so that the
-    // layout effects always observe a consistent tree.
+  // TODO: Move the tearing checks to an earlier, pre-commit phase so that the
+  // layout effects always observe a consistent tree.
+  if (checkIfSnapshotChanged(inst)) {
+    // Force a re-render.
+    forceStoreRerender(fiber);
+  }
+}
+
+function subscribeToStore(fiber, inst, subscribe) {
+  const handleStoreChange = () => {
+    // The store changed. Check if the snapshot changed since the last time we
+    // read from the store.
     if (checkIfSnapshotChanged(inst)) {
       // Force a re-render.
       forceStoreRerender(fiber);
     }
-  }, [subscribe, nextSnapshot, 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.
-        forceStoreRerender(fiber);
-      }
-    };
-    // 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 nextSnapshot;
+  };
+  // Check for changes right before subscribing. Subsequent changes will be
+  // detected in the subscription handler.
+  // TODO: Once updateStoreInstance is moved to the passive phase, we can rely
+  // on that check instead of checking again here.
+  handleStoreChange();
+  // Subscribe to the store and return a clean-up function.
+  return subscribe(handleStoreChange);
 }
 
 function checkIfSnapshotChanged(inst) {