Merge pull request #2 from jotaijs/atom-tree

dmaskasky · web-flow · commit 33fdf9b452d5 · 2025-01-24T16:07:09.000-08:00
feat: atomTree
diff --git a/README.md b/README.md
@@ -109,3 +109,107 @@ const todoFamily = atomFamily(
 ### Codesandbox
 
 <CodeSandbox id="huxd4i" />
+
+---
+
+## atomTree
+
+The **atomTree** utility provides a hierarchical way to create, reuse, and remove Jotai atoms. Each atom is associated with a unique path, which is an array of unknown types. When you request the same path multiple times, `atomTree` ensures that the same atom instance is returned. You can also remove a specific atom or an entire subtree of atoms when they are no longer needed.
+
+Use `atomTree` when you anticipate a large number of potential paths and want to:
+
+- **Reuse the same atom** for repeated paths.  
+- **Clean up** unwanted atoms easily, including entire subtrees.
+
+```js
+import { atom } from 'jotai'
+import { atomTree } from 'jotai-family'
+
+// Create a tree instance, passing a factory function
+// that takes a path array and returns a new atom.
+const tree = atomTree((path) => atom(path.join('-')))
+
+// Create or retrieve the atom at ['foo', 'bar']
+const atomA = tree(['foo', 'bar'])
+const atomB = tree(['foo', 'bar'])
+
+// atomA and atomB are the same instance.
+console.log(atomA === atomB) // true
+
+// Remove the atom at ['foo', 'bar']
+// (and optionally remove its entire subtree)
+tree.remove(['foo', 'bar'])
+```
+
+### API
+
+#### Creating the tree
+
+Creates a new hierarchical tree of Jotai atoms. It accepts a **initializePathAtom** function that receives a path array and returns an atom. The returned function can be used to create, retrieve, and remove atoms at specific paths.
+
+```ts
+function atomTree<Path, AtomType>(
+  initializePathAtom: (path: Path) => AtomType
+): {
+  (path: Path): AtomType
+  remove(path: Path, removeSubTree?: boolean): void
+  getSubTree(path: Path): Node<AtomType> | undefined
+  getNodePath(path: Path): Node<AtomType>[]
+}
+
+type Node<AtomType> = {
+  atom?: AtomType
+  children?: Map<PathSegment, Node<AtomType>>
+}
+```
+
+### Creating Path Atoms
+```ts
+tree(path: Path): AtomType
+```
+Creates (or retrieves) an atom at the specified path. Subsequent calls with the same path return the same atom instance.
+
+### Removing Path Atoms
+```ts
+tree.remove(path: Path, removeSubTree = false): void
+```
+Removes the atom at the specified path. If `removeSubTree` is `true`, all child paths under that path are also removed.
+
+This method removes the atom at the specified path. If `removeSubTree` is `true`, it also removes all child paths under that path.
+
+### Retrieving A Subtree
+```ts
+tree.getSubTree(path: Path): Node<AtomType> | undefined
+```
+
+Retrieves the internal node representing the specified path. This is useful for inspecting the tree structure. The node structure is as follows:
+
+```ts
+type Node<AtomType> = {
+  atom?: AtomType
+  children?: Map<PathSegment, Node<AtomType>>
+}
+```
+
+### Retrieving A Node Path
+```ts
+tree.getNodePath(path: Path): Node<AtomType>[]
+```
+Returns an array of node objects from the root node to the node at the specified path, inclusive.
+
+## Usage Example
+
+```js
+import { atom } from 'jotai'
+import { atomTree } from 'jotai-family'
+
+const btree = atomTree((path) => atom(`Data for path: ${path}`))
+
+// Create or retrieve the atom at [true, false]
+const userAtom = btree([true, false])
+
+console.log(store.get(userAtom)) // 'Data for path: true,false'
+
+// Remove the atom (and optionally remove its subtree)
+btree.remove([true,false])
+```
diff --git a/src/atomTree.ts b/src/atomTree.ts
@@ -0,0 +1,116 @@
+import { Atom } from 'jotai'
+
+type Node<AtomType extends Atom<unknown> = Atom<unknown>> = {
+  children?: Map<unknown, Node<AtomType>>
+  atom?: AtomType
+}
+
+type AtomTree<Path extends unknown[], AtomType extends Atom<unknown>> = {
+  (path: Path): AtomType
+  remove(path?: Path, removeSubTree?: boolean): void
+  getSubTree(path?: Path): Node<AtomType>
+  getNodePath(path?: Path): Node<AtomType>[]
+}
+
+/**
+ * Creates a hierarchical structure of Jotai atoms.
+ *
+ * @template AtomType - The type of atom returned by the initialization function.
+ * @param initializePathAtom - A function that takes a path array and returns an Atom.
+ * @returns A function for creating and managing hierarchical atoms (with additional methods).
+ */
+export function atomTree<
+  Path extends unknown[],
+  AtomType extends Atom<unknown>,
+>(initializePathAtom: (path: Path) => AtomType): AtomTree<Path, AtomType> {
+  const root: Node<AtomType> = {}
+  const defaultPath = [] as unknown[] as Path
+
+  /**
+   * Creates or retrieves an atom at the specified path in the hierarchy.
+   *
+   * @param path - Array of keys representing the location in the hierarchy.
+   * @returns The Jotai atom for the specified path.
+   */
+  const createAtom: AtomTree<Path, AtomType> = (path) => {
+    let node = root
+    for (const key of path) {
+      node.children ??= new Map()
+      if (!node.children.has(key)) {
+        node.children.set(key, {})
+      }
+      node = node.children.get(key)!
+    }
+    node.atom ??= initializePathAtom(path)
+    return node.atom
+  }
+
+  /**
+   * Removes an atom (and optionally its subtree) at the specified path.
+   *
+   * @param path - Array of keys representing the location in the hierarchy. Defaults to [] (root).
+   * @param removeSubTree - If true, removes all children of that path as well. Defaults to false.
+   * @throws Error if the path does not exist.
+   */
+  createAtom.remove = (path = defaultPath, removeSubTree = false) => {
+    const nodePath = createAtom.getNodePath(path)
+    const node = nodePath.at(-1)!
+    delete node.atom
+    if (removeSubTree) {
+      delete node.children
+    }
+    // delete empty subtrees from bottom to top
+    for (let i = nodePath.length - 1; i >= 0; i--) {
+      const node = nodePath[i]!
+      if (!node.children?.size && i > 0) {
+        const parentNode = nodePath[i - 1]!
+        parentNode.children!.delete(path[i]!)
+        if (!parentNode.children!.size) {
+          delete parentNode.children
+        }
+      } else {
+        break
+      }
+    }
+  }
+
+  /**
+   * Retrieves the internal node (subtree) at a specified path.
+   *
+   * @param path - Array of keys representing the location in the hierarchy. Defaults to [] (root).
+   * @returns A Node object with possible `children` and `atom`.
+   * @throws Error if the path does not exist.
+   */
+  createAtom.getSubTree = (path = defaultPath) => {
+    let node: Node<AtomType> | undefined = root
+    for (const key of path) {
+      node = node.children?.get(key)
+      if (!node) {
+        throw new Error('Path does not exist')
+      }
+    }
+    return node
+  }
+
+  /**
+   * Retrieves the internal node (subtree) at a specified path.
+   *
+   * @param path - Array of keys representing the location in the hierarchy. Defaults to [] (root).
+   * @returns An array of Node objects representing the path.
+   * @throws Error if the path does not exist.
+   */
+  createAtom.getNodePath = (path = defaultPath) => {
+    const nodePath = [root]
+    let node: Node<AtomType> | undefined = root
+    for (const key of path) {
+      node = node.children?.get(key)
+      if (!node) {
+        throw new Error('Path does not exist')
+      }
+      nodePath.push(node)
+    }
+    return nodePath
+  }
+
+  return createAtom
+}
diff --git a/tests/atomTree.test.ts b/tests/atomTree.test.ts
@@ -0,0 +1,106 @@
+import { atom, createStore } from 'jotai/vanilla'
+import { describe, expect, it, vi } from 'vitest'
+import { atomTree } from '../src/atomTree'
+
+describe('atomTree', () => {
+  it('creates atoms at a given path', function test() {
+    const store = createStore()
+    const initializeAtom = vi.fn((path: number[]) => atom(path.join('-')))
+    const tree = atomTree(initializeAtom)
+
+    const atomA = tree([1, 2, 3])
+    expect(initializeAtom).toHaveBeenCalledTimes(1)
+    expect(store.get(atomA)).toBe('1-2-3')
+
+    const atomB = tree([1, 2, 3])
+    expect(initializeAtom).toHaveBeenCalledTimes(1)
+    // Should return the same atom instance if called with the same path
+    expect(atomA).toBe(atomB)
+    expect(store.get(atomB)).toBe('1-2-3')
+
+    const atomC = tree([1, 3])
+    expect(initializeAtom).toHaveBeenCalledTimes(2)
+    expect(store.get(atomC)).toBe('1-3')
+  })
+
+  it('removes an atom at a given path', function test() {
+    const store = createStore()
+    const initializeAtom = (path: (number | string)[]) => atom(path.join('-'))
+    const tree = atomTree(initializeAtom)
+
+    // Create an atom
+    const myAtom = tree([1, 'test'])
+    expect(store.get(myAtom)).toBe('1-test')
+
+    // Remove it
+    tree.remove([1, 'test'])
+    // Re-creating should yield a new atom
+    const newAtom = tree([1, 'test'])
+    expect(newAtom).not.toBe(myAtom)
+  })
+
+  it('removes an atom subtree correctly', function test() {
+    const store = createStore()
+    const initializeAtom = (path: string[]) => atom(path.join('-'))
+    const tree = atomTree(initializeAtom)
+
+    const atomA = tree(['parent', 'childA'])
+    expect(store.get(atomA)).toBe('parent-childA')
+
+    const atomB = tree(['parent', 'childB'])
+    expect(store.get(atomB)).toBe('parent-childB')
+
+    // Remove whole subtree under 'parent'
+    tree.remove(['parent'], true)
+
+    // Re-create them; they should be new atoms
+    const atomA2 = tree(['parent', 'childA'])
+    const atomB2 = tree(['parent', 'childB'])
+    expect(atomA2).not.toBe(atomA)
+    expect(atomB2).not.toBe(atomB)
+  })
+
+  it('throws if path does not exist when removing without prior creation', () => {
+    const tree = atomTree((path) => atom(path.join('-')))
+    expect(() => tree.remove(['nonexistent'])).toThrowError(
+      'Path does not exist'
+    )
+  })
+
+  it('retrieves subtree node correctly', function test() {
+    const store = createStore()
+    const tree = atomTree((path: string[]) => atom(path.join('-')))
+
+    const myAtom = tree(['foo', 'bar'])
+    expect(store.get(myAtom)).toBe('foo-bar')
+
+    // getSubTree should give us the correct node
+    const subTree = tree.getSubTree(['foo', 'bar'])
+    expect(subTree.atom).toBe(myAtom)
+  })
+
+  it('throws when retrieving subtree that does not exist', function test() {
+    const tree = atomTree((path: string[]) => atom(path.join('-')))
+    expect(() => tree.getSubTree(['no', 'such', 'thing'])).toThrowError(
+      'Path does not exist'
+    )
+  })
+
+  it('retrieves the node path correctly', function test() {
+    const tree = atomTree((path) => atom(path.join('-')))
+
+    tree(['branch', 'leaf'])
+    const nodePath = tree.getNodePath(['branch', 'leaf'])
+    expect(nodePath).toHaveLength(3)
+    expect(nodePath[0]!.children?.has('branch')).toBe(true)
+    expect(nodePath[1]!.children?.has('leaf')).toBe(true)
+    expect(nodePath[2]!.atom).toBeDefined()
+  })
+
+  it('throws when retrieving node path that does not exist', function test() {
+    const tree = atomTree((path: unknown[]) => atom(path.join('-')))
+    expect(() => tree.getNodePath(['no', 'such', 'thing'])).toThrowError(
+      'Path does not exist'
+    )
+  })
+})
