Function useTreeData

packages/react-stately/src/data/useTreeData.ts:152–474

(options: TreeOptions<T>)

 * update the data over time.
 */
export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData<T> {
  let {
    initialItems = [],
    initialSelectedKeys,
    getKey = (item: any) => item.id ?? item.key,
    getChildren = (item: any) => item.children
  } = options;

  // We only want to compute this on initial render.
  let [tree, setItems] = useState<TreeDataState<T>>(() => buildTree(initialItems, new Map()));
  let {items, nodeMap} = tree;

  let [selectedKeys, setSelectedKeys] = useState(new Set<Key>(initialSelectedKeys || []));

  function buildTree(
    initialItems: T[] | null = [],
    map: Map<Key, TreeNode<T>>,
    parentKey?: Key | null
  ) {
    if (initialItems == null) {
      initialItems = [];
    }
    return {
      items: initialItems.map(item => {
        let node: TreeNode<T> = {
          key: getKey(item),
          parentKey: parentKey ?? null,
          value: item,
          children: null
        };

        node.children = buildTree(getChildren(item), map, node.key).items;
        map.set(node.key, node);
        return node;
      }),
      nodeMap: map
    };
  }

  function updateTree(
    items: TreeNode<T>[],
    key: Key | null,
    update: (node: TreeNode<T>) => TreeNode<T> | null,
    originalMap: Map<Key, TreeNode<T>>
  ) {
    let node = key == null ? null : originalMap.get(key);
    if (node == null) {
      return {items, nodeMap: originalMap};
    }
    let map = new Map<Key, TreeNode<T>>(originalMap);

    // Create a new node. If null, then delete the node, otherwise replace.
    let newNode = update(node);
    if (newNode == null) {
      deleteNode(node, map);
    } else {
      addNode(newNode, map);
    }