atomWithActions and create?

[beorn]

Hi,

I'm commenting a bit on the Zustand discussion boards (https://github.com/pmndrs/zustand/discussions/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 }

[dai-shi]

It's not about TypeScript issue. Combining actions would complicate the implementation. Combining computeds and values would be easier. Computeds must be "functions".

const [useValue, useAction] = create({
  a: 1,
  b: 2,
  sum: (state) => state.a + state.b, // But, this doesn't work with Type inference
}, {
  incA: () => (state) => ({ ...state, a: state.a + 1 }),
  incB: () => (state) => ({ ...state, b: state.b + 1 }),
});

Not very pretty?

[dai-shi]

If we were to design a new API now, it should be type inferable, I believe.

check out https://stately.ai/docs/xstate-store for reference.

[dai-shi]

const useStore = create({ a: 1, b: 2 })
  .computeds({ sum: (state) => state.a + state.b })
  .computeds({ sum2: (state) => state.sum * 2 })
  .actions({
    addA: (n: number) => (state) => ({ ...state, a: state.a + n })
  });

// usage
const a = useStore('a')
const addA = useStore('addA') // I need to think about feasibility
const a2 = useStore((state) => state.a * 2) // this isn't very good in performance, prefer computeds

[beorn]

Looks very promising! Would it be possible for sum2 & sum to be in the same computeds definition/call? And would it enable fine-grained tracking of computeds dependencies to only re-generate if dependencies change?

I agree complete type inference is a huge plus.

[dai-shi]

Would it be possible for sum2 & sum to be in the same computeds definition/call?

I don't think type inference works with the same call.

Using string keys, it should be fine-grained.

On second thought, returning a hook directly with the "chain" can be a bit troublesome. How about this?

import { createStore, useStoreValue, useStoreAction } from 'jotai-zustand';

const { values, actions } = createStore({ a: 1, b: 2 })
  .computeds({
    sum: (state) => state.a + state.b,
   })
  .computeds({
    sum2: (state) => state.sum * 2,
  })
  .actions({
    addA: (n: number) => (state) => ({ ...state, a: state.a + n }),
  });

// usage
const a = useStoreValue(values.a)
const sum2 = useStoreValue(values.sum2)
const addA = useStoreAction(actions.addA)

[beorn]

I think it's solves the problems (type inferred, supports computeds, direct subscriptions, central store).

But I do think it would at least seem less complex if the usage and definition APIs were a bit flatter, like in Zustand:

// definition
const store = createStore({ a: 1, b: 2 })
  .computeds({ // could these be part of the state definition? (function in state = computed)
    sum: (state) => state.a + state.b,
  })
  .computeds({
    sum2: (state) => state.sum * 2,
  })
  .actions({ // could .sum2 and other actions (.addB()) be available in the state here?
    addA: (n: number) => (state) => ({ ...state, a: state.sum2 + n }),
  })

// usage
const a = useStore(store.a)
const sum2 = useStore(store.sum2)
const addA = useStore(store.addA)

I guess you could also use this fluid syntax to create sub-stores that could depend on the upstream/parent stores.

[dai-shi]

// could these be part of the state definition? (function in state = computed)

I thought it's impossible to infer types, but as we don't use createState function like Zustand, I'm not sure. We need to try it.

// could .sum2 and other actions (.addB()) be available in the state here?

It's available for read, but write should be prohibited because otherwise, it breaks computeds.

// if type inference works:
const store = createStore({
  a: 1,
  b: 2,
  sum: (state) => state.a + state.b,
  addA: (n: number) => (state) => ({ a: state.sum + n })
});

// now the difficulty is sum and addA are both functions. So, even if we can solve with JS, it's not TS friendly.

I think original approach might be better:

const store = createStore(
  // values
  {
    a: 1,
    b: 2,
    sum: (state) => state.a + state.b,
  },
  // actions
   {
    addA: (n: number) => (state) => ({ a: state.sum + n })
  },
);

[beorn]

This does seem to work to some extent (see gist) - but it relies on proxies etc - maybe not very Zustand-ish (more Valtio-ish) :-/ And doesn't support selector access, and probably doesn't handle deeper structures.

