How to manage native mutable data structures inside of atoms?

[marcebdev]

Examples would be Set & Map structures. Obviously when this has shown up for other reference based structures (arrays & objects) the obvious answer is to make it immutable by rebuilding a whole new array since both these items are very simple to make immutable.

However, when using a Set or Map this isn't really as viable. Of course with a Map you could always fall back to a simple object, but in the scope of a Set, it has unique functionality that's needed.

In theory when we call a setter with mySetValue.add(newVal) we'd want a way to flag the atom that it's had an update. One way to do this I'm guessing is to use atomWithRefresh to re-run the getter even if it's a bit roundabout, but I'm unclear on how to make this work.

type SetAction =
  | { type: 'add'; value: string }
  | { type: 'remove'; value: string }

const selectedSetAtom = atom<Set<string>>(new Set<string>())

export const selectedItemsAtom = atomWithRefresh<Set<string>, [SetAction], void>(
  (get) => get(selectedSetAtom),
  (get, set, action: SetAction) => {
    const currentState = get(selectedSetAtom)

    if (action.type === 'add') {
      currentState.add(action.value)
      set(selectedItemsAtom)
    } else if (action.type === 'remove') {
      currentState.delete(action.value)
      set(selectedItemsAtom)
    }
  },
)

ImmutableJS does seem to have an immutable version of a Set, but based on the nature of how our set is used as read-only it'd be preferable to simply have some pattern for re-running the getter like with the refresh pattern.

This would also speak to when you have large chunks of data with objects you'd rather mutate in a structure way i.e. using Maps. I can't really see any large discussion about this point, but it seems like an important one.

[dai-shi]

Mutable data structure doesn't follow the React contract, and it may behave weirdly with Concurrent React.

That being said, what we can do with atoms is to create a version. The idea is similar to atomWithRefresh.

const internalSetAtom = atomWithLazy({
  set: new Set(),
  version: 0,
});
const mySetAtom = atom(
  (get) => get(internalSetAtom).set,
  (get, set, action: SetAction) => {
    const current = get(internalSetAtom)
    if (action.type === 'add') {
      current.set.add(action.value)
      set({ ...current, version: current.version + 1 })
    } else if (action.type === 'remove') {
      current.set.delete(action.value)
      set({ ...current, version: current.version + 1 })
    }
  },
);

Actually, in this example, we don't need version property. 😄

[marcebdev]

FWIW I used this as a template and ended up coming up with this:

I did not realize at the time that the Set constructor can rebuild from another set aka (new Set(oldSet)) making immutability not much of a concern, probably the constructor has it's own optimizations. I'd imagine unless you had a very very busy set mutability wouldn't actually be much of a concern (especially given it's already a pretty fast data structure)

import { atomWithReducer } from 'jotai/utils'

type MapKey = string | number

type MapAction<T> =
  | { type: 'set'; value: [key: MapKey, value: T] }
  | { type: 'delete'; value: MapKey }
  | { type: 'clear' }

export type MapTemplate<V> = Map<MapKey, V>

// we don't want to allow mutating the state value itself
type MapState<V> = Pick<
  MapTemplate<V>,
  'get' | 'has' | 'size' | 'keys' | 'values' | 'entries'
>

export function createMapAtom<V>(initialState: MapTemplate<V> = new Map()) {
  return atomWithReducer<MapState<V>, MapAction<V>>(
    initialState,
    (prev, action) => {
      if (!action) throw new Error('Invalid action')

      // abort forcing an update if value will be unchanged
      if (action.type === 'delete' && !prev.has(action.value)) return prev

      // Mutative actions
      const currentState = new Map(prev as MapTemplate<V>)
      if (action.type === 'set') {
        const [key, value] = action.value
        currentState.set(key, value)
      } else if (action.type === 'delete') {
        currentState.delete(action.value)
      } else if (action.type === 'clear') {
        currentState.clear()
      }

      // Return immutable operation
      return currentState
    },
  )
}

[dai-shi]

That's much better as it follows the immutability contract.