use-select-atom.ts

import { Atom } from 'jotai';
import { selectAtom, useAtomValue} from 'jotai/utils';
import { useMemo } from 'react';
import { createTrackedSelector } from 'react-tracked';

type ResolveType<T> = T extends Promise<infer V> ? V : T

export function useSelectAtom<Value, Slice>(
    atom: Atom<Value>,
    selector: (v: ResolveType<Value>) => Slice,
    equalityFn: (a: Slice, b: Slice) => boolean = Object.is
): ResolveType<Slice> {
    const selectedAtom = useMemo(() => selectAtom(atom, selector, equalityFn), [atom, selector, equalityFn]);

    return useAtomValue(selectedAtom)
}

export function createAtomSelector<Value>(atom: Atom<Value>): <Slice>(selector: (value: Value) => Slice) => Slice {
    return <Slice>(selector: (v: ResolveType<Value>) => Slice) => {
        return useSelectAtom(atom, selector);
    }
}

export function createTrackedAtom<Value>(atom: Atom<Value>) {
    return createTrackedSelector(createAtomSelector(atom));
}