UNPKG

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