// full impl https://gist.github.com/beorn/2726d3d3ffe661bb1c3666e2da51c787
// Create the store
const store = createStore({
  a: 1,
  b: 2,
  get sum() {
    return this.a + this.b
  },
  get sum2x() {
    return this.sum * 2
  },
  addA(n: number) {
    this.a = this.a + this.sum + n
  },
})

// In a React component
export default function MyComponent() {
  const a = useStoreValue(store.a)
  const b = useStoreValue(store.b)
  const sum = useStoreValue(store.sum)
  const sum2x = useStoreValue(store.sum2x)
  const addA = useStoreValue(store.addA)

  return (
    <div>
      <h1>Brainstorming 9</h1>
      <div>a: {a}</div>
      <div>b: {b}</div>
      <div>sum: {sum}</div>
      <div>sum2x: {sum2x}</div>
      <button onClick={() => addA(5)}>Add to A</button>
    </div>
  )
}```

[dai-shi]

maybe not very Zustand-ish (more Valtio-ish)

Yeah, we should avoid Proxies in this project.

[dai-shi]

btw, your last example resembles with https://github.com/zustandjs/zustand-valtio

[dai-shi]

If we can use getter and this, I'd go with something like:

import { createStore, useStore } from 'jotai-zustand';

const store = createStore({
  a: 1,
  b: 2,
  get sum() {
    return this.a + this.b
  },
  get sum2x() {
    return this.sum * 2
  },
  addA: (n: number) => (state) => ({ a: state.sum + n }), // not sure if `state` type can be inferred
})

// In a React component
export default function MyComponent() {
  const a = useStore(store.a)
  const sum2x = useStore(store.sum2x)
  const addA = useStore(store.addA)

  return (
    <div>
      <div>a: {a}</div>
      <div>sum2x: {sum2x}</div>
      <button onClick={() => addA(5)}>Add to A</button>
    </div>
  )
}```

[dai-shi]

Or,

  addA(n: number) {
    return { a: this.sum + n };
  }

But, I have never seen this style before.

[beorn]

Ah, yes, looks a lot like zustand-valtio :) I agree it'd be best to avoid magic like proxies.

I don't think state can be inferred if you do addA: (n: number) => (state) => ..., but I was able to get it working with your last example:

function MyComponent() {
  const a = useStoreValue(store.a)
  const b = useStoreValue(store.b)
  const sum = useStoreValue(store.sum)
  const sum2x = useStoreValue(store.sum2x)
  const addA = useStoreValue(store.addA)

  return (
    <div>
      <h1>Brainstorming 11</h1>
      <div>a: {a}</div>
      <div>b: {b}</div>
      <div>sum: {sum}</div>
      <div>sum2x: {sum2x}</div>
      <button onClick={() => addA(5)}>Add to A</button>
    </div>
  )
}

const store = createStore({
  a: 1,
  b: 2,
  addA(n: number) {
    return { a: this.a + n }
  },
  get sum(): number {
    return this.a + this.b
  },
  get sum2x(): number {
    return this.sum * 2
  },
})

See full gist. It doesn't use proxies, but it does track dependencies through property access.

[dai-shi]

store.addA(5) doesn't work with Jotai. It has to be useStoreValue(store.addA) or useStoreAction(store.addA).

I still wonder if relying on this is the best fit with jotai-zustand. We only use this because it allows type inference, right?

[beorn]

Apologies, yes, should be useStoreValue on the action too. I updated the gist and the code above.

And, yes, using this allows for type inference, but I agree it doesn't seem to fit well. I had problems getting typescript to infer types for sum: (state) => state.a + state.b — but maybe there's a way.

This at least works:

type StoreShape = {
  a: number
  b: number
  addA(n: number): Partial<StoreShape>
  sum: { get: () => number }
  sum2x: { get: () => number }
}

export const store = createStore<StoreShape>({
  a: 1,
  b: 2,
  addA: (n: number) => (state) => ({ a: state.a + n }),
  sum: { get: () => (state) => state.a + state.b },
  sum2x: { get: () => (state) => state.sum * 2 },
})

