added parentLookup

packages/react/src/components/NodeWrapper/index.tsx

@@ -42,7 +42,7 @@ export function NodeWrapper<NodeType extends Node>({
   nodeOrigin,
   onError,
 }: NodeWrapperProps<NodeType>) {
-  const { node, positionAbsoluteX, positionAbsoluteY, zIndex, isParent } = useStore((s) => {
+  const { node, positionAbsoluteX, positionAbsoluteY, internals } = useStore((s) => {
     const node = s.nodeLookup.get(id)! as InternalNode<NodeType>;
 
     const positionAbsolute = nodeExtent
@@ -55,8 +55,9 @@ export function NodeWrapper<NodeType extends Node>({
       // so we we need to force a re-render when some change
       positionAbsoluteX: positionAbsolute.x,
       positionAbsoluteY: positionAbsolute.y,
-      zIndex: node.internals.z,
-      isParent: node.internals.isParent,
+      internals: node.internals,
+      // zIndex: node.internals.z,
+      // isParent: node.internals.isParent,
     };
   }, shallow);
 
@@ -213,14 +214,14 @@ export function NodeWrapper<NodeType extends Node>({
         {
           selected: node.selected,
           selectable: isSelectable,
-          parent: isParent,
+          parent: internals.isParent,
           draggable: isDraggable,
           dragging,
         },
       ])}
       ref={nodeRef}
       style={{
-        zIndex,
+        zIndex: internals.z,
         transform: `translate(${positionAbsoluteOrigin.x}px,${positionAbsoluteOrigin.y}px)`,
         pointerEvents: hasPointerEvents ? 'all' : 'none',
         visibility: initialized ? 'visible' : 'hidden',
@@ -254,7 +255,7 @@ export function NodeWrapper<NodeType extends Node>({
           targetPosition={node.targetPosition}
           dragging={dragging}
           dragHandle={node.dragHandle}
-          zIndex={zIndex}
+          zIndex={internals.z}
           {...nodeDimensions}
         />
       </Provider>

packages/react/src/store/index.ts

@@ -38,14 +38,14 @@ const createRFStore = ({
     (set, get) => ({
       ...getInitialState({ nodes, edges, width, height, fitView, defaultNodes, defaultEdges }),
       setNodes: (nodes: Node[]) => {
-        const { nodeLookup, nodeOrigin, elevateNodesOnSelect } = get();
+        const { nodeLookup, parentLookup, nodeOrigin, elevateNodesOnSelect } = get();
         // setNodes() is called exclusively in response to user actions:
         // - either when the `<ReactFlow nodes>` prop is updated in the controlled ReactFlow setup,
         // - or when the user calls something like `reactFlowInstance.setNodes()` in an uncontrolled ReactFlow setup.
         //
         // When this happens, we take the note objects passed by the user and extend them with fields
         // relevant for internal React Flow operations.
-        adoptUserNodes(nodes, nodeLookup, { nodeOrigin, elevateNodesOnSelect });
+        adoptUserNodes(nodes, nodeLookup, parentLookup, { nodeOrigin, elevateNodesOnSelect });
 
         set({ nodes });
       },

packages/react/src/store/initialState.ts

@@ -29,13 +29,14 @@ const getInitialState = ({
   fitView?: boolean;
 } = {}): ReactFlowStore => {
   const nodeLookup = new Map();
+  const parentLookup = new Map();
   const connectionLookup = new Map();
   const edgeLookup = new Map();
   const storeEdges = defaultEdges ?? edges ?? [];
   const storeNodes = defaultNodes ?? nodes ?? [];
 
   updateConnectionLookup(connectionLookup, edgeLookup, storeEdges);
-  adoptUserNodes(storeNodes, nodeLookup, {
+  adoptUserNodes(storeNodes, nodeLookup, parentLookup, {
     nodeOrigin: [0, 0],
     elevateNodesOnSelect: false,
   });
@@ -59,6 +60,7 @@ const getInitialState = ({
     transform,
     nodes: storeNodes,
     nodeLookup,
+    parentLookup,
     edges: storeEdges,
     edgeLookup,
     connectionLookup,

packages/react/src/types/store.ts

@@ -54,6 +54,7 @@ export type ReactFlowStore<NodeType extends Node = Node, EdgeType extends Edge =
   transform: Transform;
   nodes: NodeType[];
   nodeLookup: NodeLookup<InternalNode<NodeType>>;
+  parentLookup: Map<string, InternalNode<NodeType>[]>;
   edges: Edge[];
   edgeLookup: EdgeLookup<EdgeType>;
   connectionLookup: ConnectionLookup;

packages/svelte/src/lib/store/initial-store.ts

@@ -81,7 +81,8 @@ export const getInitialStore = ({
   fitView?: boolean;
 }) => {
   const nodeLookup: NodeLookup = new Map();
-  adoptUserNodes(nodes, nodeLookup, {
+  const parentLookup = new Map();
+  adoptUserNodes(nodes, nodeLookup, parentLookup, {
     nodeOrigin: [0, 0],
     elevateNodesOnSelect: false,
     checkEquality: false
@@ -104,8 +105,9 @@ export const getInitialStore = ({
 
   return {
     flowId: writable<string | null>(null),
-    nodes: createNodesStore(nodes, nodeLookup),
+    nodes: createNodesStore(nodes, nodeLookup, parentLookup),
     nodeLookup: readable<NodeLookup<InternalNode>>(nodeLookup),
+    parentLookup: readable<Map<string, InternalNode[]>>(parentLookup),
     edgeLookup: readable<EdgeLookup<Edge>>(edgeLookup),
     visibleNodes: readable<InternalNode[]>([]),
     edges: createEdgesStore(edges, connectionLookup, edgeLookup),

packages/svelte/src/lib/store/utils.ts

@@ -127,7 +127,8 @@ export type NodeStoreOptions = {
 // The user only passes in relative positions, so we need to calculate the absolute positions based on the parent nodes.
 export const createNodesStore = (
   nodes: Node[],
-  nodeLookup: NodeLookup<InternalNode>
+  nodeLookup: NodeLookup<InternalNode>,
+  parentLookup: Map<string, InternalNode[]>
 ): {
   subscribe: (this: void, run: Subscriber<Node[]>) => Unsubscriber;
   update: (this: void, updater: Updater<Node[]>) => void;
@@ -141,7 +142,7 @@ export const createNodesStore = (
   let elevateNodesOnSelect = true;
 
   const _set = (nds: Node[]): Node[] => {
-    adoptUserNodes(nds, nodeLookup, {
+    adoptUserNodes(nds, nodeLookup, parentLookup, {
       elevateNodesOnSelect,
       defaults,
       checkEquality: false

packages/system/src/utils/store.ts

@@ -27,7 +27,7 @@ export function updateAbsolutePositions<NodeType extends NodeBase>(
     elevateNodesOnSelect: true,
     defaults: {},
   },
-  parentNodeIds?: Set<string>
+  parentLookup?: Map<string, InternalNodeBase<NodeType>[]>
 ) {
   const selectedNodeZ: number = options?.elevateNodesOnSelect ? 1000 : 0;
 
@@ -38,7 +38,7 @@ export function updateAbsolutePositions<NodeType extends NodeBase>(
       throw new Error(`Parent node ${parentId} not found`);
     }
 
-    if (parentId || node.internals.isParent || parentNodeIds?.has(id)) {
+    if (parentId || node.internals.isParent || parentLookup?.has(id)) {
       const parentNode = parentId ? nodeLookup.get(parentId) : null;
       const { x, y, z } = calculateXYZPosition(
         node,
@@ -47,20 +47,21 @@ export function updateAbsolutePositions<NodeType extends NodeBase>(
           ...node.position,
           z: (isNumeric(node.zIndex) ? node.zIndex : 0) + (node.selected ? selectedNodeZ : 0),
         },
-        parentNode?.origin || options.nodeOrigin
+        parentNode?.origin ?? options.nodeOrigin
       );
 
       const currPosition = node.internals.positionAbsolute;
       const positionChanged = x !== currPosition.x || y !== currPosition.y;
 
-      node.internals.positionAbsolute = positionChanged ? { x, y } : currPosition;
-      node.internals.z = z;
+      node.internals = {
+        ...node.internals,
+        positionAbsolute: positionChanged ? { x, y } : currPosition,
+        z,
+      };
 
-      if (parentNodeIds !== undefined) {
-        node.internals.isParent = !!parentNodeIds?.has(id);
+      if (parentLookup !== undefined) {
+        node.internals.isParent = !!parentLookup.has(id);
       }
-
-      nodeLookup.set(id, node);
     }
   }
 }
@@ -75,6 +76,7 @@ type UpdateNodesOptions<NodeType extends NodeBase> = {
 export function adoptUserNodes<NodeType extends NodeBase>(
   nodes: NodeType[],
   nodeLookup: Map<string, InternalNodeBase<NodeType>>,
+  parentLookup: Map<string, InternalNodeBase<NodeType>[]>,
   options: UpdateNodesOptions<NodeType> = {
     nodeOrigin: [0, 0] as NodeOrigin,
     elevateNodesOnSelect: true,
@@ -84,20 +86,19 @@ export function adoptUserNodes<NodeType extends NodeBase>(
 ) {
   const tmpLookup = new Map(nodeLookup);
   nodeLookup.clear();
+  parentLookup.clear();
+
   const selectedNodeZ: number = options?.elevateNodesOnSelect ? 1000 : 0;
-  const parentNodeIds = new Set<string>();
+  // const parentNodeIds = new Set<string>();
 
   nodes.forEach((userNode) => {
     const currentStoreNode = tmpLookup.get(userNode.id);
 
-    if (userNode.parentId) {
-      parentNodeIds.add(userNode.parentId);
-    }
-
+    let internalNode = currentStoreNode!;
     if (options.checkEquality && userNode === currentStoreNode?.internals.userNode) {
       nodeLookup.set(userNode.id, currentStoreNode);
     } else {
-      nodeLookup.set(userNode.id, {
+      internalNode = {
         ...options.defaults,
         ...userNode,
         measured: {
@@ -111,12 +112,22 @@ export function adoptUserNodes<NodeType extends NodeBase>(
           userNode,
           isParent: false,
         },
-      });
+      };
+      nodeLookup.set(userNode.id, internalNode);
+    }
+
+    if (userNode.parentId) {
+      const childNodes = parentLookup.get(userNode.parentId);
+      if (childNodes) {
+        childNodes.push(internalNode);
+      } else {
+        parentLookup.set(userNode.parentId, [internalNode]);
+      }
     }
   });
 
-  if (parentNodeIds.size > 0) {
-    updateAbsolutePositions(nodeLookup, options, parentNodeIds);
+  if (parentLookup.size > 0) {
+    updateAbsolutePositions(nodeLookup, options, parentLookup);
   }
 }