Fix tree line drawing by using a callback ref

In a non-dev build (no initial double-renders), the tree lines would not be
rendered correctly from the parent of a binop to its first child, because the
first child would be rendered before the parent, and the parent's ref hadn't
been set yet at that time. Switched from a normal ref to a callback-based ref
with explicit React state update to make sure that the child gets to know about
the parent's (later) rendered div element.

Signed-off-by: Julius Volz <julius.volz@gmail.com>

web/ui/mantine-ui/src/pages/query/TreeNode.tsx

@@ -4,7 +4,6 @@ import {
   useEffect,
   useLayoutEffect,
   useMemo,
-  useRef,
   useState,
 } from "react";
 import ASTNode, { nodeType } from "../../promql/ast";
@@ -60,7 +59,7 @@ const TreeNode: FC<{
   node: ASTNode;
   selectedNode: { id: string; node: ASTNode } | null;
   setSelectedNode: (Node: { id: string; node: ASTNode } | null) => void;
-  parentRef?: React.RefObject<HTMLDivElement>;
+  parentEl?: HTMLDivElement | null;
   reportNodeState?: (childIdx: number, state: NodeState) => void;
   reverse: boolean;
   // The index of this node in its parent's children.
@@ -69,13 +68,21 @@ const TreeNode: FC<{
   node,
   selectedNode,
   setSelectedNode,
-  parentRef,
+  parentEl,
   reportNodeState,
   reverse,
   childIdx,
 }) => {
   const nodeID = useId();
-  const nodeRef = useRef<HTMLDivElement>(null);
+
+  // A normal ref won't work properly here because the ref's `current` property
+  // going from `null` to defined won't trigger a re-render of the child
+  // component, since it's not a React state update. So we manually have to
+  // create a state update using a callback ref. See also
+  // https://tkdodo.eu/blog/avoiding-use-effect-with-callback-refs
+  const [nodeEl, setNodeEl] = useState<HTMLDivElement | null>(null);
+  const nodeRef = useCallback((node: HTMLDivElement) => setNodeEl(node), []);
+
   const [connectorStyle, setConnectorStyle] = useState<CSSProperties>({
     borderColor:
       "light-dark(var(--mantine-color-gray-4), var(--mantine-color-dark-3))",
@@ -97,10 +104,10 @@ const TreeNode: FC<{
 
   // Select the node when it is mounted and it is the root of the tree.
   useEffect(() => {
-    if (parentRef === undefined) {
+    if (parentEl === undefined) {
       setSelectedNode({ id: nodeID, node: node });
     }
-  }, [parentRef, setSelectedNode, nodeID, node]);
+  }, [parentEl, setSelectedNode, nodeID, node]);
 
   // Deselect node when node is unmounted.
   useEffect(() => {
@@ -173,16 +180,18 @@ const TreeNode: FC<{
 
   // Update the size and position of tree connector lines based on the node's and its parent's position.
   useLayoutEffect(() => {
-    if (parentRef === undefined) {
+    if (parentEl === undefined) {
       // We're the root node.
       return;
     }
 
-    if (parentRef.current === null || nodeRef.current === null) {
+    if (parentEl === null || nodeEl === null) {
+      // Either of the two connected nodes hasn't been rendered yet.
       return;
     }
-    const parentRect = parentRef.current.getBoundingClientRect();
+
-    const nodeRect = nodeRef.current.getBoundingClientRect();
+    const parentRect = parentEl.getBoundingClientRect();
+    const nodeRect = nodeEl.getBoundingClientRect();
     if (reverse) {
       setConnectorStyle((prevStyle) => ({
         ...prevStyle,
@@ -202,7 +211,7 @@ const TreeNode: FC<{
         borderTopLeftRadius: undefined,
       }));
     }
-  }, [parentRef, reverse, nodeRef, setConnectorStyle]);
+  }, [parentEl, nodeEl, reverse, nodeRef, setConnectorStyle]);
 
   // Update the node info state based on the query result.
   useEffect(() => {
@@ -264,7 +273,7 @@ const TreeNode: FC<{
       pos="relative"
       align="center"
     >
-      {parentRef && (
+      {parentEl !== undefined && (
         // Connector line between this node and its parent.
         <Box pos="absolute" display="inline-block" style={connectorStyle} />
       )}
@@ -391,7 +400,7 @@ const TreeNode: FC<{
             node={children[0]}
             selectedNode={selectedNode}
             setSelectedNode={setSelectedNode}
-            parentRef={nodeRef}
+            parentEl={nodeEl}
             reverse={true}
             childIdx={0}
             reportNodeState={childReportNodeState}
@@ -403,7 +412,7 @@ const TreeNode: FC<{
             node={children[1]}
             selectedNode={selectedNode}
             setSelectedNode={setSelectedNode}
-            parentRef={nodeRef}
+            parentEl={nodeEl}
             reverse={false}
             childIdx={1}
             reportNodeState={childReportNodeState}
@@ -422,7 +431,7 @@ const TreeNode: FC<{
             node={child}
             selectedNode={selectedNode}
             setSelectedNode={setSelectedNode}
-            parentRef={nodeRef}
+            parentEl={nodeEl}
             reverse={false}
             childIdx={idx}
             reportNodeState={childReportNodeState}