// Component should now have correct types
export default function MyComponent() {
  const a = useStoreValue(store.a) // number
  const b = useStoreValue(store.b) // number
  const sum = useStoreValue(store.sum) // number
  const sum2x = useStoreValue(store.sum2x) // number
  const addA = useStoreValue(store.addA) // (n: number) => void

  return (
    <div>
      <h1>Brainstorming 12</h1>
      <div>a: {a}</div>
      <div>b: {b}</div>
      <div>sum: {sum}</div>
      <div>sum2x: {sum2x}</div>
      <button onClick={() => addA(5)}>Add to A</button>
    </div>
  )
}

See gist

[dai-shi]

Yeah, if we give up type inference, that would work. So, our design options are:

• give up type inference
• give up single object initialization
• give up avoiding this

[beorn]

I think it also depends on whether we're targeting a greenfield (starting from scratch) or brownfield (building on something) type of scenario.

If starting with Zustand, then I think the last solution is probably best: it enhances Zustand with direct subscriptions and computeds by extending, not changing, the syntax. It should be backwards compatible.

If starting with Jotai and/or the atomWithActions package, then it's a bit more open which way would make more sense. I guess it depends on how you think "actions" should fit with Jotai, and whether you'd like to be more similar to Zustand or Jotai (and target devs that prefer one over the other).

If considering this a completely greenfield project then I think it would make sense to value type inference higher — and it's more open whether it's better to have it be a new project or an extension of something that exists. In any case, you are the best person to judge where/if this kind of idea fits into the existing state management solution landscape :)

There's still the possibility that we may find ways to get type inference through TS tricks, or that TS adds support for inferring types for the kind of scenario we need — I suspect it creates some kind of recursive type problem (the argument type and the return types are co-dependent).

In any case there are several other choices to be made — e.g., for computeds, whether and how to track dependencies, and for the "observables", whether to expose them (to components and/or to the store definition functions) as some well-known reactive object like atoms, observables, signals, etc, or to just abstract all that away behind a Zustand useStore and useStoreValue hooks. I think it would also be cool to allow composition of stores.

[dai-shi]

My goal in this project is to bring zustand flavor to jotai atoms. So, this is essentially jotai and the technique behind it is atoms, not custom observables.

So, the big difference from zustand is actions are technically separate things, unlike zustand which has actions in its "state" object. It makes sense to separate state and actions.

const store = createStore({
  a: 1,
  b: 2,
  sum: (state) => state.a + state.2 // or use `this`, or use `this` with getter
}, {
  addA: (n: number) => (state) => ({ a: state.a + n }),
})

// it can be
// const [useValue, useAction] = create(...)

[beorn]

Okay, focusing on that goal, a few thoughts:

It is possible to separate out, at run-time and type inference time, values, functions, and property getters — so it is technically possible to define it all together, like we did previously, and by convention treat values as state, functions as actions, and property getters as computeds. But maybe it's a better fit with the rest of the ecosystem to have it separated out, or there could be some benefit in terms of complexity and type inference ability gained?

In any case, what are the features you'd like to see of the ones we've discussed? From my understanding they are:

• Type inference (may have to use this)
• Direct subscriptions (e.g., useStore.sum)
• Computeds (e.g., sum2x)
• Computeds with automatic dependency tracking
• Composability with other stores or atoms

My situation is that I (and it seems like most other people, judging by its popularity) like Zustand's simple central store model, and I would love to see some pathway to the above features (at least computeds and direct subscriptions). I presume that if it's not an extension of Zustand it'll be an extension of Jotai, so I'm assuming the developer tooling ecosystem would be good in any case.

[dai-shi]

My second goal is to provide a solution for people who think Zustand is more scalable (or suitable for large scale apps) than Jotai. So, the central store syntax would be important.

import { createStore, useStoreValue, useStoreAction } from 'jotai-zustand';

const store = createStore({
  a: 1,
  b: 2,
  get sum() { return this.a + this.b },
  get sum2() { return this.sum * 2 },
  addA: (n: number) => (state) => ({ a: state.a + n }),
});

const sum2 = useStoreValue(store.sum2);
const addA = useStoreAction(store.addA);

