ideating with overwritable setAtom pattern

pmndrs · Nov 14, 2024 · aa9cdb2 · aa9cdb2
1 parent 2e41018
commit aa9cdb2
Showing 1 changed file with 15 additions and 12 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -600,22 +600,25 @@ const buildStore = (
       }
       if (isActuallyWritableAtom(atom)) {
         const mounted = atomState.m
+        let isSync = true
+        const createSetAtom = (pending: Pending) => {
+          return (...args: unknown[]) => {
+            try {
+              return writeAtomState(pending, atom, ...args)
+            } finally {
+              if (!isSync) {
+                flushPending(pending)
+              }
+            }
+          }
+        }
         addPendingFunction(pending, () => {
-          let isSync = true
-          let currPending = pending
+          let setAtom = createSetAtom(pending)
           try {
-            const onUnmount = atomOnMount(atom, (...args) => {
-              try {
-                return writeAtomState(currPending, atom, ...args)
-              } finally {
-                if (!isSync) {
-                  flushPending(currPending)
-                }
-              }
-            })
+            const onUnmount = atomOnMount(atom, (...args) => setAtom(...args))
             if (onUnmount) {
               mounted.u = (pending) => {
-                currPending = pending
+                setAtom = createSetAtom(pending)
                 isSync = true
                 try {
                   onUnmount()