fix: delay all state updates after mount (#158)

Author: lukasbach

.changeset/heavy-carpets-smash.md

@@ -0,0 +1,6 @@
+---
+"@headless-tree/react": minor
+"@headless-tree/core": minor
+---
+
+all state updates (like setSelectedItems) will not propagate while the component is unmounted. This happened before for `tree.setState()` calls directly, but not individual state atoms like `setSelectedItems`. When calling `createTree()` directly (instead of `useTree()`), `tree.setMounted(true)` needs to be called once after mount. No changes are necessary when using the React-based `useTree()` integration. (#158)

packages/core/src/core/create-tree.ts

@@ -11,6 +11,7 @@ import { treeFeature } from "../features/tree/feature";
 import { ItemMeta } from "../features/tree/types";
 import { buildStaticInstance } from "./build-static-instance";
 import { throwError } from "../utilities/errors";
+import type { TreeDataRef } from "../features/main/types";
 
 const verifyFeatures = (features: FeatureImplementation[] | undefined) => {
   const loadedFeatures = features?.map((feature) => feature.key);
@@ -161,17 +162,34 @@ export const createTree = <T>(
         config.setState?.(state); // TODO this cant be right... This doesnt allow external state updates
         // TODO this is never used, remove
       },
+      setMounted: ({}, isMounted) => {
+        const ref = treeDataRef as TreeDataRef;
+        treeDataRef.current.isMounted = isMounted;
+        if (isMounted) {
+          ref.waitingForMount?.forEach((cb) => cb());
+          ref.waitingForMount = [];
+        }
+      },
       applySubStateUpdate: <K extends keyof TreeState<any>>(
         {},
         stateName: K,
         updater: Updater<TreeState<T>[K]>,
       ) => {
-        state[stateName] =
-          typeof updater === "function" ? updater(state[stateName]) : updater;
-        const externalStateSetter = config[
-          stateHandlerNames[stateName]
-        ] as Function;
-        externalStateSetter?.(state[stateName]);
+        const apply = () => {
+          state[stateName] =
+            typeof updater === "function" ? updater(state[stateName]) : updater;
+          const externalStateSetter = config[
+            stateHandlerNames[stateName]
+          ] as Function;
+          externalStateSetter?.(state[stateName]);
+        };
+        const ref = treeDataRef.current as TreeDataRef;
+        if (ref.isMounted) {
+          apply();
+        } else {
+          ref.waitingForMount ??= [];
+          ref.waitingForMount.push(apply);
+        }
       },
       // TODO rebuildSubTree: (itemId: string) => void;
       rebuildTree: () => {

packages/core/src/features/main/types.ts

@@ -10,6 +10,11 @@ import {
 } from "../../types/core";
 import { ItemMeta } from "../tree/types";
 
+export interface TreeDataRef {
+  isMounted?: boolean;
+  waitingForMount?: (() => void)[];
+}
+
 export type InstanceTypeMap = {
   itemInstance: ItemInstance<any>;
   treeInstance: TreeInstance<any>;
@@ -51,6 +56,8 @@ export type MainFeatureDef<T = any> = {
     rebuildTree: () => void;
     /** @deprecated Experimental feature, might get removed or changed in the future. */
     scheduleRebuildTree: () => void;
+    /** @internal */
+    setMounted: (isMounted: boolean) => void;
   };
   itemInstance: {
     registerElement: (element: HTMLElement | null) => void;

packages/core/src/test-utils/test-tree.ts

@@ -78,6 +78,7 @@ export class TestTree<T = string> {
   get instance() {
     if (!this.treeInstance) {
       this.treeInstance = createTree(this.config);
+      this.treeInstance.setMounted(true);
       this.treeInstance.rebuildTree();
     }
     return this.treeInstance;

packages/react/src/use-tree.tsx

@@ -1,42 +1,20 @@
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useState } from "react";
 import { TreeConfig, TreeState, createTree } from "@headless-tree/core";
 
-/* A bug with NextJS was reported in the past where very quick state updates (i.e. data
- * loader returning withing milliseconds) will cause the state update to occur before
- * mount, resulting in a console warning. This alleviates this.
- * We should monitor if this remains a problem in the future, maybe we can eventually
- * remove this workaround...
- */
-const useApplyAfterMount = () => {
-  const isMounted = useRef(false);
-  const callbacks = useRef<(() => void)[]>([]);
-
-  useEffect(() => {
-    isMounted.current = true;
-    callbacks.current.forEach((callback) => callback());
-  }, []);
-
-  const apply = (callback: () => void) => {
-    if (isMounted.current) {
-      callback();
-    } else {
-      callbacks.current.push(callback);
-    }
-  };
-
-  return apply;
-};
-
 export const useTree = <T,>(config: TreeConfig<T>) => {
-  const apply = useApplyAfterMount();
   const [tree] = useState(() => ({ current: createTree(config) }));
   const [state, setState] = useState<Partial<TreeState<T>>>(() =>
     tree.current.getState(),
   );
 
   useEffect(() => {
+    (tree.current as any).setMounted(true);
     tree.current.rebuildTree();
-  }, [tree]); // runs only once after mount
+    return () => {
+      // eslint-disable-next-line react-hooks/exhaustive-deps
+      (tree.current as any).setMounted(false);
+    };
+  }, [tree]);
 
   tree.current.setConfig((prev) => ({
     ...prev,
@@ -46,10 +24,8 @@ export const useTree = <T,>(config: TreeConfig<T>) => {
       ...config.state,
     },
     setState: (state) => {
-      apply(() => {
-        setState(state);
-        config.setState?.(state);
-      });
+      setState(state);
+      config.setState?.(state);
     },
   }));