[utils] Add infrastructure to support nested stores

[michaldudak]

A couple of changes to Store and ReactStore:

• BREAKING CHANGE: renamed ReactStore's getElementSetter to useStateSetter as it calls React.useCallback internally.
• Added a notifyAll method to the Store. It's useful to notify listeners of a store when a nested store is updated. See the "supports nested stores as state values" test for usage example.
• Added a select method to get a non-reactive store value. It's a shortcut to calling a selector manually.
• Added an observeSelector method to be notified about changes to a particular slice of the state.
• Added support for parametrized selectors in useState and select methods.

[pkg-pr-new]

vite-css-base-ui-example

pnpm add https://pkg.pr.new/mui/base-ui/@base-ui-components/react@3037

pnpm add https://pkg.pr.new/mui/base-ui/@base-ui-components/utils@3037

commit: 2a1abf0

[mui-bot]

Bundle size report

Bundle	Parsed size	Gzip size
@base-ui-components/react	🔺+397B(+0.10%)	🔺+128B(+0.10%)

Details of bundle changes

---

Check out the code infra dashboard for more information about this PR.

[netlify]

✅ Deploy Preview for base-ui ready!

Name	Link
🔨 Latest commit	2a1abf0
🔍 Latest deploy log	https://app.netlify.com/projects/base-ui/deploys/690487eba834e000080d08b3
😎 Deploy Preview	https://deploy-preview-3037--base-ui.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

---

To edit notification comments on pull requests, go to your Netlify project configuration.

[romgrk]

Added writeInterceptors to ReactStore. They fire whenever a value is about to be set and let users update the value or trigger any side effects (like calling an onChange event). This is also useful to register to a nested store's updates (required in [Menu] Keep state in a store #3022).

Could you expand on why this logic is required for the menu store?

[michaldudak]

Menus can be nested within other menus, menubars, etc. Each level of the menu has its own store. In some cases, when a menu is nested, the root menu's state should take precedence over the local state (for example, if the menubar is disabled, the menu within must be disabled as well).

This code allows me to subscribe to the parent store's changes as if they were local state updates (no matter how many levels deep the menu is). Thanks to interceptors, I can also set the state in the parent menus (as in https://github.com/mui/base-ui/pull/3022/files#diff-b9e964ca92468413a6386a3abbb38d86ccc231ff12fa1ca3c3919de0d2fdc443R80-R89)

[romgrk]

I don't like this new concept because it introduces complexity in a piece of code that I'd like to keep as simple as possible. Isn't there another way to solve that problem? Could a React context be used to store the first parent menu's store?

If a value needs to depend on either a parent or local store, I'd rather we subscribe to both stores (possibly with a custom hook to make it ergonomic) rather than add notifyAll(), which feels a bit hacky to me.

[michaldudak]

I can't use React context here, as with detached triggers, the root element can be outside of the parent's React tree:

<Menubar.Root>
  <Menu.Trigger handle={myMenu} />
</Menubar.Root>
<Menu.Root handle={myMenu} />

The Store instance (within the myMenu handle) is the only point of communication between the Trigger and Root.

I can create a Base UI-specific subclass of ReactStore that will handle these cases instead if you don't want it in the shared code. But the interceptors I introduced here can also be used to trigger change callbacks automatically (which could be useful for controlled props), so perhaps it's worth having this in the shared code. For raw performance without any bells and whistles, pure Store can always be used.

[romgrk]

Could we solve that use case with existing store effects instead of write interceptors? From the use-case in your link above, essentially what I see is "if this value changes, do that", which is already covered by existing API. And for values that depend on both local and parent store, because the parent store is also in the state, you can write selectors like this:

const disabledSelector = createSelector(
  state => state.parent.store?.state.disabled ?? state.disabled
)

And we could add just an API to link stores' reactivity together, something like this, but I think notifyAll() doesn't need to change its state ref, it just needs to call the listeners. With selectors like the one above, it would work.

[romgrk]

If I subscribe to allowMouseEnter changes and update the parent store after it changes, wouldn't it trigger another render

I don't think it would. The store's update notifies all its listeners, one of which would update the parent store, which would notify its listeners. Both of those may schedule a re-render in React, but that doesn't happen synchronously (afaiu), both of those updates happen within a root Store.setState() call that doesn't yield to React. The Store is currently designed to run listeners (& effects) synchronously.

So yes I would just add a subscription to update the parent when allowMouseEnter changes.

[michaldudak]

If I understand it correctly, there's no way currently to be notified if just a specific part of the state changes, right? I'd have to keep the previous value of the parent myself and compare it with the new state coming from subscribe callback.

We could create something like this method:

observe(selector, onChange) {
  let prevValue = selector(this.state);

  return store.subscribe((nextState) => {
    const nextValue = selector(nextState);
    if (!Object.is(prevValue, nextValue)) {
      prevValue = nextValue;
      onChange(nextValue);
    }
  });
}

[michaldudak]

Actually, this won't work like that - if we observe the result of selectors and the selector gets a value from a nested store, we won't ever fire the listener function as only the local state gets updated:

const selectors = {
  count: (state: ChildState) => state.parent?.state.count ?? state.count,
  parent: (state: ChildState) => state.parent,
};

