perf: implement weakAtomFamily (#117)

.changeset/tiny-worms-warn.md

@@ -0,0 +1,5 @@
+---
+"@reactive-dot/react": patch
+---
+
+Implemented `atomFamily` using `WeakMap`: this change allows for the automatic garbage collection of unused values, improving memory efficiency and overall performance.

packages/react/src/stores/accounts.ts

@@ -1,4 +1,4 @@
-import { withAtomFamilyErrorCatcher } from "../utils/jotai.js";
+import { weakAtomFamily, withAtomFamilyErrorCatcher } from "../utils/jotai.js";
 import { chainSpecDataAtomFamily } from "./client.js";
 import { walletsAtom } from "./wallets.js";
 import {
@@ -7,9 +7,9 @@ import {
   type PolkadotAccount,
 } from "@reactive-dot/core";
 import type { Atom } from "jotai";
-import { atomFamily, atomWithObservable } from "jotai/utils";
+import { atomWithObservable } from "jotai/utils";
 
-export const accountsAtom = atomFamily(
+export const accountsAtom = weakAtomFamily(
   (chainId: ChainId): Atom<PolkadotAccount[] | Promise<PolkadotAccount[]>> =>
     withAtomFamilyErrorCatcher(
       accountsAtom,

packages/react/src/stores/block.ts

@@ -1,18 +1,19 @@
+import { weakAtomFamily } from "../utils/jotai.js";
 import { clientAtomFamily } from "./client.js";
 import { type ChainId, getBlock } from "@reactive-dot/core";
-import { atomFamily, atomWithObservable } from "jotai/utils";
+import { atomWithObservable } from "jotai/utils";
 import { from } from "rxjs";
 import { switchMap } from "rxjs/operators";
 
-export const finalizedBlockAtomFamily = atomFamily((chainId: ChainId) =>
+export const finalizedBlockAtomFamily = weakAtomFamily((chainId: ChainId) =>
   atomWithObservable((get) =>
     from(get(clientAtomFamily(chainId))).pipe(
       switchMap((client) => getBlock(client, { tag: "finalized" })),
     ),
   ),
 );
 
-export const bestBlockAtomFamily = atomFamily((chainId: ChainId) =>
+export const bestBlockAtomFamily = weakAtomFamily((chainId: ChainId) =>
   atomWithObservable((get) =>
     from(get(clientAtomFamily(chainId))).pipe(
       switchMap((client) => getBlock(client, { tag: "best" })),

packages/react/src/stores/client.ts

@@ -1,10 +1,10 @@
+import { weakAtomFamily } from "../utils/jotai.js";
 import { chainConfigsAtom } from "./config.js";
 import type { ChainId } from "@reactive-dot/core";
 import { getClient, ReDotError } from "@reactive-dot/core";
 import { atom } from "jotai";
-import { atomFamily } from "jotai/utils";
 
-export const clientAtomFamily = atomFamily((chainId: ChainId) =>
+export const clientAtomFamily = weakAtomFamily((chainId: ChainId) =>
   atom(async (get) => {
     const chainConfig = get(chainConfigsAtom)[chainId];
 
@@ -16,15 +16,15 @@ export const clientAtomFamily = atomFamily((chainId: ChainId) =>
   }),
 );
 
-export const chainSpecDataAtomFamily = atomFamily((chainId: ChainId) =>
+export const chainSpecDataAtomFamily = weakAtomFamily((chainId: ChainId) =>
   atom(async (get) => {
     const client = await get(clientAtomFamily(chainId));
 
     return client.getChainSpecData();
   }),
 );
 
-export const typedApiAtomFamily = atomFamily((chainId: ChainId) =>
+export const typedApiAtomFamily = weakAtomFamily((chainId: ChainId) =>
   atom(async (get) => {
     const config = get(chainConfigsAtom)[chainId];
 

packages/react/src/stores/query.ts

@@ -1,4 +1,4 @@
-import { withAtomFamilyErrorCatcher } from "../utils/jotai.js";
+import { weakAtomFamily, withAtomFamilyErrorCatcher } from "../utils/jotai.js";
 import { stringify } from "../utils/vanilla.js";
 import { typedApiAtomFamily } from "./client.js";
 import {
@@ -10,10 +10,10 @@ import {
   type QueryInstruction,
 } from "@reactive-dot/core";
 import { atom, type Atom, type WritableAtom } from "jotai";
-import { atomFamily, atomWithObservable, atomWithRefresh } from "jotai/utils";
+import { atomWithObservable, atomWithRefresh } from "jotai/utils";
 import { from, switchMap, type Observable } from "rxjs";
 
-const instructionPayloadAtomFamily = atomFamily(
+const instructionPayloadAtomFamily = weakAtomFamily(
   (param: {
     chainId: ChainId;
     instruction: Exclude<
@@ -76,7 +76,7 @@ export function getQueryInstructionPayloadAtoms(
 
 // TODO: should be memoized within render function instead
 // https://github.com/pmndrs/jotai/discussions/1553
-export const queryPayloadAtomFamily = atomFamily(
+export const queryPayloadAtomFamily = weakAtomFamily(
   (param: { chainId: ChainId; query: Query }): Atom<unknown> =>
     withAtomFamilyErrorCatcher(
       queryPayloadAtomFamily,

packages/react/src/utils/jotai.ts

@@ -1,24 +1,102 @@
 import { QueryError } from "@reactive-dot/core";
 import type { Atom, Getter } from "jotai";
-import type { AtomFamily } from "jotai/vanilla/utils/atomFamily";
 import { Observable } from "rxjs";
 import { catchError } from "rxjs/operators";
 
+type WeakAtomFamily<TParam, TAtomType> = ((param: TParam) => TAtomType) & {
+  remove: (param: TParam) => void;
+};
+
+export function weakAtomFamily<TParam, TAtomType extends Atom<unknown>>(
+  initializeAtom: (param: TParam) => TAtomType,
+  areEqual?: (a: TParam, b: TParam) => boolean,
+): WeakAtomFamily<TParam, TAtomType> {
+  // TODO: should use `Symbol` as `WeakMap` key instead
+  // https://bugzilla.mozilla.org/show_bug.cgi?id=1710433
+  const keys = new Map<TParam, object>();
+  const atoms = new WeakMap<object, TAtomType>();
+
+  const getKey = (param: TParam) => {
+    if (areEqual === undefined) {
+      return keys.get(param);
+    }
+
+    for (const [key, value] of keys) {
+      if (areEqual(key, param)) {
+        return value;
+      }
+    }
+
+    return undefined;
+  };
+
+  const deleteKey = (param: TParam) => {
+    if (areEqual === undefined) {
+      return keys.delete(param);
+    }
+
+    for (const key of keys.keys()) {
+      if (areEqual(key, param)) {
+        return keys.delete(param);
+      }
+    }
+
+    return false;
+  };
+
+  return Object.assign(
+    (param: TParam) => {
+      const key = getKey(param);
+
+      if (key !== undefined) {
+        const atom = atoms.get(key);
+
+        if (atom !== undefined) {
+          return atom;
+        }
+      }
+
+      const newKey = {};
+      keys.set(param, newKey);
+
+      const newAtom = initializeAtom(param);
+      atoms.set(newKey, newAtom);
+
+      return newAtom;
+    },
+    {
+      remove: (param: TParam) => {
+        const key = getKey(param);
+
+        if (key === undefined) {
+          return;
+        }
+
+        deleteKey(param);
+        atoms.delete(key);
+      },
+    },
+  );
+}
+
 export class AtomFamilyError extends QueryError {
   constructor(
-    readonly atomFamily: AtomFamily<unknown, unknown>,
+    readonly atomFamily: WeakAtomFamily<unknown, unknown>,
     readonly param: unknown,
     message: string | undefined,
     options?: ErrorOptions,
   ) {
     super(message, options);
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  static fromAtomFamilyError<TError, TAtomFamily extends AtomFamily<any, any>>(
+  static fromAtomFamilyError<
+    TError,
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    TAtomFamily extends WeakAtomFamily<any, any>,
+  >(
     error: TError,
     atomFamily: TAtomFamily,
-    param: TAtomFamily extends AtomFamily<infer Param, infer _>
+    param: TAtomFamily extends WeakAtomFamily<infer Param, infer _>
       ? Param
       : unknown,
     message?: string,
@@ -36,7 +114,7 @@ export function withAtomFamilyErrorCatcher<
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   TAtomCreator extends (read: TRead, ...args: any[]) => Atom<unknown>,
 >(
-  atomFamily: AtomFamily<TParam, unknown>,
+  atomFamily: WeakAtomFamily<TParam, unknown>,
   param: TParam,
   atomCreator: TAtomCreator,
 ): TAtomCreator {