Implement splitAtomRecord

[arturaz]

Implementation:

/**
 * Split an atom that contains a {@link Record} into an atom that contains an atom for each key-value pair.
 *
 * Take note that, by default, if a record key disappears and then reappears later, you will get a new
 * {@link PrimitiveAtom} atom for it.
 *
 * You can disable this behavior by setting `prune: false`, but make sure the set of your keys is bounded,
 * otherwise old atoms will never be removed and the memory usage will grow unboundedly.
 *
 * Note: with `prune: true`, if someone holds onto a per-key atom after the key is removed, it will remain
 * "stale" and will not automatically start targeting the re-added key.
 */
export function splitAtomRecordUnsafe<Key extends string, Item>(
    recordAtom: PrimitiveAtom<Record<Key, Item>>,
    opts?: { prune?: boolean },
): Atom<Record<Key, PrimitiveAtom<Item>>> {
    const prune = opts?.prune ?? true

    // Stable per-key atoms (like splitAtom reuses atoms by key)
    const perKeyAtomCache = new Map<Key, PrimitiveAtom<Item>>()

    // Stable output record object if keys unchanged (like mapping reuse when keyList unchanged)
    let lastKeysSig: string | undefined
    let lastRecordOut: Record<Key, PrimitiveAtom<Item>> | undefined

    function getOrCreateItemAtom(key: Key): PrimitiveAtom<Item> {
        const existing = perKeyAtomCache.get(key)
        if (existing) return existing

        // splitAtom returns stale value for reads when item disappears, if possible
        let hasLastValue = false
        let lastValue!: Item

        const itemAtom = atom(
            (get) => {
                const rec = get(recordAtom)
                if (key in rec) {
                    const v = rec[key]
                    hasLastValue = true
                    lastValue = v
                    return v
                }
                if (hasLastValue) {
                    return lastValue
                }
                throw new Error(`splitAtomRecord: key "${s(key)}" missing for read`)
            },
            (get, set, action: SetStateAction<Item>) => {
                const prevRec = get(recordAtom)
                if (!(key in prevRec)) {
                    // splitAtom throws on write if item is out of bounds
                    throw new Error(`splitAtomRecord: key "${s(key)}" missing for write`)
                }

                const prevItem = prevRec[key]
                const nextItem = setStateActionFold<Item, Item>(
                    action,
                    (value) => value,
                    (fn) => fn(prevItem),
                )

                // Preserve identity if nothing changes
                if (Object.is(prevItem, nextItem)) return
                set(recordAtom, { ...prevRec, [key]: nextItem })
            },
        ) as PrimitiveAtom<Item>

        perKeyAtomCache.set(key, itemAtom)
        return itemAtom
    }

    return atom((get) => {
        const rec = get(recordAtom)

        // keep deterministic order
        const keys = Object.keys(rec).sort() as Key[]
        const keysSig = keys.join('\u0000')

        if (lastRecordOut && lastKeysSig === keysSig) {
            return lastRecordOut
        }

        // Optional: prune cache entries that no longer exist
        // (keeps memory bounded if keys churn)
        if (prune) {
            for (const cachedKey of perKeyAtomCache.keys()) {
                if (!Object.hasOwn(rec, cachedKey)) perKeyAtomCache.delete(cachedKey)
            }
        }

        const out = {} as Record<Key, PrimitiveAtom<Item>>
        for (const key of keys) {
            out[key] = getOrCreateItemAtom(key)
        }

        lastKeysSig = keysSig
        lastRecordOut = out
        return out
    })
}

/** Maps {@link React.SetStateAction} to other value.  */
export function setStateActionFold<A, R>(
    action: React.SetStateAction<A>,
    whenValue: (value: A) => R,
    whenFunction: (fn: (previous: A) => A) => R,
) {
    if (typeof action === 'function') {
        const actionFn = action as (prev: A) => A
        return whenFunction(actionFn)
    } else {
        const a: A = action
        return whenValue(a)
    }
}

Tests:

