atomWithActions and create?

[beorn]

Hi,

I'm commenting a bit on the Zustand discussion boards (pmndrs/zustand#2830), trying to find a way to get great performance with many components listening to one Zustand store. I'm trying to not only reduce re-renders, but also reducing the number of selector calls that will have to be done for each component, i.e., to have components subscribe on a more granular basis to parts of the state. A solution is of course to use atoms/signals instead, but I like the more centralized model of a Zustand store, so I've been looking for ways to do it without rewriting everything in Jotai or resorting to convoluted workarounds.

Jotai-zustand could of course be a good solution, and I can use atomWithStore and create derived atoms and then useAtom to create a more direct subscription to that atom's value (saving a lot of work done by selector calls).

But then looking at your code base I came across a few interesting functions that are exported, and that could hint at some other solution, namely create and atomWithActions — these seem to create a Zustand-like useStore interface backed by an atom, is that right?

I thought perhaps I could create a hybrid/mutant Zustand store that directly supports deriving atoms from it (or within it) — and while the atom/store they create seems to have a useStore type interface, it doesn't have the getStore/setStore of the vanilla stores, and it doesn't expose the underlying atoms?

What is the purpose of these, and do you have any thoughts on how to achieve the performance of Jotai Atoms (avoid selector calls) while keeping the programming model of Zustand?

[beorn]

I was hoping something like this might be possible, essentially creating a storeAtom with actions and with other actions:

interface State {
  allItems: Item[]
  firstItemAtom: Atom<Item | undefined> | undefined
}
interface Actions {
  selectItems: (items: Item[]) => void
}

const storeAtom = atomWithActions<State, Actions>(
  {
    allItems: [],
    firstItemAtom: undefined, // computed (read-only) - and independently subscribable without selector
  },
  (set, get) => {
    return {
      firstItemAtom: atom<Item | undefined>((_set, _get) => {
        console.log("jotai/computed")
        const items = get().selectedItems
        return items.length > 0 ? items[0] : undefined
      }),
      selectItems(allItems) {
        set({ allItems })
      },
    }
  }
)

... but I cannot get the firstItemAtom atom initialized in a way so it depends on the underlying state atom's state and is updated when it updates. I guess it would be best if the underlying state atom was exposed when the underlying state was created, e.g.,

const storeAtom = atomWithActions<State & Actions>(
  (set, get) => {
    return {
      allItems: [],
      firstItemAtom: atom<Item | undefined>((_set, _get) => {
        console.log("jotai/computed")
        const items = get(this).selectedItems // perhaps bind state atom to this, or make it available as _get()
        return items.length > 0 ? items[0] : undefined
      }),
      selectItems(allItems) {
        set({ allItems })
      }
    }
  }
)

In any case, when using it I guess it would be something like this:

// render optimized - expose sub-state (sub-atoms) directly for direct subscription
function Example() {
  // const { allItems, firstItemAtom, selectItems } = useAtomValue(storeAtom)
  const firstItem = useAtomValue(storeAtom.firstItem)
  return <p>First item {firstItem}</button>
}

[dai-shi]

A solution is of course to use atoms/signals instead, but I like the more centralized model of a Zustand store, so I've been looking for ways to do it without rewriting everything in Jotai or resorting to convoluted workarounds.

That's an interesting requirement.

namely create and atomWithActions — these seem to create a Zustand-like useStore interface backed by an atom, is that right?

Yeah, that's kind of an experiment. kind of trying to prove Jotai includes Zustand capability-wise. ref: https://x.com/dai_shi/status/1833089238482698515

But, if your issue is performance, I don't think it helps, because it's backed by one atom.

What is the purpose of these, and do you have any thoughts on how to achieve the performance of Jotai Atoms (avoid selector calls) while keeping the programming model of Zustand?

Sounds too tough if we want a real compatibility, but if it's just about the programming model, it might be feasible. I haven't thought about it.

[beorn]

Yeah, my main concern isn't compatibility, but performance while preserving the centralized state programming model of Zustand. Using atomWithActions, if I could create sub-atoms that depend on the base state (but that of course can be independently subscribed to), then I could achieve that.

[dai-shi]

Can we start with a very small problem?
Let's assume we have only two number values in a store (and maybe some actions to update those numbers), our goal is atomic subscription, right? It means, if one of the values is changed, the other value doesn't trigger notifying subscribers.

How would it look like with the programming model in your mind?

[beorn]

Yes, that's it. I guess there are many ways to solve this, I was just wondering what the best/recommended way to do it is, while leveraging just Zustand and/or Jotai.

[dai-shi]

With only Jotai, it would be just two separate atoms. If I understand it correctly, that's not the programming model in your mind.

Another idea with Jotai is to define a store atom and derived smaller atoms.

const stateAtom = atom({ a: 0, b: 0 });
const actionsAtom = atom(null, (get, set, action: { type: 'INC_A' } | { type: 'INC_B' }) => {
  if (action.type === 'INC_A') {
    set(stateAtom, (prev) => ({ ...prev, a: prev.a + 1 }));
  }
  if (action.type === 'INC_B') {
    set(stateAtom, (prev) => ({ ...prev, b: prev.b + 1 }));
  }
});

export const useDispatch = () => {
  const dispatch = useSetAtom(actionsAtom);
  return dispatch;
}

