feat(useTreeLoader): get more data about the tree

Author: sbesselsen

README.md

@@ -16,6 +16,7 @@ With `useTree` you can focus on how to render your tree structure and forget abo
   - [useTreeContent()](#usetreecontent)
   - [useTreeLoader()](#usetreeloader)
 - [Rendering a tree](#rendering-a-tree)
+  - [Type RootTree<T>](#type-roottreet)
   - [Type Tree<T>](#type-treet)
   - [Type TreeNode<T>](#type-treenodet)
 - [Typescript](#typescript)
@@ -89,7 +90,7 @@ This interface describes the current display state of a tree. It contains two pr
 ### useTreeController()
 Get an object that lets you control the current tree (from context). Provides these methods:
 
-* `updateState(updater: (oldState: TreeState) => TreeState): void`: update the tree state with your own updater function.
+* `updateState(updater: (oldState: TreeState, tree: RootTree<T>) => TreeState): void`: update the tree state with your own updater function.
 * `setExpanded(id: string, expanded?: boolean): void`: expanded or collapse a tree node.
 * `toggleExpanded(id: string): void`: toggle the expanded state of a tree node.
 * `setActiveId(id: string | null): void`: set which (if any) tree node is active.
@@ -105,13 +106,17 @@ Same as `useTreeController()`, but only for one tree node. You should, for insta
 When used inside the context of a tree, returns the current data of the tree. This can be used if you have several components nested inside your `<TreeContainer>` that all want to display the tree (or parts thereof).
 
 ### useTreeLoader()
-Use this if you want full control and don't want to use `TreeContainer`. This hook takes a `TreeSource` and a `TreeState` and returns the most up-to-date tree data structure. It will load data from the source if necessary and re-render as that data comes in.
+Use this if you want full control and don't want to use `TreeContainer`. This hook takes a `TreeSource` and a `TreeState` and returns the most up-to-date tree data structure with type `RootTree<T>`. It will load data from the source if necessary and re-render as that data comes in.
 
 ## Rendering a tree
 When rendering a tree through the `rootElement` of `TreeContainer` or by passing the result of `useTreeLoader()` directly to your component, your component should accept these data types. We will assume that your `TreeSource` is of type `TreeSource<T>` where `T` is your own type that contains your own properties for tree nodes.
 
 Root element properties:
-* `tree: Tree<T>`
+* `tree: RootTree<T>`
+
+### Type RootTree<T>
+* All properties from `Tree<T>`
+* `allNodes: {[k: string]: TreeNode<T>}`: all currently loaded tree nodes, indexed by ID.
 
 ### Type Tree<T>
 * `isLoading: boolean`: whether the items are still being loaded.

src/TreeContainer.tsx

@@ -15,14 +15,14 @@ interface TreeContainerProps<T> {
 
 export function TreeContainer<T>(props: PropsWithChildren<TreeContainerProps<T>>, context?: any): ReactElement | null {
     const { source, defaultState, state, onStateChange, rootElement, children } = props;
-    const controller = useRef<TreeController>(treeControllerFromUpdateState(noopUpdateState));
+    const controller = useRef<TreeController<unknown>>(treeControllerFromUpdateState(noopUpdateState));
     const [innerState, setInnerState] = useBinding(defaultState, state, onStateChange, {});
 
     const tree = useTreeLoader(source, innerState);
 
     controller.current.updateState = useCallback((updater) => {
-        setInnerState(updater(innerState));
-    }, [innerState, setInnerState]);
+        setInnerState(updater(innerState, tree));
+    }, [innerState, setInnerState, tree]);
 
     return (
         <TreeContentContext.Provider value={tree}>

src/types.ts

@@ -26,3 +26,7 @@ export type TreeNode<T> = TreeSourceNode<T> & {
 };
 
 export type Tree<T> = LoadableArray<TreeNode<T>>;
+
+export type RootTree<T> = Tree<T> & {
+    allNodes: {[k: string]: TreeNode<T>};
+};

src/use-tree-controller.ts

@@ -1,8 +1,10 @@
 import { createContext, useContext, useMemo } from 'react';
-import { TreeSourceNode, TreeState } from './types';
+import { RootTree, TreeSourceNode, TreeState } from './types';
 
-export interface TreeController {
-    updateState(f: (st: TreeState) => TreeState): void;
+type TreeStateUpdater<T> = (state: TreeState, tree: RootTree<T>) => TreeState;
+
+export interface TreeController<T> {
+    updateState(f: TreeStateUpdater<T>): void;
     toggleExpanded(id: string): void;
     setExpanded(id: string, expanded?: boolean): void;
     setActiveId(id: string | null): void;
@@ -14,8 +16,8 @@ export interface TreeNodeController {
     setActive(active?: boolean): void;
 }
 
-export function treeControllerFromUpdateState(updateState: (f: (st: TreeState) => TreeState) => void): TreeController {
-    const obj: Partial<TreeController> = {
+export function treeControllerFromUpdateState<T>(updateState: (f: TreeStateUpdater<T>) => void): TreeController<T> {
+    const obj: Partial<TreeController<T>> = {
         updateState,
     };
     obj.setExpanded = (id: string, expanded?: boolean) => {
@@ -33,16 +35,16 @@ export function treeControllerFromUpdateState(updateState: (f: (st: TreeState) =
     obj.setActiveId = (id: string | null) => {
         obj.updateState!((st) => ({ ...st, activeId: id }));
     };
-    return obj as TreeController;
+    return obj as TreeController<T>;
 }
 
-export const noopUpdateState = (f: (st: TreeState) => TreeState) => { /* noop */ };
+export const noopUpdateState = <T, >(f: TreeStateUpdater<T>) => { /* noop */ };
 
 export const noopTreeController = treeControllerFromUpdateState(noopUpdateState);
 
-export const TreeControllerContext = createContext<TreeController>(noopTreeController);
+export const TreeControllerContext = createContext<TreeController<unknown>>(noopTreeController);
 
-export function useTreeController(): TreeController {
+export function useTreeController(): TreeController<unknown> {
     return useContext(TreeControllerContext);
 }
 

src/use-tree-loader.ts

@@ -1,5 +1,5 @@
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import { LoadableArray, Tree, TreeNode, TreeSource, TreeSourceNode, TreeState } from './types';
+import { LoadableArray, RootTree, TreeNode, TreeSource, TreeSourceNode, TreeState } from './types';
 
 interface StringMap<V> {
     [k: string]: V;
@@ -29,7 +29,7 @@ function valuesEqual(arr1: unknown[], arr2: unknown[]): boolean {
     return true;
 }
 
-export function useTreeLoader<T>(source: TreeSource<T>, state: TreeState): Tree<T> {
+export function useTreeLoader<T>(source: TreeSource<T>, state: TreeState): RootTree<T> {
     const [rootNodes, setRootNodes] = useState<LoadableArray<TreeSourceNode<T>>>({ isLoading: true, items: [] });
     const [children, setChildren] = useState<StringMap<LoadableArray<TreeSourceNode<T>>>>({});
     const [trails, setTrails] = useState<StringMap<Array<TreeSourceNode<T>>>>({});
@@ -135,6 +135,7 @@ export function useTreeLoader<T>(source: TreeSource<T>, state: TreeState): Tree<
         return {
             items: rootNodes.items.map(buildOutputNode),
             isLoading: rootNodes.isLoading,
+            allNodes: statefulNodes.current,
         };
     }, [activeId, expandedIds, rootNodes, children, activeTrailIds, statefulNodes]);
 }