Fork updateSyncExternalStore impl in update and rerender

facebook · Nov 8, 2022 · 71e60d0 · 71e60d0
1 parent 1e3e30d
commit 71e60d0
Show file tree

Hide file tree

Showing 2 changed files with 94 additions and 10 deletions.
diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js
@@ -1593,13 +1593,14 @@ function mountSyncExternalStore<T>(
   return nextSnapshot;
 }
 
-function updateSyncExternalStore<T>(
+function updateSyncExternalStoreImpl<T>(
+  hook: Hook,
+  prevSnapshot: T,
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
   getServerSnapshot?: () => 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
   // always synchronous.
@@ -1615,7 +1616,6 @@ function updateSyncExternalStore<T>(
       }
     }
   }
-  const prevSnapshot = (currentHook || hook).memoizedState;
   const snapshotChanged = !is(prevSnapshot, nextSnapshot);
   if (snapshotChanged) {
     hook.memoizedState = nextSnapshot;
@@ -1666,6 +1666,44 @@ function updateSyncExternalStore<T>(
   return nextSnapshot;
 }
 
+function updateSyncExternalStore<T>(
+  subscribe: (() => void) => () => void,
+  getSnapshot: () => T,
+  getServerSnapshot?: () => T,
+): T {
+  const hook = updateWorkInProgressHook();
+  const prevSnapshot = hook.memoizedState;
+  return updateSyncExternalStoreImpl(
+    hook,
+    prevSnapshot,
+    subscribe,
+    getSnapshot,
+    getServerSnapshot,
+  );
+}
+
+function rerenderSyncExternalStore<T>(
+  subscribe: (() => void) => () => void,
+  getSnapshot: () => T,
+  getServerSnapshot?: () => T,
+): T {
+  const hook = updateWorkInProgressHook();
+  const prevSnapshot =
+    currentHook === null
+      ? // This is a rerender during a mount.
+        hook.memoizedState
+      : // This is a rerender during an update.
+        currentHook.memoizedState;
+
+  return updateSyncExternalStoreImpl(
+    hook,
+    prevSnapshot,
+    subscribe,
+    getSnapshot,
+    getServerSnapshot,
+  );
+}
+
 function pushStoreConsistencyCheck<T>(
   fiber: Fiber,
   getSnapshot: () => T,
@@ -2798,7 +2836,7 @@ const HooksDispatcherOnRerender: Dispatcher = {
   useDeferredValue: rerenderDeferredValue,
   useTransition: rerenderTransition,
   useMutableSource: updateMutableSource,
-  useSyncExternalStore: updateSyncExternalStore,
+  useSyncExternalStore: rerenderSyncExternalStore,
   useId: updateId,
 
   unstable_isNewReconciler: enableNewReconciler,
@@ -3443,7 +3481,11 @@ if (__DEV__) {
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       updateHookTypesDev();
-      return updateSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
+      return rerenderSyncExternalStore(
+        subscribe,
+        getSnapshot,
+        getServerSnapshot,
+      );
     },
     useId(): string {
       currentHookNameInDev = 'useId';

diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js
@@ -1593,13 +1593,14 @@ function mountSyncExternalStore<T>(
   return nextSnapshot;
 }
 
-function updateSyncExternalStore<T>(
+function updateSyncExternalStoreImpl<T>(
+  hook: Hook,
+  prevSnapshot: T,
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
   getServerSnapshot?: () => 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
   // always synchronous.
@@ -1615,7 +1616,6 @@ function updateSyncExternalStore<T>(
       }
     }
   }
-  const prevSnapshot = (currentHook || hook).memoizedState;
   const snapshotChanged = !is(prevSnapshot, nextSnapshot);
   if (snapshotChanged) {
     hook.memoizedState = nextSnapshot;
@@ -1666,6 +1666,44 @@ function updateSyncExternalStore<T>(
   return nextSnapshot;
 }
 
+function updateSyncExternalStore<T>(
+  subscribe: (() => void) => () => void,
+  getSnapshot: () => T,
+  getServerSnapshot?: () => T,
+): T {
+  const hook = updateWorkInProgressHook();
+  const prevSnapshot = hook.memoizedState;
+  return updateSyncExternalStoreImpl(
+    hook,
+    prevSnapshot,
+    subscribe,
+    getSnapshot,
+    getServerSnapshot,
+  );
+}
+
+function rerenderSyncExternalStore<T>(
+  subscribe: (() => void) => () => void,
+  getSnapshot: () => T,
+  getServerSnapshot?: () => T,
+): T {
+  const hook = updateWorkInProgressHook();
+  const prevSnapshot =
+    currentHook === null
+      ? // This is a rerender during a mount.
+        hook.memoizedState
+      : // This is a rerender during an update.
+        currentHook.memoizedState;
+
+  return updateSyncExternalStoreImpl(
+    hook,
+    prevSnapshot,
+    subscribe,
+    getSnapshot,
+    getServerSnapshot,
+  );
+}
+
 function pushStoreConsistencyCheck<T>(
   fiber: Fiber,
   getSnapshot: () => T,
@@ -2798,7 +2836,7 @@ const HooksDispatcherOnRerender: Dispatcher = {
   useDeferredValue: rerenderDeferredValue,
   useTransition: rerenderTransition,
   useMutableSource: updateMutableSource,
-  useSyncExternalStore: updateSyncExternalStore,
+  useSyncExternalStore: rerenderSyncExternalStore,
   useId: updateId,
 
   unstable_isNewReconciler: enableNewReconciler,
@@ -3443,7 +3481,11 @@ if (__DEV__) {
     ): T {
       currentHookNameInDev = 'useSyncExternalStore';
       updateHookTypesDev();
-      return updateSyncExternalStore(subscribe, getSnapshot, getServerSnapshot);
+      return rerenderSyncExternalStore(
+        subscribe,
+        getSnapshot,
+        getServerSnapshot,
+      );
     },
     useId(): string {
       currentHookNameInDev = 'useId';