react-tracked and Zustand v5, no longer working for my use case

[zeorin]

I've created my own custom useTrackedStore hook:

import { createTrackedSelector } from 'react-tracked';
import { useStore } from 'zustand';
import type { StoreApi } from 'zustand/vanilla';

type ReadonlyStoreApi<T> = Pick<
  StoreApi<T>,
  'getState' | 'getInitialState' | 'subscribe'
>;

function identity<T>(x: T): T;
function identity(): undefined;
function identity(x?: unknown) {
  return x;
}

export const useTrackedStore = <const State, const Selected = State>(
  store: ReadonlyStoreApi<State>,
  selector: (state: State) => Selected = identity
): Selected => {
  const useSelector = <const Tracked>(
    tracker: (state: Selected) => Tracked
  ): Tracked => useStore(store, (state) => tracker(selector(state)));
  return createTrackedSelector(useSelector)();
};

I created it so that I can use this with a store that is passed via React Context. It works really well with v4, but with v5 I get a warning about getSnapShot not returning a stable reference, and then an infinite loop.

I have read the migration guide, which suggests using useShallow, but that doesn't seem to help at all, at least, not when I use it like this:

const state = useTrackedStore(store, useShallow(selector));

[dai-shi]

  return createTrackedSelector(useSelector)();

I thought that that would never work. It seems to create new reference. It might have worked by chance.
Can you use createTrackedSelector globally?

[zeorin]

It works because the actual core React hooks under the hood are still called according to the rules of hooks.

I cannot use createTrackedSelector globally for this use case because I'm not using it for global state. I'm passing the store down to my useTrackedStore call sites via React Context.

See my comment on @TkDodo's Zustand and React Context article for more info on how and why I'm using it. TL;DR: in a "legacy" project I'm using Zustand + React Tracked to dramatically improve the performance of Formik by wrapping its state in a store and using the tracked selector (incidentally, I have since gotten the render time per keystroke down to <0.1ms).

[zeorin]

My useTrackedStore hook will indeed return a new reference each time, but that is the same for the global version from React Tracked's docs. The returned object is a proxy that is created anew on each render, and I'm not expecting that to be referentially stable.

[dai-shi]

The global version doesn't create a hook inside render, does it?

[zeorin]

const useFoo = () => {
	// … React hooks called here according to rules of hooks
}

const Component = () => {
	const foo = useFoo();
	return null
}

is equivalent to

const useFoo = () => {
	const useBar = () => {
		// … React hooks called here according to rules of hooks
	}
	return useBar()
}

const Component = () => {
	const foo = useFoo();
	return null
}

which is equivalent to

const useFoo = () => {
	const useBar = () => {
		// … React hooks called here according to rules of hooks
	}
	return useBar
}

const Component = () => {
	const foo = useFoo()();
	return null
}

The React Compiler linter doesn't complain about this and it actually works with the compiler, too.

I've sometimes used this if I needed to both provide a type parameter and also infer some. I actually cribbed that trick from Zustand. I was surprised that it worked with hooks, but when I thought more deeply about it, as far as React knows, during component rendering, the same core hooks are called in the same order every render. How they come to be called is not that important as long as they're called in the same order every time.

[dai-shi]

In that case, it might be equivalent.

Back to your original problem, if it's equivalent, I hope you can create a very small reproduction with "global version" which should help getting helps.

[mooalot]

Can you provide an example of it not working? Mine works just fine.

[zeorin]

I fixed it.

TL;DR: instead of applying the tracking proxy to the value returned by the selector, I apply it to the value returned by useStoreWithEqualityFn, and use proxy-compare's isChanged as part of the equality function.

import { useCallback, useMemo } from 'react';

import { createProxy, isChanged } from 'proxy-compare';
import { useStoreWithEqualityFn } from 'zustand/traditional';
import type { StoreApi } from 'zustand/vanilla';

type ReadonlyStoreApi<T> = Pick<
  StoreApi<T>,
  'getState' | 'getInitialState' | 'subscribe'
>;

function identity<T>(x: T): T;
function identity(): undefined;
function identity(x?: unknown) {
  return x;
}

export const useTrackedStore = <const State, const Selected = State>(
  api: ReadonlyStoreApi<State>,
  selector: (state: State) => Selected = identity,
  isEqual: NoInfer<(next: Selected, prev: Selected) => boolean> = Object.is
): Selected => {
  const affected = useMemo(() => new WeakMap(), []);

  const isTrackedEqual = useCallback(
    (next: Selected, prev: Selected) =>
      isEqual(next, prev) &&
      !isChanged(
        prev,
        next,
        affected,
        new WeakMap(),
        isEqual as (a: unknown, b: unknown) => boolean
      ),
    [affected, isEqual]
  );

  const slice = useStoreWithEqualityFn(api, selector, isTrackedEqual);

  const proxyCache = useMemo(() => new WeakMap(), []);

  return createProxy(slice, affected, proxyCache);
};

This also has the benefit of making unwrapping of the proxy much simpler (e.g. for useDebugValue), as I can just use proxy-compare's getUntracked.

(Note that in practice I've written my own useStoreWithEqualityFn because Zustand's one is still using useSyncExternalStoreWithSelector from use-sync-external-store/shim/with-selector and not use-sync-external-store/with-selector, it calls useDebugValue directly on the return value, I'm not sure whether that affects affected in any way, and it doesn't memoize the selectors.)

[zeorin]

@dai-shi would this be a welcome addition to Zustand or React Tracked? Happy to submit a PR.

[dai-shi]

I think it's an advanced usage. Not sure if it's worth a guide in Zustand docs.
Maybe it's good to add a guide in proxy-compare docs (new folder).

[mooalot]

I implemented something like this improve performance just as zeorin did. My overall performance (per state update) increased by about 20x. I think some addition of this to zustand would be a great idea.

[dai-shi]

If you think it's valuable for other users, would you like to create a third-party package? It would be nice if there were a single solution that covers both use cases.