From 3aec3475962852c0a9518a61489a07d411e60a1c Mon Sep 17 00:00:00 2001
From: Daishi Kato <dai-shi@users.noreply.github.com>
Date: Fri, 29 Mar 2024 18:12:58 +0900
Subject: [PATCH] fix: avoid slow types (#2472)

---
 src/react/Provider.ts                   |  5 ++++-
 src/react/utils/useResetAtom.ts         |  6 +++---
 src/vanilla/store.ts                    | 26 ++++++++++++++++++-------
 src/vanilla/utils/atomWithObservable.ts |  6 ------
 src/vanilla/utils/atomWithReset.ts      |  5 ++++-
 src/vanilla/utils/freezeAtom.ts         |  4 ++--
 6 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/src/react/Provider.ts b/src/react/Provider.ts
index 5819561dea..3c57cee40c 100644
--- a/src/react/Provider.ts
+++ b/src/react/Provider.ts
@@ -4,7 +4,10 @@ import { createStore, getDefaultStore } from '../vanilla.ts'
 
 type Store = ReturnType<typeof createStore>
 
-const StoreContext = createContext<Store | undefined>(undefined)
+type StoreContextType = ReturnType<typeof createContext<Store | undefined>>
+const StoreContext: StoreContextType = createContext<Store | undefined>(
+  undefined,
+)
 
 type Options = {
   store?: Store
diff --git a/src/react/utils/useResetAtom.ts b/src/react/utils/useResetAtom.ts
index baa7a77639..1d13125350 100644
--- a/src/react/utils/useResetAtom.ts
+++ b/src/react/utils/useResetAtom.ts
@@ -5,10 +5,10 @@ import type { WritableAtom } from '../../vanilla.ts'
 
 type Options = Parameters<typeof useSetAtom>[1]
 
-export function useResetAtom(
-  anAtom: WritableAtom<unknown, [typeof RESET], unknown>,
+export function useResetAtom<T>(
+  anAtom: WritableAtom<unknown, [typeof RESET], T>,
   options?: Options,
-) {
+): () => T {
   const setAtom = useSetAtom(anAtom, options)
   const resetAtom = useCallback(() => setAtom(RESET), [setAtom])
   return resetAtom
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 36fa46cf22..453e682c2a 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -8,7 +8,7 @@ type OnUnmount = () => void
 type Getter = Parameters<AnyAtom['read']>[0]
 type Setter = Parameters<AnyWritableAtom['write']>[1]
 
-const isSelfAtom = (atom: AnyAtom, a: AnyAtom) =>
+const isSelfAtom = (atom: AnyAtom, a: AnyAtom): boolean =>
   atom.unstable_is ? atom.unstable_is(a) : a === atom
 
 const hasInitialValue = <T extends Atom<AnyValue>>(
@@ -20,7 +20,7 @@ const isActuallyWritableAtom = (atom: AnyAtom): atom is AnyWritableAtom =>
   !!(atom as AnyWritableAtom).write
 
 type CancelPromise = (next?: Promise<unknown>) => void
-const cancelPromiseMap = new WeakMap<Promise<unknown>, CancelPromise>()
+const cancelPromiseMap: WeakMap<Promise<unknown>, CancelPromise> = new WeakMap()
 
 const registerCancelPromise = (
   promise: Promise<unknown>,
@@ -124,6 +124,8 @@ type Mounted = {
   u?: OnUnmount
 }
 
+type MountedAtoms = Set<AnyAtom>
+
 // for debugging purpose only
 type StoreListenerRev2 = (
   action:
@@ -134,7 +136,19 @@ type StoreListenerRev2 = (
     | { type: 'restore'; flushed: Set<AnyAtom> },
 ) => void
 
-type MountedAtoms = Set<AnyAtom>
+type Store = {
+  get: <Value>(atom: Atom<Value>) => Value
+  set: <Value, Args extends unknown[], Result>(
+    atom: WritableAtom<Value, Args, Result>,
+    ...args: Args
+  ) => Result
+  sub: (atom: AnyAtom, listener: () => void) => () => void
+  dev_subscribe_store?: (l: StoreListenerRev2, rev: 2) => () => void
+  dev_get_mounted_atoms?: () => IterableIterator<AnyAtom>
+  dev_get_atom_state?: (a: AnyAtom) => AtomState | undefined
+  dev_get_mounted?: (a: AnyAtom) => Mounted | undefined
+  dev_restore_atoms?: (values: Iterable<readonly [AnyAtom, AnyValue]>) => void
+}
 
 /**
  * Create a new store. Each store is an independent, isolated universe of atom
@@ -152,7 +166,7 @@ type MountedAtoms = Set<AnyAtom>
  *
  * @returns A store.
  */
-export const createStore = () => {
+export const createStore = (): Store => {
   const atomStateMap = new WeakMap<AnyAtom, AtomState>()
   const mountedMap = new WeakMap<AnyAtom, Mounted>()
   const pendingStack: Set<AnyAtom>[] = []
@@ -824,11 +838,9 @@ export const createStore = () => {
   }
 }
 
-type Store = ReturnType<typeof createStore>
-
 let defaultStore: Store | undefined
 
-export const getDefaultStore = () => {
+export const getDefaultStore = (): Store => {
   if (!defaultStore) {
     defaultStore = createStore()
     if (import.meta.env?.MODE !== 'production') {
diff --git a/src/vanilla/utils/atomWithObservable.ts b/src/vanilla/utils/atomWithObservable.ts
index fa24605456..6c2c8f2293 100644
--- a/src/vanilla/utils/atomWithObservable.ts
+++ b/src/vanilla/utils/atomWithObservable.ts
@@ -4,12 +4,6 @@ import type { Atom, Getter, WritableAtom } from '../../vanilla.ts'
 type Timeout = ReturnType<typeof setTimeout>
 type AnyError = unknown
 
-declare global {
-  interface SymbolConstructor {
-    readonly observable: symbol
-  }
-}
-
 type Subscription = {
   unsubscribe: () => void
 }
diff --git a/src/vanilla/utils/atomWithReset.ts b/src/vanilla/utils/atomWithReset.ts
index 8dff76b7e5..2d9ae5929c 100644
--- a/src/vanilla/utils/atomWithReset.ts
+++ b/src/vanilla/utils/atomWithReset.ts
@@ -13,7 +13,10 @@ type WithInitialValue<Value> = {
   init: Value
 }
 
-export function atomWithReset<Value>(initialValue: Value) {
+export function atomWithReset<Value>(
+  initialValue: Value,
+): WritableAtom<Value, [SetStateActionWithReset<Value>], void> &
+  WithInitialValue<Value> {
   type Update = SetStateActionWithReset<Value>
   const anAtom = atom<Value, [Update], void>(
     initialValue,
diff --git a/src/vanilla/utils/freezeAtom.ts b/src/vanilla/utils/freezeAtom.ts
index 6918045760..e779c2d4aa 100644
--- a/src/vanilla/utils/freezeAtom.ts
+++ b/src/vanilla/utils/freezeAtom.ts
@@ -30,7 +30,7 @@ export function freezeAtom<AtomType extends Atom<unknown>>(
 
 export function freezeAtomCreator<
   CreateAtom extends (...params: never[]) => Atom<unknown>,
->(createAtom: CreateAtom) {
+>(createAtom: CreateAtom): CreateAtom {
   return ((...params: never[]) => {
     const anAtom = createAtom(...params)
     const origRead = anAtom.read
@@ -38,5 +38,5 @@ export function freezeAtomCreator<
       return deepFreeze(origRead.call(this, get, options))
     }
     return anAtom
-  }) as CreateAtom
+  }) as never
 }