Sam/use selector hacking

package.json

@@ -42,7 +42,7 @@
     "format": "prettier --write \"{src,test}/**/*.{js,ts,tsx}\" \"docs/**/*.md\"",
     "lint": "eslint src test",
     "lint:fix": "eslint src test --fix",
-    "prepare": "yarn clean && yarn build",
+    "prepare": "yarn install && yarn clean && yarn build",
     "pretest": "yarn lint",
     "test": "vitest --run --typecheck",
     "test:watch": "vitest --watch",
@@ -62,10 +62,6 @@
       "optional": true
     }
   },
-  "dependencies": {
-    "@types/use-sync-external-store": "^0.0.6",
-    "use-sync-external-store": "^1.2.2"
-  },
   "devDependencies": {
     "@babel/cli": "^7.24.7",
     "@babel/core": "^7.24.7",

src/components/Context.ts

@@ -7,7 +7,7 @@ import type { ProviderProps } from './Provider'
 export interface ReactReduxContextValue<
   SS = any,
   A extends Action<string> = UnknownAction,
-> extends Pick<ProviderProps, 'stabilityCheck' | 'identityFunctionCheck'> {
+> extends Pick<ProviderProps, 'stabilityCheck' | 'identityFunctionCheck' | 'traceFactory'> {
   store: Store<SS, A>
   subscription: Subscription
   getServerState?: () => SS

src/components/Provider.tsx

@@ -6,6 +6,7 @@ import { createSubscription } from '../utils/Subscription'
 import { useIsomorphicLayoutEffect } from '../utils/useIsomorphicLayoutEffect'
 import type { ReactReduxContextValue } from './Context'
 import { ReactReduxContext } from './Context'
+import type { TraceFactory } from '@internal/utils/trace'
 
 export interface ProviderProps<
   A extends Action<string> = UnknownAction,
@@ -51,6 +52,11 @@ export interface ProviderProps<
    */
   identityFunctionCheck?: DevModeCheckFrequency
 
+  /** Tracing allows capturing interactions such as selector runs and
+   * memoization hits.
+   */
+  traceFactory?: TraceFactory
+
   children: ReactNode
 }
 
@@ -61,6 +67,7 @@ function Provider<A extends Action<string> = UnknownAction, S = unknown>({
   serverState,
   stabilityCheck = 'once',
   identityFunctionCheck = 'once',
+  traceFactory,
 }: ProviderProps<A, S>) {
   const contextValue = React.useMemo(() => {
     const subscription = createSubscription(store)
@@ -70,8 +77,9 @@ function Provider<A extends Action<string> = UnknownAction, S = unknown>({
       getServerState: serverState ? () => serverState : undefined,
       stabilityCheck,
       identityFunctionCheck,
+      traceFactory,
     }
-  }, [store, serverState, stabilityCheck, identityFunctionCheck])
+  }, [store, serverState, stabilityCheck, identityFunctionCheck, traceFactory])
 
   const previousState = React.useMemo(() => store.getState(), [store])
 

src/components/connect.tsx

@@ -39,14 +39,6 @@ import type {
 } from './Context'
 import { ReactReduxContext } from './Context'
 
-import type { uSES } from '../utils/useSyncExternalStore'
-import { notInitialized } from '../utils/useSyncExternalStore'
-
-let useSyncExternalStore = notInitialized as uSES
-export const initializeConnect = (fn: uSES) => {
-  useSyncExternalStore = fn
-}
-
 // Define some constant arrays just to avoid re-creating these
 const EMPTY_ARRAY: [unknown, number] = [null, 0]
 const NO_SUBSCRIPTION_ARRAY = [null, null]
@@ -726,7 +718,7 @@ function connect<
       let actualChildProps: Record<string, unknown>
 
       try {
-        actualChildProps = useSyncExternalStore(
+        actualChildProps = React.useSyncExternalStore(
           // TODO We're passing through a big wrapper that does a bunch of extra side effects besides subscribing
           subscribeForReact,
           // TODO This is incredibly hacky. We've already processed the store update and calculated new child props,

src/exports.ts

@@ -40,6 +40,8 @@ export type { UseStore } from './hooks/useStore'
 
 export type { Subscription } from './utils/Subscription'
 
+export type { Trace } from "./utils/trace"
+
 export * from './types'
 
 /**

src/hooks/useSelector.ts

@@ -4,12 +4,12 @@ import { React } from '../utils/react'
 import type { ReactReduxContextValue } from '../components/Context'
 import { ReactReduxContext } from '../components/Context'
 import type { EqualityFn, NoInfer } from '../types'
-import type { uSESWS } from '../utils/useSyncExternalStore'
-import { notInitialized } from '../utils/useSyncExternalStore'
 import {
   createReduxContextHook,
   useReduxContext as useDefaultReduxContext,
 } from './useReduxContext'
+import type { Trace } from '@internal/exports'
+import { useEffect } from 'react'
 
 /**
  * The frequency of development mode checks.
@@ -93,6 +93,7 @@ export interface UseSelector<StateType = unknown> {
   <TState extends StateType = StateType, Selected = unknown>(
     selector: (state: TState) => Selected,
     equalityFnOrOptions?: EqualityFn<Selected> | UseSelectorOptions<Selected>,
+    name?: string,
   ): Selected
 
   /**
@@ -118,11 +119,6 @@ export interface UseSelector<StateType = unknown> {
   >() => UseSelector<OverrideStateType>
 }
 
-let useSyncExternalStoreWithSelector = notInitialized as uSESWS
-export const initializeUseSelector = (fn: uSESWS) => {
-  useSyncExternalStoreWithSelector = fn
-}
-
 const refEquality: EqualityFn<any> = (a, b) => a === b
 
 /**
@@ -147,6 +143,7 @@ export function createSelectorHook(
     equalityFnOrOptions:
       | EqualityFn<NoInfer<Selected>>
       | UseSelectorOptions<NoInfer<Selected>> = {},
+    name?: string,
   ): Selected => {
     const { equalityFn = refEquality, devModeChecks = {} } =
       typeof equalityFnOrOptions === 'function'
@@ -172,8 +169,33 @@ export function createSelectorHook(
       getServerState,
       stabilityCheck,
       identityFunctionCheck,
+      traceFactory,
     } = useReduxContext()
 
+    const trace = React.useRef<{ 
+      trace: Trace, 
+      selector: (state: TState) => Selected,
+      name?: string,
+    }>();
+
+    if (trace.current && 
+         (trace.current.selector !== selector || 
+          trace.current.name !== name)) {
+      trace.current?.trace.onSelectorUnmount?.();
+      trace.current = undefined;
+    }
+
+    if (!trace.current && traceFactory) {
+      trace.current = { trace: traceFactory(name), selector, name };
+    }
+
+    useEffect(() => {
+      return () => {
+        trace.current?.trace.onSelectorUnmount?.();
+        trace.current = undefined;
+      }
+    }, []);
+
     const firstRun = React.useRef(true)
 
     const wrappedSelector = React.useCallback<typeof selector>(
@@ -252,6 +274,7 @@ export function createSelectorHook(
       getServerState || store.getState,
       wrappedSelector,
       equalityFn,
+      trace.current?.trace,
     )
 
     React.useDebugValue(selectedState)
@@ -290,3 +313,154 @@ export function createSelectorHook(
  * }
  */
 export const useSelector = /*#__PURE__*/ createSelectorHook()
+
+
+
+
+export type uSES = typeof React.useSyncExternalStore
+export type uSESWS = typeof useSyncExternalStoreWithSelector
+
+
+// Copied from use-sync-external-store
+function useSyncExternalStoreWithSelector<Snapshot, Selection>(
+  subscribe: (onStoreChange: () => void) => () => void,
+  getSnapshot: () => Snapshot,
+  getServerSnapshot: undefined | null | (() => Snapshot),
+  selector: (snapshot: Snapshot) => Selection,
+  isEqual?: (a: Selection, b: Selection) => boolean,
+  trace?: Trace
+): Selection {
+  // Use this to track the rendered snapshot.
+  type refType = { hasValue: boolean, value: Selection | null };
+  const instRef = React.useRef<refType|null>(null);
+  let inst: refType | null = null;
+
+  if (instRef.current === null) {
+    inst = {
+      hasValue: false,
+      value: null
+    };
+    instRef.current = inst;
+  } else {
+    inst = instRef.current;
+  }
+
+  const subscribeWrapper = React.useCallback((onStoreChange: () => void): () => void => {
+    trace?.onSubscribe?.();
+    const cleanup = subscribe(() => {
+      trace?.onStoreChange?.();
+      onStoreChange();
+    });
+
+    return () => {
+      trace?.onSubscribeCleanup?.();
+      cleanup();
+    }
+  }, [trace]);
+
+  const [getSelection, getServerSelection] = React.useMemo(function () {
+    // Track the memoized state using closure variables that are local to this
+    // memoized instance of a getSnapshot function. Intentionally not using a
+    // useRef hook, because that state would be shared across all concurrent
+    // copies of the hook/component.
+    let hasMemo = false;
+    let memoizedSnapshot: Snapshot | null = null;
+    let memoizedSelection: Selection | null = null;
+
+    const memoizedSelector = function (nextSnapshot: Snapshot) {
+      if (!hasMemo) {
+        // The first time the hook is called, there is no memoized result.
+        hasMemo = true;
+        memoizedSnapshot = nextSnapshot;
+
+        trace?.onSelectorCallStart?.();
+        const _nextSelection = selector(nextSnapshot);
+        trace?.onSelectorCallEnd?.();
+
+        if (isEqual !== undefined) {
+          // Even if the selector has changed, the currently rendered selection
+          // may be equal to the new selection. We should attempt to reuse the
+          // current value if possible, to preserve downstream memoizations.
+          if (inst?.hasValue) {
+            const currentSelection = inst.value;
+            const eq = isEqual(currentSelection!, _nextSelection);
+            trace?.onIsEqualCall?.(eq);
+
+            if (eq) {
+              memoizedSelection = currentSelection;
+              return currentSelection;
+            }
+          }
+        }
+
+        memoizedSelection = _nextSelection;
+        return _nextSelection;
+      } // We may be able to reuse the previous invocation's result.
+
+
+      // We may be able to reuse the previous invocation's result.
+      const prevSnapshot = memoizedSnapshot;
+      const prevSelection = memoizedSelection;
+
+      const isEq = Object.is(prevSnapshot, nextSnapshot);
+
+      trace?.onObjectIsEqualCall?.(isEq);
+
+      if (isEq) {
+        // The snapshot is the same as last time. Reuse the previous selection.
+        return prevSelection;
+      } // The snapshot has changed, so we need to compute a new selection.
+
+
+      // The snapshot has changed, so we need to compute a new selection.
+      trace?.onSelectorCallStart?.();
+      const nextSelection = selector(nextSnapshot); // If a custom isEqual function is provided, use that to check if the data
+      trace?.onSelectorCallEnd?.();
+      // has changed. If it hasn't, return the previous selection. That signals
+      // to React that the selections are conceptually equal, and we can bail
+      // out of rendering.
+
+      // If a custom isEqual function is provided, use that to check if the data
+      // has changed. If it hasn't, return the previous selection. That signals
+      // to React that the selections are conceptually equal, and we can bail
+      // out of rendering.
+
+      if (isEqual !== undefined) {
+        const eq = isEqual(prevSelection!, nextSelection);
+        trace?.onIsEqualCall?.(eq);
+        if (eq)
+          return prevSelection;
+      }
+
+      memoizedSnapshot = nextSnapshot;
+      memoizedSelection = nextSelection;
+      return nextSelection;
+    }; // Assigning this to a constant so that Flow knows it can't change.
+
+    // Assigning this to a constant so that Flow knows it can't change.
+    const maybeGetServerSnapshot = getServerSnapshot === undefined ? null : getServerSnapshot;
+
+    const getSnapshotWithSelector = function () {
+      trace?.onGetSnapshot?.();
+      return memoizedSelector(getSnapshot());
+    };
+
+    const getServerSnapshotWithSelector = maybeGetServerSnapshot === null ? undefined : function () {
+      return memoizedSelector(maybeGetServerSnapshot());
+    };
+    return [getSnapshotWithSelector, getServerSnapshotWithSelector];
+  }, [getSnapshot, getServerSnapshot, selector, isEqual, trace]);
+
+  const value = React.useSyncExternalStore(subscribeWrapper, getSelection!, getServerSelection);
+
+  React.useEffect(function () {
+    if (inst) {
+      inst.hasValue = true;
+      inst.value = value;
+    }
+  }, [value]);
+
+  React.useDebugValue(value);
+
+  return value!;
+}

src/index.ts

@@ -1,15 +1,3 @@
-// The primary entry point assumes we are working with React 18, and thus have
-// useSyncExternalStore available. We can import that directly from React itself.
-// The useSyncExternalStoreWithSelector has to be imported, but we can use the
-// non-shim version. This shaves off the byte size of the shim.
-
-import * as React from 'react'
-import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/with-selector.js'
-
-import { initializeUseSelector } from './hooks/useSelector'
-import { initializeConnect } from './components/connect'
-
-initializeUseSelector(useSyncExternalStoreWithSelector)
-initializeConnect(React.useSyncExternalStore)
+// Assume React >= 18 for useSyncExternalStore
 
 export * from './exports'

src/utils/trace.ts

@@ -0,0 +1,19 @@
+export interface Trace {
+  onSelectorUnmount?: () => void
+
+  onSelectorCallStart?: () => void
+  onSelectorCallEnd?: () => void
+
+  onIsEqualCall?: (equal: boolean) => void
+  onObjectIsEqualCall?: (equal: boolean) => void
+
+  onSubscribe?: () => void
+  onSubscribeCleanup?: () => void
+  onStoreChange?: () => void
+
+  onGetSnapshot?: () => void
+}
+
+export interface TraceFactory {
+  (name?: string): Trace
+}