const stateAAtom = atom((get) => get(stateAtom).a);
const stateBAtom = atom((get) => get(stateAtom).b);

export const useStateA = () => useAtomValue(stateAAtom);
export const useStateB = () => useAtomValue(stateBAtom);

[beorn]

And if we try to solve this by being closer to the Zustand model with a central store and functions instead of flux/redux actions? I understand that you sometimes really want to define your state in a distributed and tree shakable way (like a Jotai solution would allow for), but very often keeping it all encapsulated into one store/object makes it easier to confidently reason about (more Zustand-ish), and that's important.

I think it'd be something like this all holistically integrated into one easy-to-understand Zustand store with:

• actions
• base state
• computed values (aka derived values, or selectors, or atoms)

And support using the store not only through selectors, but also through direct subscriptions to the computed values, so there's no overhead if the computed values doesn't change.

[dai-shi]

How about this?

import { create } from 'jotai-zustand';

const [useSelector, useAction, useComputed] = create({
  state: { a: 1, b: 2 },
  actions: {
    incA: () => (state) => ({ ...state, a: state.a + 1 }),
    incB: () => (state) => ({ ...state, b: state.b + 1 }),
  },
  computeds: {
    a: (state) => state.a,
    sum: (state) => state.a + state.b,
  },
});

// usage
const a2 = useSelector((state) => state.a * 2);
const incA = useAction('incA');
const sum = useComputed('sum');

[beorn]

I think functionally that would be there, as long as computeds (perhaps even state) could be directly/atomically subscribed to, and ideally computeds could depend on other computeds.

But aesthetically I think it would be better with a smaller API surface in terms of number of hooks, and perhaps avoid using labeled strings when JS properties could be used, e.g.,

import { create } from 'jotai-zustand';

const [useStore, store] = create({
  state: { a: 1, b: 2 },
  actions: {
    incA: () => (state) => ({ ...state, a: state.a + 1 }),
    incB: () => (state) => ({ ...state, b: state.b + 1 }),
  },
  computeds: {
    a: (state) => state.a,
    sum: (state) => state.a + state.b,
    sumDoubled: (state) => 2 * state.sum
  },
});

// usage
const a2 = useStore(store.a2)
const incA = useStore(store.incA)
const sum2x = useStore(store.sumDoubled)
const a2selectorStyle = useStore(state => state.a2)

Or something like that...?

But I guess this may be even nicer:

import { create } from 'jotai-zustand';

const [useStore, store] = create((set, get) => ({
  a: 1,
  b: 2
  get sum() { return get().a + get().b }, // computed
  get dblSum() { return get().sum * 2 }, // computed that depends on computed
  sum3x: get().sum * 3, // or could this way of doing computeds be possible?
  incA: () => (state) => ({ ...state, a: state.a + 1 }),
  incB: () => (state) => ({ ...state, b: state.b + 1 }),
}));

// usage
const a2 = useStore(store.a2)
const incA = useStore(store.incA)
const sum2x = useStore(store.dblSum) // subscription only "fires" if dblSum value changes
const a2selectorStyle = useStore(state => state.a2)

[dai-shi]

perhaps avoid using labeled strings when JS properties could be used

That would be pretty difficult, or extremely tricky.

Why did you close the issue?

[beorn]

Sorry, I didn't mean to close the issue. Pressed the wrong button!

Using labels isn't a big deal, just (I think at least) a slight improvement in DX. How about the other aspects?

[dai-shi]

We need to explicitly separate actions from values (or original state or computeds).

Having a single hook for actions and values is technically possible, but it's not typescript friendly. It's kind of a bad aspect of Zustand, because its API is designed before TypeScript era.

[beorn]

I think explicitly separating out state, actions, computeds still gets the job done. So that would be possible?

In terms of combining them, would it help to make computeds be values too (but tracked and therefore updated), so only actions are functions, and then use typescript utility types to separate the actions from the values/comptueds? E.g.,

type StoreState = {
  a: number
  b: number
  sum: number
  sum3x: number
  incA: () => void
  incB: () => void
}

const [useStore, store] = create<StoreState>((set, get) => ({
  a: 1,
  b: 2,
  sum: get().a + get().b,
  sum3x: get().sum * 3,
  incA: () => (state) => ({ ...state, a: state.a + 1 }),
  incB: () => (state) => ({ ...state, b: state.b + 1 })
}))

// usage
const a2 = useStore(store, "a2")
const incA = useStore(store, "incA")
const sum3x = useStore(store, "sum3x") // subscription only "fires" if sum value changes
const a2selectorStyle = useStore((state) => state.a2)

// separating out types
type FunctionProperties<T> = {
  [K in keyof T]: T[K] extends Function ? K : never
}[keyof T]
type NonFunctionProperties<T> = {
  [K in keyof T]: T[K] extends Function ? never : K
}[keyof T]

type FunctionsOnly<T> = Pick<T, FunctionProperties<T>>
type ValuesOnly<T> = Pick<T, NonFunctionProperties<T>>

type F = FunctionsOnly<StoreState>
// ^ { incA: () => void; incB: () => void }

type V = ValuesOnly<StoreState>
// ^ { a: number; b: number; sum: number; sum3x: number }