All first level properties become atoms, so it's direct subscription to them. No nesting are supported.

Composability with other stores or atoms

If that's important. We should make store.foo an atom. If that's the case, useStoreValue === useAtomValue and useStoreAction === useSetAtom. So, it makes a lot of sense.

[beorn]

I think that would be nice, and fits with the bringing Zustand flavor to Jotai while still remaining atom based.

Is there a difference between this and a factory method that creates a set of atoms? Is it in essence just a wrapper that give you a Zustand-like syntax to create the below:

import { Atom, atom, useAtom, WritableAtom, PrimitiveAtom } from "jotai"

// /** Store */
type StoreShape = {
  a: PrimitiveAtom<number>
  b: PrimitiveAtom<number>
  sum: Atom<number>
  sum2x: Atom<number>
  addA: WritableAtom<null, [number], void>
}

export const store: StoreShape = {
  a: atom(1),
  b: atom(2),
  sum: atom((get) => get(store.a) + get(store.b)),
  sum2x: atom((get) => get(store.sum) * 2),
  addA: atom(null, (get, set, update) => set(store.a, get(store.a) + update)),
}

// Component should now have correct types
export default function MyComponent() {
  const [a] = useAtom(store.a) // number
  const [b] = useAtom(store.b) // number
  const [sum] = useAtom(store.sum) // number
  const [sum2x] = useAtom(store.sum2x) // number
  const [, addA] = useAtom(store.addA) // (n: number) => void

  return (
    <div>
      <h1>Brainstorming 13</h1>
      <div>a: {a}</div>
      <div>b: {b}</div>
      <div>sum: {sum}</div>
      <div>sum2x: {sum2x}</div>
      <button onClick={() => addA(5)}>Add to A</button>
    </div>
  )
}

[dai-shi]

Yeah, that's exactly what I was thinking.

For addA, this might be easier for the factory function:

-   addA: (n: number) => (state) => ({ a: state.a + n }),
+   addA: (n: number) => ({ a: (prev) => prev + n }),

Or, maybe it doesn't matter much. I wonder what would be a Zustand flavored api.

[beorn]

Ok, I think it's cool. :)

Note that to differentiate between getters at the type level, you need a bit of hackery, but it's okay; you can rely on the fact that getters create readonly properties, and those are different than read-write properties — see IsEqual.

I'm also thinking that it may make sense to just treat all non-functions (non-actions = computeds or state) the same in terms of instrumenting for dependency tracking, so that you could have state that has both setters and getters. The only difference between computeds and state is that computeds only has getters, and therefore is readonly.

[dai-shi]

I'm also thinking that it may make sense to just treat all non-functions (non-actions = computeds or state) the same in terms of instrumenting for dependency tracking, so that you could have state that has both setters and getters.

I'm not entirely sure if I get it.

Is it in essence just a wrapper that give you a Zustand-like syntax to create the below:

Would you like to implement it and send a PR? Or, do you want me or someone to work on it?

[beorn]

I can take a stab at a PR, but I'd like to do a bit more analysis / design clarification first... :)

Goals

Benefit from central Zustand-ish store definition syntax, while keeping Jotai atom benefits such as direct subscriptions (no selectors) and derived values.

Store atom

Just trying to get the nomenclature clear. I don't think end-users have to necessarily see all of these terms.

Previously we've been saying store, and while I think we're clear it is meant in a Zustand store type of way, I realize in Jotai store already means a storage container for atoms, so it's going to potentially be confusing to just say store. Instead we could call it StoreAtom, which would be different from the atom store. (Alternatively Jotai's store could be renamed storage, which is also a common way to name such a thing, but it'd be a breaking change.)

Definition and properties

The StoreAtom defines a State using an object (StoreAtomDefinition) which uses the syntax we have been talking about above; using getters and potentially this. The definition supports three kinds of store atom properties:

• actions — defined using regular functions that can take any arguments and must return void or Partial<State>
• computeds — derived state defined using javascript property getters that create readonly properties
• state — any non-function that is not a computed is considered base state

For example:

import { createStoreAtom } from 'jotai-zustand';