describe('splitAtomRecord', () => {
    function setup(
        fn: (
            recordAtom: PrimitiveAtom<Record<string, number>>,
            split: Atom<Record<string, PrimitiveAtom<number>>>,
        ) => void,
    ) {
        const recordAtom = atom<Record<string, number>>({
            a: 1,
            b: 2,
            c: 3,
        })
        const split = splitAtomRecord(recordAtom)
        return () => fn(recordAtom, split)
    }

    it(
        'should have correct keys',
        setup((_recordAtom, split) => {
            const atoms = testJotaiStore.get(split)
            expectTypeSafe(objectKeys(atoms)).toStrictEqual(['a', 'b', 'c'])
        }),
    )

    it(
        'should have correct values for the split atom',
        setup((_recordAtom, split) => {
            const atoms = testJotaiStore.get(split)

            expectTypeSafe(testJotaiStore.get(atoms.a!)).toStrictEqual(1)
            expectTypeSafe(testJotaiStore.get(atoms.b!)).toStrictEqual(2)
            expectTypeSafe(testJotaiStore.get(atoms.c!)).toStrictEqual(3)
        }),
    )

    it(
        'should propagate values from underlying record to split atoms',
        setup((recordAtom, split) => {
            const atoms = testJotaiStore.get(split)
            testJotaiStore.set(recordAtom, { a: 3, b: 2, c: 1 })

            expectTypeSafe(testJotaiStore.get(atoms.a!)).toStrictEqual(3)
            expectTypeSafe(testJotaiStore.get(atoms.b!)).toStrictEqual(2)
            expectTypeSafe(testJotaiStore.get(atoms.c!)).toStrictEqual(1)
        }),
    )

    it(
        'should propagate values from split atoms to underlying record',
        setup((recordAtom, split) => {
            const atoms = testJotaiStore.get(split)

            testJotaiStore.set(atoms.a!, 3)
            expectTypeSafe(testJotaiStore.get(recordAtom)).toStrictEqual({ a: 3, b: 2, c: 3 })

            testJotaiStore.set(atoms.b!, 2)
            expectTypeSafe(testJotaiStore.get(recordAtom)).toStrictEqual({ a: 3, b: 2, c: 3 })

            testJotaiStore.set(atoms.c!, 1)
            expectTypeSafe(testJotaiStore.get(recordAtom)).toStrictEqual({ a: 3, b: 2, c: 1 })
        }),
    )

    it(
        'should keep stable per-key atom identities when values change',
        setup((recordAtom, split) => {
            const atoms1 = testJotaiStore.get(split)
            const a1 = atoms1.a!
            const b1 = atoms1.b!
            const c1 = atoms1.c!

            testJotaiStore.set(recordAtom, { a: 10, b: 20, c: 30 })

            const atoms2 = testJotaiStore.get(split)
            expectTypeSafe(atoms2.a!).toBe(a1)
            expectTypeSafe(atoms2.b!).toBe(b1)
            expectTypeSafe(atoms2.c!).toBe(c1)
        }),
    )

    it(
        'should keep stable returned record identity when keys do not change',
        setup((recordAtom, split) => {
            const out1 = testJotaiStore.get(split)

            // change only values, not keys
            testJotaiStore.set(recordAtom, { a: 2, b: 3, c: 4 })

            const out2 = testJotaiStore.get(split)
            expectTypeSafe(out2).toBe(out1)
        }),
    )

    it(
        'should return a new record object when keys change (add key)',
        setup((recordAtom, split) => {
            const out1 = testJotaiStore.get(split)

            testJotaiStore.set(recordAtom, { a: 1, b: 2, c: 3, d: 4 })

            const out2 = testJotaiStore.get(split)
            expectTypeSafe(out2).not.toBe(out1)
            expectTypeSafe(objectKeys(out2)).toStrictEqual(['a', 'b', 'c', 'd'])

            // existing keys keep their atom identity
            expectTypeSafe(out2.a!).toBe(out1.a!)
            expectTypeSafe(out2.b!).toBe(out1.b!)
            expectTypeSafe(out2.c!).toBe(out1.c!)

            // new key has correct value
            expectTypeSafe(testJotaiStore.get(out2.d!)).toStrictEqual(4)
        }),
    )

    it(
        'should return a new record object when keys change (remove key)',
        setup((recordAtom, split) => {
            const out1 = testJotaiStore.get(split)

            testJotaiStore.set(recordAtom, { a: 1, c: 3 })

            const out2 = testJotaiStore.get(split)
            expectTypeSafe(out2).not.toBe(out1)
            expectTypeSafe(objectKeys(out2)).toStrictEqual(['a', 'c'])

            // remaining keys keep identity
            expectTypeSafe(out2.a!).toBe(out1.a!)
            expectTypeSafe(out2.c!).toBe(out1.c!)
            // removed key is gone
            expect(out2.b).toBeUndefined()
        }),
    )

    it(
        'should support SetStateAction function updates on per-key atoms',
        setup((recordAtom, split) => {
            const atoms = testJotaiStore.get(split)

            testJotaiStore.set(atoms.a!, (prev) => prev + 41)
            expectTypeSafe(testJotaiStore.get(recordAtom)).toStrictEqual({ a: 42, b: 2, c: 3 })

            testJotaiStore.set(atoms.b!, (prev) => prev * 10)
            expectTypeSafe(testJotaiStore.get(recordAtom)).toStrictEqual({ a: 42, b: 20, c: 3 })
        }),
    )

    it(
        'should throw if setting a removed key atom',
        setup((recordAtom, split) => {
            const out1 = testJotaiStore.get(split)
            const bAtom = out1.b!

            // remove b
            testJotaiStore.set(recordAtom, { a: 1, c: 3 })

            expect(() => testJotaiStore.set(bAtom, 999)).toThrow()
        }),
    )

    it(
        'should produce sorted keys (stable order)',
        setup((recordAtom, split) => {
            testJotaiStore.set(recordAtom, { c: 3, a: 1, b: 2, d: 4 })
            const out = testJotaiStore.get(split)
            expectTypeSafe(objectKeys(out)).toStrictEqual(['a', 'b', 'c', 'd'])
        }),
    )

    it(
        'should return stale value after key removal if it was previously read',
        setup((recordAtom, split) => {
            const atoms = testJotaiStore.get(split)
            const bAtom = atoms.b!

            // prime the "lastValue" cache (important!)
            expectTypeSafe(testJotaiStore.get(bAtom)).toStrictEqual(2)

            // remove b from the record
            testJotaiStore.set(recordAtom, { a: 1, c: 3 })

            // stale read behavior (matches upstream splitAtom intent)
            expectTypeSafe(testJotaiStore.get(bAtom)).toStrictEqual(2)
        }),
    )

    it(
        'should throw on read after key removal if the key atom was never successfully read',
        setup((recordAtom, split) => {
            const atoms = testJotaiStore.get(split)
            const bAtom = atoms.b!

            // remove b before any read of bAtom
            testJotaiStore.set(recordAtom, { a: 1, c: 3 })

            // no lastValue cached -> should throw (same spirit as upstream splitAtom throwing
            // if no stale value exists)
            expect(() => testJotaiStore.get(bAtom)).toThrow()
        }),
    )

    it(
        'should throw on write after key removal (even if stale read works)',
        setup((recordAtom, split) => {
            const atoms = testJotaiStore.get(split)
            const bAtom = atoms.b!

            // prime stale cache
            expectTypeSafe(testJotaiStore.get(bAtom)).toStrictEqual(2)

            // remove b
            testJotaiStore.set(recordAtom, { a: 1, c: 3 })

            // writes must throw when missing
            expect(() => testJotaiStore.set(bAtom, 999)).toThrow()
            expect(() => testJotaiStore.set(bAtom, (prev) => prev + 1)).toThrow()
        }),
    )

    it(
        'should return the last known underlying value as stale after key removal',
        setup((recordAtom, split) => {
            const atoms = testJotaiStore.get(split)
            const bAtom = atoms.b!

            // initial read
            expectTypeSafe(testJotaiStore.get(bAtom)).toStrictEqual(2)

            // update underlying value
            testJotaiStore.set(recordAtom, { a: 1, b: 200, c: 3 })
            expectTypeSafe(testJotaiStore.get(bAtom)).toStrictEqual(200)

            // remove key
            testJotaiStore.set(recordAtom, { a: 1, c: 3 })

            // stale read should return *last observed* value
            expectTypeSafe(testJotaiStore.get(bAtom)).toStrictEqual(200)
        }),
    )
})

[dai-shi]

Thanks for the suggestion. I think Record might be easier than Array. In any case, we don't add new features in the main repo, but grow our ecosystem. So, it'd be good for a separate library.