atomWith: Storage And Compare

[PhilipGrefe]

I'd like to create an atomWithStorageAndCompare that

• is in sync with the storage via the subscribe callback
• only updates itself if the new value is different from the previous value.

The goal is to prevent rerenders after db changes if the actual atom value didn't change.

I assume it's not possible with atomWithStorage and that I have to create my own atomWithStorageAndCompare. However, I'm unable to achieve that. Any help is much appreciated 🙏

--
Below, I've create an atomWithStorage for visualisation/clarification of my goal:

const myAtom = atomWithStorage(
  'myAtom',
  [],
  {
    getItem: getFromDB,
    setItem: () => {}, // not necessary for this example
    removeItem: () => {}, // not necessary for this example
    subscribe: (key, callback, initialValue) => {
      const listener = addDatabaseChangeListener(async () => {
        // compare previous and next values
        const prevValue = "?" // how to get this value 🤔
        const nextValue = await getFromDB()
        if (areEqual(prevValue, nextValue)) return

        callback(nextValue)
        })
      
      return () => {
        listener.remove()
      }
    },
  },
)

[dai-shi]

I have to create my own atomWithStorageAndCompare.

In general, yes.

However, I'm unable to achieve that.

In your case, what would be easier it to use atomWithStorage and selectAtom. Doesn't it work?

[PhilipGrefe]

SelectAtom is a great idea!

Unfortunately, selectAtom from utils does not work in my case because I also need the write function on the original atom. AFAIK selectAtom only gives the read function.

After lots of trial and error, I figured I could "fork" the selectAtom to support the write direction as well. I've done it like so:

import { atom, Atom, WritableAtom } from 'jotai'

const getCached = <T>(c: () => T, m: WeakMap<object, T>, k: object): T => (m.has(k) ? m : m.set(k, c())).get(k) as T
const cache1 = new WeakMap()
const memo3 = <T>(create: () => T, dep1: object, dep2: object, dep3: object): T => {
  const cache2 = getCached(() => new WeakMap(), cache1, dep1)
  const cache3 = getCached(() => new WeakMap(), cache2, dep2)
  return getCached(create, cache3, dep3)
}

export function selectAtomWithWrite<Value, Args extends unknown[], Result, Slice>(
  anAtom: WritableAtom<Value, Args, Result>, // it's a writableAtom now
  selector: (v: Value, prevSlice?: Slice) => Slice,
  equalityFn?: (a: Slice, b: Slice) => boolean,
): WritableAtom<Slice, Args, Result>

export function selectAtomWithWrite<Value, Args extends unknown[], Result, Slice>(
  anAtom: WritableAtom<Value, Args, Result>,
  selector: (v: Value, prevSlice?: Slice) => Slice,
  equalityFn: (prevSlice: Slice, slice: Slice) => boolean = Object.is,
) {
  return memo3(
    () => {
      const EMPTY = Symbol()
      const selectValue = ([value, prevSlice]: readonly [Value, Slice | typeof EMPTY]) => {
        if (prevSlice === EMPTY) {
          return selector(value)
        }
        const slice = selector(value, prevSlice)
        return equalityFn(prevSlice, slice) ? prevSlice : slice
      }
      const derivedAtom: Atom<Slice | typeof EMPTY> & {
        init?: typeof EMPTY
      } = atom(
        (get) => {
          const prev = get(derivedAtom)
          const value = get(anAtom)
          return selectValue([value, prev] as const)
        },
        (get, set, ...args: Args) => {
          set(anAtom, ...args)
        },
      )
      // HACK to read derived atom before initialization
      derivedAtom.init = EMPTY
      return derivedAtom
    },
    anAtom,
    selector,
    equalityFn,
  )
}

Usage in a component

const fooAtom = selectAtomWithWrite(
  barAtom, // some atomWithStorage. I'm using a db with drizzle - hence it's potentially async
  (bar) => {
    if (bar instanceof Promise) { // I did not run into 'bar' being a promise (yet). If so, this could be resolved by returning the resolved value
      throw new Error('event is a promise')
    }
    return bar.foo
  },
  deepEqual,
)

[dai-shi]

I also need the write function on the original atom. AFAIK selectAtom only gives the read function.

It's by design. You should composite atoms again.

const baseAtom = atom(...)
const selectedAtom = selectAtom(baseAtom, ..., ...)
const finalAtom = atom(
  (get) => get(selectedAtom),
  (_get, set, arg) => set(baseAtom, arg)
)

[PhilipGrefe]

Thanks a lot for the comment - the composition works perfectly.