const storeAtom = createStoreAtom({
  // state
  a: 1,
  b: 2,
  // computeds
  get sum() { return this.a + this.b; },
  get sum2() { return this.sum * 2; },
  // actions
  addA(n: number) { return { a: this.a + n }}
  // addA(n: number) { this.a += n; }
  // -- Alternatives that are not type inferable:
  // addA: (n: number) => (state) => ({ a: state.a + n }),
  // addA: (n: number) => ({ a: (prev) => prev + n }),
});

Realization into atoms

The createStoreAtom function takes the definition and realizes it into an object (StoreAtomRealization), mapping each of the properties on the definition object to atom properties on the realization object, e.g.,

// approximate realization
storeAtom = {
  a: atom(1),
  b: atom(2),
  sum: atom((get) => get(storeAtom.a) + get(storeAtom.b)),
  sum2x: atom((get) => get(storeAtom.sum) * 2),
  addA: atom(null, (get, set, update) => set(store.a, get(storeAtom.a) + update)),
};

Usage

These store atoms are just Jotai atoms that can be used as normal:

import { useAtom, useAtomValue, useSetAtom } from 'jotai';

const a = useAtomValue(store.a); // number
const b = useAtomValue(store.b); // number
const sum = useAtomValue(store.sum); // number
const sum2x = useAtomValue(store.sum2x); // number
const addA = useSetAtom(store.addA); // (n: number) => void

Recomputation

Computeds are recomputed when state it depends on is updated.

Cyclic dependencies will throw an Error.

Options

1. Actions syntax: a. type inferable using this and either mutating and/or returning Partial<State> (type inferrable) b. using methods that accept get/set (more similar to Zustand)
2. Computeds syntax - using get/set or this
3. Allow store state to be accessed using a selector, similarly to Zustand const ab = useStoreAtom(state => state.a + state.b)
4. Allow property setters as well as property getters (so you can have derived settable state)
5. Allow using atoms within atom store definition
6. Allow nested stores/state
7. Allow async state, computeds, actions, selectors
8. Realize computeds using atom composition OR using single atom that has its own dependency tracking on atom store state (just not sure how to do this yet, will have to prototype)

[dai-shi]

I'd like to do a bit more analysis / design clarification first...

Sure.

Benefit from central Zustand-ish store definition syntax, while keeping Jotai atom benefits such as direct subscriptions (no selectors) and derived values.

Sounds good!

Store atom

As the new function returns a set of atoms, it should be called "storeAtoms". However, I'm not sure if we should call it "store" at all. Some shorter names would be nice. "key", "bag", "def", hmm.

addA(n: number) { return { a: this.a + n }}

As we are using this anyway for getters, I think this seems the best fit.

Cyclic dependencies will throw an Error.

For now, it will be infinite loops, which errors.

Allow store state to be accessed using a selector

I think that's necessary to cover some edge cases. We need to ideate the api.

Allow property setters as well as property getters

I'd like to avoid it unless there's a huge demand. Hm, but on the second thought, it's actually very cute. On the third thought, it requires Proxies, so let's not do that.

Allow using atoms within atom store definition

I'd like to avoid it. It sounds mismatching with "store" model.

Allow nested stores/state

This would be too complicated. It depends on a use case. Can be a separate util?

Allow async state, computeds, actions, selectors

Async would be really nice. But, maybe not for the initial release.

Realize computeds using atom composition OR using single atom that has its own dependency tracking on atom store state

Not sure if I understand how it looks like. Isn't it like "Allow using atoms within atom store definition"?

[beorn]

Here's an updated spec/README — I'm trying out some nomenclature (atomic store, computeds => derived state) to see if it makes it easier/simpler.

It struck me that it is presumably possible to make this atomic store completely Zustand compatible, and it would probably be possible to wrap a Zustand store to make it an atomic store — it wouldn't have derived state, but you could add that if you wanted to.

Atomic Store

An atomic store is a type inferred central store defined using a State object with properties of these types:

• actions that update the state (defined using methods)
• derived state (defined using getters)
• basic state (all other properties)

The store exposes each of the properties as an appropriate Jotai atom which you can then consume/use to interact with the state in the store.

