UNPKG

1.85 kBJavaScriptView Raw
1import { produce } from 'immer';
2import { atom, useAtom } from 'jotai';
3import { useCallback } from 'react';
4
5function atomWithImmer(initialValue) {
6 const anAtom = atom(initialValue, (get, set, fn) => set(anAtom, produce(get(anAtom), typeof fn === "function" ? fn : () => fn)));
7 return anAtom;
8}
9
10function useImmerAtom(anAtom, scope) {
11 const [state, setState] = useAtom(anAtom, scope);
12 const setStateWithImmer = useCallback((fn) => setState(produce(fn)), [setState]);
13 return [state, setStateWithImmer];
14}
15
16const getWeakCacheItem = (cache, deps) => {
17 do {
18 const [dep, ...rest] = deps;
19 const entry = cache.get(dep);
20 if (!entry) {
21 return;
22 }
23 if (!rest.length) {
24 return entry[1];
25 }
26 cache = entry[0];
27 deps = rest;
28 } while (deps.length);
29};
30const setWeakCacheItem = (cache, deps, item) => {
31 do {
32 const [dep, ...rest] = deps;
33 let entry = cache.get(dep);
34 if (!entry) {
35 entry = [ new WeakMap()];
36 cache.set(dep, entry);
37 }
38 if (!rest.length) {
39 entry[1] = item;
40 return;
41 }
42 cache = entry[0];
43 deps = rest;
44 } while (deps.length);
45};
46const createMemoizeAtom = () => {
47 const cache = /* @__PURE__ */ new WeakMap();
48 const memoizeAtom = (createAtom, deps) => {
49 const cachedAtom = getWeakCacheItem(cache, deps);
50 if (cachedAtom) {
51 return cachedAtom;
52 }
53 const createdAtom = createAtom();
54 setWeakCacheItem(cache, deps, createdAtom);
55 return createdAtom;
56 };
57 return memoizeAtom;
58};
59
60const memoizeAtom = createMemoizeAtom();
61function withImmer(anAtom) {
62 return memoizeAtom(() => {
63 const derivedAtom = atom((get) => get(anAtom), (get, set, fn) => set(anAtom, produce(get(anAtom), typeof fn === "function" ? fn : () => fn)));
64 return derivedAtom;
65 }, [anAtom]);
66}
67
68export { atomWithImmer, useImmerAtom, withImmer };