// ...

store.observe('count', (newCount) => {
  store.parent.set('count', newCount);
});

// ...

store.set('count', 42) // -> this won't cause the store.parent.set to execute

We'd have to have another function that observes the values of state fields, not selector return functions (https://github.com/mui/base-ui/blob/09904b55fddef4655f0923b733253a84d0d28d99/packages/utils/src/store/ReactStore.ts).

Alternatively, we can keep observing selectors' return values and set up two-way data syncing, as in https://github.com/mui/base-ui/blob/a713ccc9ff558864b7aaecb854c7da238d46ac26/packages/utils/src/store/ReactStore.ts and

base-ui/packages/utils/src/store/ReactStore.test.tsx

Lines 253 to 274 in a713ccc

	let unsubscribeParentHandler: () => void;
	const onParentUpdated = (
	newParent: ReactStore<ParentState> | undefined,
	_: ReactStore<ParentState> | undefined,
	store: ReactStore<ChildState, any, any>,
	) => {
	if (!newParent) {
	unsubscribeParentHandler?.();
	return;
	}
	unsubscribeParentHandler = newParent.observe('count', (newCount) => {
	store.set('count', newCount);
	});
	};
	const onCountUpdated = (
	newCount: number,
	_: number,
	store: ReactStore<ChildState, any, any>,
	) => {
	store.state.parent?.set('count', newCount);
	};

. I like it less, though, as the setup is less straightforward and you have to make sure selectors don't read from nested stores.

[romgrk]

We'd have to have another function that observes the values of state fields, not selector return functions

That's just because of the name collision in your selectors. The one you named count does more than what its name implies. You could add another selector, e.g. localCount: state => state.count, and then observing that one would make that snippet of code work properly.

I have some hesitations about calling selectors via strings, because that syntax may cause confusion with other APIs like .set() where the string refers to an actual field rather than a selector. In this part of your example, 'count' refers to different things:

store.observe('count', (newCount) => {
  store.parent.set('count', newCount);
});

[michaldudak]

It's because when parent is set, I want to completely ignore the local property (that's the case in Base UI). The caller doesn't really need to know if the value comes from the local state or another store. The localCount selector would exist only to set up a listener, and it would be invalid to use it from application code.

I have some hesitations about calling selectors via strings, because that syntax may cause confusion with other APIs like .set()

I can see that, but on the other hand, it's more convenient to call store.useState('foo') than store.useState(selectors.foo) (plus it also prevents developers from creating ad-hoc selectors such as store.useState((state) => state.foo) as I mentioned in another discussion).

One way to make it clearer would be to rename methods that operate on selectors to contain the "selector" word: useState -> useSelector, observe -> observeSelector

[romgrk]

Is the prototype call still necessary?

[michaldudak]

Yes - it prevents an unnecessary call when ReactStore.update or ReactStore.set is called.
Without it the control flow looks like this (for update):
ReactStore.update -> Store.update -> ReactStore.setState -> Store.setState.
Introducing it eliminates the ReactStore.setState call.

[romgrk]

Let's leave it but the controlled props thingy will need a refactoring/performance pass. The check for controlled props is made 2 times when calling ReactStore.update. Originally I aimed to make .setState the only method that mutates .state, the other methods being built on top of it. Also I've been wondering, is it really necessary to add a defensive check for controlled props? Would it be possible to add that check only in dev?

[michaldudak]

The check for controlled props is made 2 times when calling ReactStore.update|

How so? ReactStore.update calls Store.update that in turn calls Store.setState.
Store.setState is the only place directly mutating state. All the other methods call it.

Also I've been wondering, is it really necessary to add a defensive check for controlled props? Would it be possible to add that check only in dev?

It wouldn't work when setState is called from the outside. But perhaps we can make it protected? Do you have a need to completely replace the state?

Side note - this doesn't touch this PR directly. I'd appreciate it if we focused on what's being changed here so I can proceed with other PRs and continue the discussion in parallel.

[romgrk]

How so?

Here and here. But I guess with this change the 2 checks will be reduced to just one in .setState().

It wouldn't work when setState is called from the outside.

But the check is only there in case someones inadvertedly tries to update a controlled value right? Could we consider moving that check to typescript by declaring controlled props and removing them from all methods' typings except useControlledProp?

[michaldudak]

Could we consider moving that check to typescript by declaring controlled props and removing them from all methods' typings except useControlledProp?

Not really, as we don't know which props are really controlled. We know which are "controllable," but it's impossible to know at the type level whether callers used a controlled or uncontrolled variant.

[romgrk]

I just realized that if we expose APIs that act on state keys as *State, the setState method becomes confusing. Tbh I'm not very enthusiast about having methods like this, I'd rather we keep one way of doing things: selectors. I would also prefer we invoke them without passing through strings - but if I have to pick one I'd rather have "selectors as function or string" over "selectors and state keys as strings". The "selectors as function or string" also lets you have selectors that aren't exposed in the .selectors interface, e.g. .observe(internalSelectors.count).

[michaldudak]

OK, having both functions and keys seems like a good compromise

[romgrk]

If we agree to go with selectors as the only way to read the state, then we can rename this to just .observe().