withCompare: A composable version of atomWithCompare

[nwrahnema]

This is a pretty simple creation, but I'm new to Jotai and this took me a while to figure out so I figured I'd share.

Context: I'm syncing an array with the URL location using atomWithLocation - a simplified version of my code:

const locationAtom = atomWithLocation()
const atomWithLocationStore = <T>(key: string, defaultValue: T) => atom(
  (get) => {
    const serialized = get(locationAtom).searchParams?.get(key) ?? undefined
    if (serialized === undefined) {
      return defaultValue
    }
    return JSON.parse(serialized) as T
  },
  (_get, set, update: T) => {
    set(
      locationAtom,
      (prev) => {
        const searchParams = new URLSearchParams(prev.searchParams)
        searchParams.set(key, JSON.stringify(update))
        return { ...prev, searchParams }
      }
    )
  }
)

const numbersAtom = atomWithLocationStore("numbers", [])

The issue with this is that JSON.parse will create a new object on every read, so for reference types like arrays this will trigger a rerender whenever the URL changes (since atomWithLocation is subscribed to the entire URL, not just this query parameter).

After fiddling around with atomWithCompare, useReducerAtom, and useAtomEffect I eventually discovered selectAtom which leads to a very simple solution:

import _ from "lodash"

function withCompare<Value, Args extends unknown[], Result>(
  anAtom: WritableAtom<Value, Args, Result>,
  isEqual: (prev: Value, next: Value) => boolean
) {
  const withCompareAtom = selectAtom<Value, Value>(anAtom, _.identity, isEqual)
  return atom(withCompareAtom.read, anAtom.write)
}

const numbersAtom = withCompare(atomWithLocationStore("numbers", []), _.isEqual)

And we can make it a bit more generic by accepting any type of atom 😁

function withCompare<Value, Args extends unknown[], Result>(
  anAtom: WritableAtom<Value, Args, Result>,
  isEqual: (prev: Value, next: Value) => boolean
): WritableAtom<Value, Args, Result>

function withCompare<Value>(
  anAtom: Atom<Value>,
  isEqual: (prev: Value, next: Value) => boolean
): Atom<Value>

function withCompare<Value, Args extends unknown[], Result>(
  anAtom: Atom<Value> | WritableAtom<Value, Args, Result>,
  isEqual: (prev: Value, next: Value) => boolean
) {
  const withCompareAtom = selectAtom<Value, Value>(anAtom, _.identity, isEqual)

  const read = withCompareAtom.read
  const write = "write" in anAtom ? anAtom.write : undefined
  return write === undefined ? atom(read) : atom(read, write)
}

[dmaskasky]

Looks interesting. Please consider opening a PR. I think it could be a good fit in recipes or in jotai/utils. @dai-shi

[dai-shi]

I'd suggest something like this:

const identity = (x) => x;
const withCompare = (anAtom, isEqual) => {
  const selected = selectAtom(anAtom, identity, isEqual);
  return atom((get) => get(selected), (get, set, ...args) => set(anAtom, ...args);
};

So, your second version looks almost good.

btw, we don't generally recommend isEqual and teach selectAtom is for escape hatches.

It might be good for a recipe, but not for utils in which we only add some primitive ones.
Moreover, my recent suggestion is not to add anything in utils. Basically, we should start anything with jotaijs/* first and see how it adapts.

[dai-shi]

btw, we don't generally recommend isEqual

The issue with this is that JSON.parse will create a new object

I think the general idea to avoid, if possible, is to use JSON.parse in write, instead of read.

[nwrahnema]

Thanks for the review! I can put a PR up for a recipe.