Prevent query graph flickering during relayout

query-graphs/README.md

@@ -48,10 +48,10 @@ The library intentionally exposes low-level loaders (`json`, `xml`) as generic f
 ## The Renderer
 
 `QueryGraph` (`src/ui/QueryGraph.tsx`) is the top-level component rendering a `TreeDescription`.
-It assigns a stable id to every node, seeds the interaction store from each node's `expandedByDefault` flag, and observes on-screen node sizes with a `ResizeObserver`.
+It assigns a stable id to every node, creates a graph-local rendering store seeded from each node's `expandedByDefault` flag, and retains the node dimensions measured by react-flow.
 
 `tree-layout.ts` positions the tree with [`d3-flextree`](https://github.com/Klortho/d3-flextree) on top of `d3-hierarchy`, then translates the result into react-flow nodes and edges.
-Layout is driven by the **measured** DOM size of each node, so it runs in two passes: the first render uses a placeholder size and, once the `ResizeObserver` reports real dimensions, the tree re-lays-out with the correct sizes.
+Layout is driven by the **measured** DOM size of each node, so it runs in two passes: react-flow measures new nodes after their first render, then the tree re-lays-out with the correct sizes. Those measurements are retained in the controlled node objects so react-flow does not re-initialize them on every layout.
 Edge thickness is scaled from `edgeWidth`, and `crosslinks` are added as extra edges.
 
 `QueryNode` (`src/ui/QueryNode.tsx`) draws a single node.
@@ -83,12 +83,12 @@ These are the touches that make a plan readable at a glance:
 
 ### Interaction State
 
-`store.ts` is a [Zustand](https://github.com/pmndrs/zustand) store (with the `immer` and `devtools` middleware) holding all mutable view state, so React components stay purely declarative.
+Each `QueryGraph` owns a [Zustand](https://github.com/pmndrs/zustand) store holding its mutable rendering state, so multiple graphs do not interfere with one another.
 It tracks three things, and the distinction between the first two is the key subtlety:
 
 * `expandedNodes` — which nodes have their **property detail panel** open.
 * `expandedSubtrees` — which nodes reveal their **`collapsedChildren`** in the graph.
-* `nodeDimensions` — the measured head/body size of each node, fed back into layout.
+* `nodeDimensions` — react-flow's measurements, retained across controlled-node layout updates.
 
 ## Tech Debt
 

query-graphs/package.json

@@ -47,7 +47,6 @@
     "classcat": "^5.0.4",
     "d3-flextree": "^2.1.2",
     "d3-hierarchy": "^3.1.2",
-    "immer": "^11.1.18",
     "react": "^19.2.8",
     "react-dom": "^19.2.8",
     "zustand": "^5.0.0"

query-graphs/src/ui/QueryGraph.tsx

@@ -1,22 +1,27 @@
+import type {NodeChange} from "@xyflow/react";
 import {ReactFlow, MiniMap, Controls, ReactFlowProvider} from "@xyflow/react";
 import "@xyflow/react/dist/base.css";
 
 import {layoutTree} from "./tree-layout";
 import type {TreeDescription, TreeNode} from "../tree-description";
 import {allChildren, visitTreeNodes} from "../tree-description";
 import type {ReactNode} from "react";
-import {useMemo, useEffect, useRef} from "react";
+import {useCallback, useMemo} from "react";
 import {QueryNode} from "./QueryNode";
 import type {QueryGraphNode} from "./QueryNode";
 import {ColoredEdge} from "./ColoredEdge";
-import {useGraphRenderingStore} from "./store";
+import {createGraphRenderingStore, GraphRenderingStoreContext, useGraphRenderingStore} from "./store";
 import "./QueryGraph.css";
 
 interface QueryGraphProps {
     treeDescription: TreeDescription;
     children: ReactNode | ReactNode[];
 }
 
+interface QueryGraphInternalProps extends QueryGraphProps {
+    nodeIdMapping: Map<TreeNode, string>;
+}
+
 function minimapNodeColor(n: QueryGraphNode): string {
     if (n.data.nodeColor) return n.data.nodeColor;
     if (n.data.iconColor) return n.data.iconColor;
@@ -31,59 +36,27 @@ const edgeTypes = {
     colored: ColoredEdge,
 };
 
-function QueryGraphInternal({treeDescription, children}: QueryGraphProps) {
-    // Assign ids to all nodes
-    const nodeIdMapping = useMemo(() => {
-        let nextId = 0;
-        const nodeIds = new Map<TreeNode, string>();
-        visitTreeNodes(
-            treeDescription.root,
-            (d) => {
-                nodeIds.set(d, "" + nextId++);
-            },
-            allChildren,
-        );
-        return nodeIds;
-    }, [treeDescription]);
-
-    // Initialize our state using the correct "expandedByDefault" state
-    const initGraphStore = useGraphRenderingStore((s) => s.init);
-    useMemo(() => {
-        const expandedSubtrees = {};
-        visitTreeNodes(
-            treeDescription.root,
-            (n) => {
-                if (n.expandedByDefault) {
-                    expandedSubtrees[nodeIdMapping.get(n)!] = true;
-                }
-            },
-            allChildren,
-        );
-        initGraphStore(expandedSubtrees);
-    }, [treeDescription, initGraphStore, nodeIdMapping]);
-
-    // Create a ResizeObserver to keep track of the sizes of the nodes
-    const resizeObserverRef = useRef<ResizeObserver | undefined>(undefined);
+function QueryGraphInternal({treeDescription, children, nodeIdMapping}: QueryGraphInternalProps) {
+    // Keep React Flow's measurements in the controlled node objects. Dropping them when
+    // recomputing the layout makes React Flow repeatedly hide and re-initialize the nodes.
+    const nodeDimensions = useGraphRenderingStore((s) => s.nodeDimensions);
     const updateNodeDimensions = useGraphRenderingStore((s) => s.updateNodeDimensions);
-    const resizeObserver = useMemo(() => {
-        resizeObserverRef.current?.disconnect();
-        const observer = new ResizeObserver(updateNodeDimensions);
-        resizeObserverRef.current = observer;
-        return observer;
-    }, [updateNodeDimensions]);
-    useEffect(() => {
-        return () => {
-            resizeObserverRef.current?.disconnect();
-        };
-    }, []);
+    const onNodesChange = useCallback(
+        (changes: NodeChange<QueryGraphNode>[]) => {
+            const updates = changes.flatMap((change) => {
+                if (change.type !== "dimensions" || change.dimensions === undefined) return [];
+                return [[change.id, change.dimensions] as const];
+            });
+            updateNodeDimensions(updates);
+        },
+        [updateNodeDimensions],
+    );
 
-    // Layout the tree, using the actual measured sizes of the DOM nodes
-    const nodeDimensions = useGraphRenderingStore((s) => s.nodeDimensions);
-    const expandedNodes = useGraphRenderingStore((s) => s.expandedNodes);
+    // Layout the tree using the dimensions measured by React Flow itself.
     const expandedSubtrees = useGraphRenderingStore((s) => s.expandedSubtrees);
     const layout = useMemo(
-        () => layoutTree(treeDescription, nodeIdMapping, nodeDimensions, expandedNodes, expandedSubtrees, resizeObserver),
-        [treeDescription, nodeIdMapping, nodeDimensions, expandedNodes, expandedSubtrees, resizeObserver],
+        () => layoutTree(treeDescription, nodeIdMapping, nodeDimensions, expandedSubtrees),
+        [treeDescription, nodeIdMapping, nodeDimensions, expandedSubtrees],
     );
 
     return (
@@ -93,6 +66,7 @@ function QueryGraphInternal({treeDescription, children}: QueryGraphProps) {
             nodeOrigin={[0.5, 0]}
             nodeTypes={nodeTypes}
             edgeTypes={edgeTypes}
+            onNodesChange={onNodesChange}
             fitView
             minZoom={0.2}
             maxZoom={1.5}
@@ -109,10 +83,33 @@ function QueryGraphInternal({treeDescription, children}: QueryGraphProps) {
     );
 }
 
+function createGraphState(treeDescription: TreeDescription) {
+    let nextId = 0;
+    const nodeIdMapping = new Map<TreeNode, string>();
+    const expandedSubtrees: Record<string, boolean> = {};
+    visitTreeNodes(
+        treeDescription.root,
+        (node) => {
+            const id = "" + nextId++;
+            nodeIdMapping.set(node, id);
+            if (node.expandedByDefault) expandedSubtrees[id] = true;
+        },
+        allChildren,
+    );
+    return {
+        nodeIdMapping,
+        graphStore: createGraphRenderingStore(expandedSubtrees),
+    };
+}
+
 export function QueryGraph(props: QueryGraphProps) {
+    const {nodeIdMapping, graphStore} = useMemo(() => createGraphState(props.treeDescription), [props.treeDescription]);
+
     return (
         <ReactFlowProvider>
-            <QueryGraphInternal {...props} />
+            <GraphRenderingStoreContext.Provider value={graphStore}>
+                <QueryGraphInternal {...props} nodeIdMapping={nodeIdMapping} />
+            </GraphRenderingStoreContext.Provider>
         </ReactFlowProvider>
     );
 }

query-graphs/src/ui/QueryNode.css

@@ -46,9 +46,6 @@
 }
 
 .qg-graph-node-head {
-    /* We don't want the node to automatically stretch/shrink with its container
-     * because that would trigger the `ResizeObserver` to  often and would cause too
-     * many layout recomputations.*/
     width: max-content;
     margin: auto;
     text-align: center;

query-graphs/src/ui/QueryNode.tsx

@@ -1,33 +1,16 @@
-import type {ReactElement, MouseEvent, RefObject} from "react";
-import {memo, useCallback, useRef, useEffect} from "react";
+import type {ReactElement, MouseEvent} from "react";
+import {memo, useCallback} from "react";
 import type {Node, NodeProps} from "@xyflow/react";
 import {Handle, Position} from "@xyflow/react";
 import cc from "classcat";
 import type {TreeNode} from "../tree-description";
 import {NodeIcon} from "./NodeIcon";
 import "./QueryNode.css";
 import {useGraphRenderingStore} from "./store";
-import {assert} from "../assert";
 
-type NodeData = TreeNode & {resizeObserver: ResizeObserver};
-
-export type QueryGraphNode = Node<NodeData, "querynode">;
-
-function useResizeObservedRef<T extends Element>(resizeObserver: ResizeObserver): RefObject<T | null> {
-    const ref = useRef<T>(null);
-    useEffect(() => {
-        assert(ref.current !== null);
-        const currNode = ref.current;
-        resizeObserver.observe(currNode);
-        return () => resizeObserver.unobserve(currNode);
-    }, [resizeObserver]);
-    return ref;
-}
+export type QueryGraphNode = Node<TreeNode, "querynode">;
 
 function QueryNode({data, id}: NodeProps<QueryGraphNode>) {
-    const bodyRef = useResizeObservedRef<HTMLDivElement>(data.resizeObserver);
-    const headRef = useResizeObservedRef<HTMLDivElement>(data.resizeObserver);
-
     const expanded = useGraphRenderingStore((s) => s.expandedNodes[id]);
     const toggleNode = useGraphRenderingStore((s) => s.toggleExpandedNode);
     const subtreeExpanded = useGraphRenderingStore((s) => s.expandedSubtrees[id]);
@@ -93,17 +76,15 @@ function QueryNode({data, id}: NodeProps<QueryGraphNode>) {
         <>
             <Handle type="target" position={Position.Top} />
             <div className={nodeClassName} onClick={onClick}>
-                <div className="qg-graph-node-head" ref={headRef}>
+                <div className="qg-graph-node-head">
                     {colorBar(data.barsAbove, "above")}
                     <NodeIcon icon={data.icon} iconColor={data.iconColor} />
                     <div className="qg-graph-node-label" style={{background: data.nodeColor}}>
                         {data.name}
                     </div>
                 </div>
                 <div className="qg-graph-node-body-wrapper nowheel">
-                    <div ref={bodyRef} className="qg-graph-node-body">
-                        {children}
-                    </div>
+                    <div className="qg-graph-node-body">{children}</div>
                 </div>
                 {colorBar(data.barsBelow, "below")}
             </div>