feat: add atomTree utility

- createAtom
- createAtom.remove
- createAtom.getSubTree
- createAtom.getNodePath

Author: dmaskasky

README.md

@@ -109,3 +109,113 @@ const todoFamily = atomFamily(
 ### Codesandbox
 
 <CodeSandbox id="huxd4i" />
+
+---
+
+# atomTree
+
+[atomTree](https://github.com/jotaijs/jotai-family/blob/main/src/atomTree.ts) is a tree structure that allows you to create and remove atoms at a given path.
+
+Use `atomTree` when you need a hierarchical way to store atoms, particularly if you expect many potential paths and want to reuse the same atom for repeated paths while also having an easy cleanup mechanism for subtrees.
+
+```js
+import { atom } from 'jotai'
+import { atomTree } from 'jotai-family'
+
+const tree = atomTree((path) => atom(path.join('-')))
+
+const atomA = tree(['foo', 'bar'])
+const atomB = tree(['foo', 'bar'])
+
+// Remove the atom at the given path
+tree.remove(['foo', 'bar'])
+```
+
+## 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 **factory** function that receives a path array and returns an atom.
+
+```ts
+type Path = string[] // Or any other array type
+type AtomType = Atom<unknown>
+
+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>[]
+}
+```
+
+- **`initializePathAtom`**: A function invoked whenever the tree needs to create a new atom. Receives the `path` as an argument and must return a Jotai atom.
+
+The returned function has four main operations:
+
+1. **`tree(path: Path): AtomType`**  
+   Creates (or retrieves) an atom at the specified path. Subsequent calls with the same path return the same atom instance.
+
+2. **`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.
+
+3. **`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>>
+   }
+   ```
+
+4. **`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])
+```

src/atomTree.ts

@@ -1,54 +1,116 @@
 import { Atom } from 'jotai'
 
-export function atomTree<Param, AtomType extends Atom<unknown>>(
-  initializeAtom: (path: [Param, ...Param[]]) => AtomType
-) {
-  type Node = { m?: Map<Param, Node>; v?: AtomType }
-  const root = new Map<Param, Node>()
-
-  const createAtom = (path: [Param, ...Param[]]): AtomType => {
-    if (path.length === 0) {
-      throw new Error('Path must have at least one key')
-    }
-    let current = root
-    for (let i = 0; i < path.length; i++) {
-      const key = path[i]!
-      let node = current.get(key)
-      if (!node) {
-        node = {}
-        current.set(key, node)
-      }
-      if (i < path.length - 1) {
-        current = node.m = new Map()
-        continue
+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 = new Array() 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, {})
       }
-      return (node.v = initializeAtom(path))
+      node = node.children.get(key)!
     }
-    throw new Error('Unreachable')
+    node.atom ??= initializePathAtom(path)
+    return node.atom
   }
-  createAtom.remove = (path: [Param, ...Param[]]) => {
-    if (path.length === 0) {
-      throw new Error('Path must have at least one key')
-    }
-    const current = root
-    const nodePath: Node[] = []
-    for (const segment of path) {
-      const node = current.get(segment)
-      if (!node) {
-        break
-      }
-      nodePath.push(node)
+
+  /**
+   * 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 nodePath[nodePath.length - 1]!.v
-    // delete empty subtrees
+    // delete empty subtrees from bottom to top
     for (let i = nodePath.length - 1; i >= 0; i--) {
-      const current = nodePath[i]!
-      if (current.m?.size === 0 && i > 0) {
-        nodePath[i - 1]!.m!.delete(path[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 = 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
 }

tests/atomTree.test.ts

@@ -0,0 +1,106 @@
+import { describe, it, expect, vi } from 'vitest'
+import { atom, createStore } from 'jotai/vanilla'
+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'
+    )
+  })
+})