Batch calls to setNodes and addNodes for better perf.

This commit is contained in:
Hayleigh Thompson
2024-01-19 16:09:51 +00:00
parent 5f5099d1fc
commit a25854646a

View File

@@ -46,122 +46,109 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
return edges.find((e) => e.id === id) as EdgeType;
}, []);
const setNodesQueue = useRef<(NodeType[] | ((nodes: NodeType[]) => NodeType[]))[]>([]);
const setNodesHasStateReplacement = useRef(false);
const [setNodesShouldFlushNodesQueue, setSetNodesShouldFlushNodesQueue] = useState(false);
const setNodes = useCallback<Instance.SetNodes<NodeType>>((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<SetElementsQueue>({ 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<Node[]>();
// const setNodesTimeout = useRef<ReturnType<typeof setTimeout>>();
// const setNodes = useCallback<Instance.SetNodes<NodeType>>((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<Edge[]>();
const setEdgesTimeout = useRef<ReturnType<typeof setTimeout>>();
const setEdges = useCallback<Instance.SetEdges<EdgeType>>((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<Instance.SetNodes<NodeType>>((payload) => {
setElementsQueue.current.nodes.push(payload);
setShouldFlushQueue(true);
}, []);
const setEdges = useCallback<Instance.SetEdges<EdgeType>>((payload) => {
setElementsQueue.current.edges.push(payload);
setShouldFlushQueue(true);
}, []);
const addNodes = useCallback<Instance.AddNodes<NodeType>>((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<NodeType>));
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<Instance.AddEdges<EdgeType>>((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<EdgeType>));
onEdgesChange(changes);
}
setElementsQueue.current.edges.push((edges) => [...edges, ...newEdges]);
setShouldFlushQueue(true);
}, []);
const toObject = useCallback<Instance.ToObject<NodeType, EdgeType>>(() => {