diff --git a/packages/react/src/hooks/useReactFlow.ts b/packages/react/src/hooks/useReactFlow.ts index 097c78b6..18de6fc6 100644 --- a/packages/react/src/hooks/useReactFlow.ts +++ b/packages/react/src/hooks/useReactFlow.ts @@ -46,122 +46,109 @@ export function useReactFlow e.id === id) as EdgeType; }, []); - const setNodesQueue = useRef<(NodeType[] | ((nodes: NodeType[]) => NodeType[]))[]>([]); - const setNodesHasStateReplacement = useRef(false); - const [setNodesShouldFlushNodesQueue, setSetNodesShouldFlushNodesQueue] = useState(false); - const setNodes = useCallback>((payload) => { - if (typeof payload === 'function' && !setNodesHasStateReplacement.current) { - setNodesQueue.current.push(payload); - setSetNodesShouldFlushNodesQueue(true); - } else { - setNodesQueue.current = [payload]; - setNodesHasStateReplacement.current = true; - setSetNodesShouldFlushNodesQueue(true); - } - }, []); + type SetElementsQueue = { + nodes: (NodeType[] | ((nodes: NodeType[]) => NodeType[]))[]; + edges: (EdgeType[] | ((edges: EdgeType[]) => EdgeType[]))[]; + }; + // A reference of all the batched updates to process before the next render. We + // want a mutable reference here so multiple synchronous calls to `setNodes` etc + // can be batched together. + const setElementsQueue = useRef({ nodes: [], edges: [] }); + // Because we're using a ref above, we need some way to let React know when to + // actually process the queue. We flip this bit of state to `true` any time we + // mutate the queue and then flip it back to `false` after flushing the queue. + const [shouldFlushQueue, setShouldFlushQueue] = useState(false); + + // Layout effects are guaranteed to run before the next render which means we + // shouldn't run into any issues with stale state or weird issues that come from + // rendering things one frame later than expected (we used to use `setTimeout`). useLayoutEffect(() => { - if (!setNodesShouldFlushNodesQueue) { - setNodesQueue.current = []; - setNodesHasStateReplacement.current = false; - + // Because we need to flip the state back to false after flushing, this should + // trigger the hook again (!). If the hook is being run again we know that any + // updates should have been processed by now and we can safely clear the queue + // and bail early. + if (!shouldFlushQueue) { + setElementsQueue.current = { nodes: [], edges: [] }; return; } - const { nodes = [], setNodes, hasDefaultNodes, onNodesChange, nodeLookup } = store.getState(); - const nextNodes = setNodesQueue.current.reduce( - (prev: NodeType[], payload) => (typeof payload === 'function' ? payload(prev) : payload), - nodes as NodeType[] - ); + if (setElementsQueue.current.nodes.length) { + const { nodes = [], setNodes, hasDefaultNodes, onNodesChange, nodeLookup } = store.getState(); - if (hasDefaultNodes) { - setNodes(nextNodes); - } else if (onNodesChange) { - const changes: NodeChange[] = getElementsDiffChanges({ items: nextNodes, lookup: nodeLookup }); - onNodesChange(changes); - } - - setNodesQueue.current = []; - setNodesHasStateReplacement.current = false; - setSetNodesShouldFlushNodesQueue(false); - }, [setNodesShouldFlushNodesQueue]); - - // // this is used to handle multiple syncronous setNodes calls - // const setNodesData = useRef(); - // const setNodesTimeout = useRef>(); - // const setNodes = useCallback>((payload) => { - // const { nodes = [], setNodes, hasDefaultNodes, onNodesChange, nodeLookup } = store.getState(); - // const nextNodes = typeof payload === 'function' ? payload((setNodesData.current as NodeType[]) || nodes) : payload; - - // setNodesData.current = nextNodes; - - // if (setNodesTimeout.current) { - // clearTimeout(setNodesTimeout.current); - // } - - // // if there are multiple synchronous setNodes calls, we only want to call onNodesChange once - // // for this, we use a timeout to wait for the last call and store updated nodes in setNodesData - // // this is not perfect, but should work in most cases - // setNodesTimeout.current = setTimeout(() => { - // if (hasDefaultNodes) { - // setNodes(nextNodes); - // } else if (onNodesChange) { - // const changes: NodeChange[] = getElementsDiffChanges({ items: setNodesData.current, lookup: nodeLookup }); - // onNodesChange(changes); - // } - - // setNodesData.current = undefined; - // }, 0); - // }, []); - - // this is used to handle multiple syncronous setEdges calls - const setEdgesData = useRef(); - const setEdgesTimeout = useRef>(); - const setEdges = useCallback>((payload) => { - const { edges = [], setEdges, hasDefaultEdges, onEdgesChange, edgeLookup } = store.getState(); - const nextEdges = typeof payload === 'function' ? payload((setEdgesData.current as EdgeType[]) || edges) : payload; - - setEdgesData.current = nextEdges; - - if (setEdgesTimeout.current) { - clearTimeout(setEdgesTimeout.current); - } - - setEdgesTimeout.current = setTimeout(() => { - if (hasDefaultEdges) { - setEdges(nextEdges); - } else if (onEdgesChange) { - const changes: EdgeChange[] = getElementsDiffChanges({ items: nextEdges, lookup: edgeLookup }); - onEdgesChange(changes); + // This is essentially an `Array.reduce` in imperative clothing. Processing + // this queue is a relatively hot path so we'd like to avoid the overhead of + // array methods where we can. + let next = nodes as NodeType[]; + for (const payload of setElementsQueue.current.nodes) { + next = typeof payload === 'function' ? payload(next) : payload; } - setEdgesData.current = undefined; - }, 0); + if (hasDefaultNodes) { + setNodes(next); + } else if (onNodesChange) { + onNodesChange( + getElementsDiffChanges({ + items: next, + lookup: nodeLookup, + }) + ); + } + + setElementsQueue.current.nodes = []; + } + + if (setElementsQueue.current.edges.length) { + const { edges = [], setEdges, hasDefaultEdges, onEdgesChange, edgeLookup } = store.getState(); + + let next = edges as EdgeType[]; + for (const payload of setElementsQueue.current.edges) { + next = typeof payload === 'function' ? payload(next) : payload; + } + + if (hasDefaultEdges) { + setEdges(next); + } else if (onEdgesChange) { + onEdgesChange( + getElementsDiffChanges({ + items: next, + lookup: edgeLookup, + }) + ); + } + + setElementsQueue.current.edges = []; + } + + // Beacuse we're using reactive state to trigger this effect, we need to flip + // it back to false. + setShouldFlushQueue(false); + }, [shouldFlushQueue]); + + const setNodes = useCallback>((payload) => { + setElementsQueue.current.nodes.push(payload); + setShouldFlushQueue(true); + }, []); + + const setEdges = useCallback>((payload) => { + setElementsQueue.current.edges.push(payload); + setShouldFlushQueue(true); }, []); const addNodes = useCallback>((payload) => { - const nodes = Array.isArray(payload) ? payload : [payload]; - const { nodes: currentNodes, hasDefaultNodes, onNodesChange, setNodes } = store.getState(); + const newNodes = Array.isArray(payload) ? payload : [payload]; - if (hasDefaultNodes) { - const nextNodes = [...currentNodes, ...nodes]; - setNodes(nextNodes); - } else if (onNodesChange) { - const changes = nodes.map((node) => ({ item: node, type: 'add' } as NodeAddChange)); - onNodesChange(changes); - } + // Queueing a functional update means that we won't worry about other calls + // to `setNodes` that might happen elsewhere. + setElementsQueue.current.nodes.push((nodes) => [...nodes, ...newNodes]); + setShouldFlushQueue(true); }, []); const addEdges = useCallback>((payload) => { - const nextEdges = Array.isArray(payload) ? payload : [payload]; - const { edges = [], setEdges, hasDefaultEdges, onEdgesChange } = store.getState(); + const newEdges = Array.isArray(payload) ? payload : [payload]; - if (hasDefaultEdges) { - setEdges([...edges, ...nextEdges]); - } else if (onEdgesChange) { - const changes = nextEdges.map((edge) => ({ item: edge, type: 'add' } as EdgeAddChange)); - onEdgesChange(changes); - } + setElementsQueue.current.edges.push((edges) => [...edges, ...newEdges]); + setShouldFlushQueue(true); }, []); const toObject = useCallback>(() => {