From 480a053a71b073d24f850ea47a4b76cc2d87008d Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Tue, 9 Jan 2024 22:22:06 +0900
Subject: [PATCH] fix behavior with strict mode

---
 examples/01_typescript/src/index.tsx |  8 +++--
 src/atomWithCache.ts                 | 51 +++++++++++-----------------
 2 files changed, 25 insertions(+), 34 deletions(-)

diff --git a/examples/01_typescript/src/index.tsx b/examples/01_typescript/src/index.tsx
index bfd6d8f..b5a068d 100644
--- a/examples/01_typescript/src/index.tsx
+++ b/examples/01_typescript/src/index.tsx
@@ -1,9 +1,13 @@
-import React from 'react';
+import React, { StrictMode } from 'react';
 import { createRoot } from 'react-dom/client';
 
 import App from './App';
 
 const ele = document.getElementById('app');
 if (ele) {
-  createRoot(ele).render(<App />);
+  createRoot(ele).render(
+    <StrictMode>
+      <App />
+    </StrictMode>,
+  );
 }
diff --git a/src/atomWithCache.ts b/src/atomWithCache.ts
index a63a4ae..9ea988d 100644
--- a/src/atomWithCache.ts
+++ b/src/atomWithCache.ts
@@ -1,5 +1,5 @@
 import { atom } from 'jotai/vanilla';
-import type { Atom, Getter } from 'jotai/vanilla';
+import type { Atom, Getter, WritableAtom } from 'jotai/vanilla';
 
 type AnyAtomValue = unknown;
 type AnyAtom = Atom<AnyAtomValue>;
@@ -29,27 +29,24 @@ export function atomWithCache<Value>(
     writeGetterAtom.debugPrivate = true;
   }
 
-  const baseAtom = atom(
-    (get, opts) => {
-      const writeGetter = get(writeGetterAtom)?.[0];
-      if (writeGetter) {
-        const index = cache.findIndex((item) =>
-          Array.from(item[2]).every(([a, v]) => {
-            const vv = writeGetter(a);
-            if (vv instanceof Promise) {
-              return false;
-            }
-            return is(v, vv);
-          }),
-        );
-        if (index >= 0) {
-          const item = cache[index] as (typeof cache)[number];
-          if (options?.shouldRemove?.(...item)) {
-            cache.splice(index, 1);
-          } else {
-            item[2].forEach((_, a) => get(a)); // touch atoms
-            return item[1] as Value;
+  const baseAtom: WritableAtom<Value, [AnyAtom], AnyAtomValue> = atom(
+    (get, { setSelf: writeGetter, ...opts }) => {
+      const index = cache.findIndex((item) =>
+        Array.from(item[2]).every(([a, v]) => {
+          const vv = writeGetter(a);
+          if (vv instanceof Promise) {
+            return false;
           }
+          return is(v, vv);
+        }),
+      );
+      if (index >= 0) {
+        const item = cache[index] as (typeof cache)[number];
+        if (options?.shouldRemove?.(...item)) {
+          cache.splice(index, 1);
+        } else {
+          item[2].forEach((_, a) => get(a)); // touch atoms
+          return item[1] as Value;
         }
       }
       const map = new Map<AnyAtom, AnyAtomValue>();
@@ -64,23 +61,13 @@ export function atomWithCache<Value>(
       }
       return value;
     },
-    (get, set, isInit) => {
-      if (isInit) {
-        set(writeGetterAtom, [get]);
-      } else {
-        set(writeGetterAtom, null);
-      }
-    },
+    (get, _set, anAtom: AnyAtom) => get(anAtom), // HACK HACK HACK
   );
 
   if (process.env.NODE_ENV !== 'production') {
     baseAtom.debugPrivate = true;
   }
 
-  baseAtom.onMount = (init) => {
-    init(true);
-    return () => init(false);
-  };
   const derivedAtom = atom((get) => get(baseAtom));
   return derivedAtom;
 }