This way you can benefit from both the conciseness and simplicity of a central Zustand-ish store definition syntax, and the Jotai atom benefits such as cached, auto-updating derived values, and direct subscriptions that doesn't require selectors.

Definition

import { create } from 'jotai-zustand'

const atomicStore = create({
  a: 1,
  b: 2,

  // derived state defined using getters
  get sum() { return this.a + this.b },
  get sum2() { return this.sum * 2 },

  // actions return Partial<State> or mutate state directly
  adda(n: number) { return { a: this.a + n } },
  addb(n: number) { this.b += n },
});
// => {
//   a: PrimitiveAtom<number>
//   b: PrimitiveAtom<number>
//   sum: Atom<number>
//   sum2: Atom<number>
//   adda: WritableAtom<null, [number], void>
//   addb: WritableAtom<null, [number], void>
// };

All method properties on the state object are considered to be actions, and they must either mutate the state in the this object directly, or return Partial<State>, which will then be merged with the existing state.

Derived state (aka computeds or computed values) are defined using getters, and are automatically updated when the state they depend on changes. Be careful not to create circular dependencies.

Usage

The store can be consumed as a set of atoms:

import { useAtom, useAtomValue, useSetAtom } from 'jotai'

export default function MyComponent() {
  const a = useAtomValue(atomicStore.a) // number
  const sum2x = useAtomValue(atomicStore.sum2) // number
  const adda = useSetAtom(atomicStore.adda) // (n: number) => void

  return (
    <div>
      <div>a: {a}</div>
      <div>sum2x: {sum2x}</div>
      <button onClick={() => adda(5)}>Add 5 to a</button>
    </div>
  )
}

Or through useStore and selectors, similarly to how Zustand works:

import { useStore } from 'jotai-zustand'
const sum = useStore(atomicStore, (state) => state.sum)
const state = useStore(atomicStore)

Using selectors is not quite as performant as using atoms. Each useStore call in each component instance will register a selector that is called on every store update. This can be expensive if you render many components that use selectors.

Component instances that use atoms has no performance penalty unless the atom they depend on changes value.

Craz idea: Generalization

The state definition object above could actually connect to and bridge to other state systems, e.g.,

import { fromZustand, fromSignal, type State } from 'jotai-zustand'
const store = create({
  zustand: fromZustand(zustandStore), // composable
  signal: fromSignal(signal$), // maybe auto-detect type
  a: 1,
  b: 2,
  get sum() { return this.zustand.var + this.signal }
})
// => State<{
//   zustand: State<...zustandStore>,
//   signal: number,
//   a: number,
//   b: number,
//   sum: readonly number
// }>
fromAtomic(store, { // extensible
  get sum2() { return this.sum * 2 }
})
// => State<{
//   zustand: State<...zustandStore>,
//   signal: number,
//   a: number,
//   b: number,
//   sum: number,
//   sum2: number
// }>

toSignals(store)
// => {
//   zustand: { var: Signal<number> },
//   a: Signal<number>,
//   b: Signal<number>,
//   signal: Signal<number>,
//   sum: Signal<number>
// }
toAtoms(store)
// => {
//   zustand: { var: atom<...> },
//   signal: atom<number>,
//   a: atom<number>,
//   b: atom<number>,
//   sum: atom<number>
// }

To do

Must explore:

• [ ] Best way to track dependencies and create atoms
• [ ] Naming :)

Also likely explore:

• [ ] Generalization to other state systems
• [ ] Zustand compatibility
  • [ ] Consume store using selectors — ideate API (the above Zustand one looks good to me, but not clear how to deal with setting basic state)
  • [ ] Also offer a setState / getState API
  • [ ] Create atomic store from a Zustand store (and allow easy addition of derived state)
• [ ] Dealing with async (state, computeds, actions, selectors)
• [ ] Allow property setters in addition to property getters

Perhaps out of scope:

• [ ] Dealing with nested stores/state (I think this would be very useful)

Out of scope:

• [ ] Also allow using atoms within the store

[beorn]

You can check out https://github.com/beorn/jotai-zustand/tree/pr/atomic-store

[dai-shi]

