1 | import { produce } from 'immer';
|
2 | import { atom, useAtom } from 'jotai';
|
3 | import { useCallback } from 'react';
|
4 |
|
5 | function 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 |
|
19 | function 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 |
|
28 | const 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 | };
|
42 | const 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 | };
|
58 | const createMemoizeAtom = () => {
|
59 | const cache = 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 |
|
72 | const memoizeAtom = createMemoizeAtom();
|
73 | function 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 |
|
89 | export { atomWithImmer, useImmerAtom, withImmer };
|