const atomicStore = create({

Nice naming.

  // actions return Partial<State> or mutate state directly
  adda(n: number) { return { a: this.a + n } },
  addb(n: number) { this.b += n },

I don't prefer the direct mutation. I don't think we should allow addb. Unless we used Proxies, direct mutation would break React contract.

Craz idea: Generalization

It's an interesting idea, but it's out of the scope of the current goal. And, it should be a separate library than jotai-zustand, because its use case is broader.

[dai-shi]

import { useStore } from 'jotai-zustand'
const sum = useStore(atomicStore, (state) => state.sum)
const state = useStore(atomicStore)

Actually, const state = useStore(atomicStore) is probably not necessary and we should basically discourage such usage.

On the other hand, supporting const foo = useStore(atomicStore.foo) might be interesting. It thought it's too difficult because we need to distinguish actions from values. But, I've got an idea:

const EMPTY = Symbol();
const valueAtom = atom('str')
const actionAtom = atom(EMPTY, (get, set, arg) => ...);

const useValueOrActionAtom = (atom) => {
  const [value, action] = useAtom(atom);
  return value === EMPTY ? action : value;
};

[beorn]

I’ll reply more tomorrow - but just a quick note before signing off today: while not using Proxy, the code does wrap each property in a getter/setter. Not sure how you feel about that?

[dai-shi]

while not using Proxy, the code does wrap each property in a getter/setter. Not sure how you feel about that?

I'm still against to it. If this.foo = 1 were possible, I'd expect this.nested.bar = 2 would be possible too. And it feels Valtio's territory to me.

Zustand's mental model & design policy is "immutable state", and no mutation style. That's another reason.

[dai-shi]

- import { useStore } from 'jotai-zustand'
+ import { useAtomicStore } from 'jotai-zustand'

might be considerable.

[beorn]

Okay, this seems clear:

• "Atomic Store" is an okay Working Title for now
• To discourage bad practices we'll disallow useAtomicStore() without any args — it's not performant
• Generalize to other state systems out-of-scope

Questions:

• useAtomicStore(atom | selector) can be used as unified way to consume the store — but for basic state (PrimitiveAtom = read/write atoms), how would you access both the getter and setter?
• Mutations does work through a setter, so I think it's technically okay. Maybe we'd still disallow for DX/best practice reasons, or are there technical reasons too?
• The code does leverage wrapping all of the props in getters, but it doesn't use Proxy — not sure if that will cause a lot of problems?

[beorn]

https://github.com/beorn/jotai-zustand/tree/pr/atomic-store

• still just exploring, but...
• added type checks
• added atomicStoreFromZustand() helper method

[dai-shi]

• useAtomicStore(atom | selector) can be used as unified way to consume the store — but for basic state (PrimitiveAtom = read/write atoms), how would you access both the getter and setter?

Actually, we may want to separate it into two hooks. useAtomicStore and useAtomicStoreSelector. But, if the unified way is more ergonomic, it's fine too.

I mean, even for the basic state, it should be read-only.

• Mutations does work through a setter, so I think it's technically okay. Maybe we'd still disallow for DX/best practice reasons, or are there technical reasons too?

I'm a bit confused with the context. Is it about useAtomStore(atom)?

• The code does leverage wrapping all of the props in getters, but it doesn't use Proxy — not sure if that will cause a lot of problems?

Is it about actions? Partly my concern was the coding, but partly it's api design. I don't think mutation style fits the mental model.

Definition

import { createAtomicStore } from 'jotai-zustand'

const atomicStore = createAtomicStore({
  a: 1,
  b: 2,

  // derived state defined using getters
  get sum() { return this.a + this.b },
  get sum2() { return this.sum * 2 },

  // actions return Partial<State>, can use `this` for reading current state
  setA(n: number) { return { a: n } },
  addB(n: number) { return { b: this.b + n } },
});
// => {
//   a: PrimitiveAtom<number>
//   b: PrimitiveAtom<number>
//   sum: Atom<number>
//   sum2: Atom<number>
//   setA: WritableAtom<unique symbol, [number], void>
//   addB: WritableAtom<unique symbol, [number], void>
// };