fix(react): handle setNodes batching in provider #4147

This commit is contained in:
moklick
2024-04-25 12:57:24 +02:00
parent 40937c38ae
commit a7b6789ce0
7 changed files with 187 additions and 104 deletions

View File

@@ -0,0 +1,80 @@
import { createContext, ReactNode, useCallback, useContext, useMemo } from 'react';
import { EdgeChange, NodeChange } from '@xyflow/system';
import { useStoreApi } from '../../hooks/useStore';
import { getElementsDiffChanges } from '../../utils';
import { Queue, QueueItem } from './types';
import type { Edge, Node } from '../../types';
import { useQueue } from './useQueue';
const BatchContext = createContext<{
nodeQueue: Queue<Node>;
edgeQueue: Queue<Edge>;
} | null>(null);
export function BatchProvider<NodeType extends Node = Node, EdgeType extends Edge = Edge>({
children,
}: {
children: ReactNode;
}) {
const store = useStoreApi<NodeType, EdgeType>();
const nodeQueueHandler = useCallback((queueItems: QueueItem<NodeType>[]) => {
const { nodes = [], setNodes, hasDefaultNodes, onNodesChange, nodeLookup } = store.getState();
// 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 queueItems) {
next = typeof payload === 'function' ? payload(next) : payload;
}
if (hasDefaultNodes) {
setNodes(next);
} else if (onNodesChange) {
onNodesChange(
getElementsDiffChanges({
items: next,
lookup: nodeLookup,
}) as NodeChange<NodeType>[]
);
}
}, []);
const nodeQueue = useQueue<NodeType>(nodeQueueHandler);
const edgeQueueHandler = useCallback((queueItems: QueueItem<EdgeType>[]) => {
const { edges = [], setEdges, hasDefaultEdges, onEdgesChange, edgeLookup } = store.getState();
let next = edges as EdgeType[];
for (const payload of queueItems) {
next = typeof payload === 'function' ? payload(next) : payload;
}
if (hasDefaultEdges) {
setEdges(next);
} else if (onEdgesChange) {
onEdgesChange(
getElementsDiffChanges({
items: next,
lookup: edgeLookup,
}) as EdgeChange<EdgeType>[]
);
}
}, []);
const edgeQueue = useQueue<EdgeType>(edgeQueueHandler);
const value = useMemo(() => ({ nodeQueue, edgeQueue }), []);
return <BatchContext.Provider value={value}>{children}</BatchContext.Provider>;
}
export function useBatchContext() {
const batchContext = useContext(BatchContext);
if (!batchContext) {
throw new Error('useBatchContext must be used within a BatchProvider');
}
return batchContext;
}

View File

@@ -0,0 +1,7 @@
export type QueueItem<T> = T[] | ((items: T[]) => T[]);
export type Queue<T> = {
get: () => QueueItem<T>[];
reset: () => void;
push: (item: QueueItem<T>) => void;
};

View File

@@ -0,0 +1,44 @@
import { useState } from 'react';
import { createQueue } from './utils';
import { useIsomorphicLayoutEffect } from '../../hooks/useIsomorphicLayoutEffect';
import { QueueItem } from './types';
export function useQueue<T>(runQueue: (items: QueueItem<T>[]) => void) {
// 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 [shouldFlush, setShouldFlush] = useState(false);
// 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 [queue] = useState(() => createQueue<T>(() => setShouldFlush(true)));
// 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`).
useIsomorphicLayoutEffect(() => {
// 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 (!shouldFlush) {
queue.reset();
return;
}
const queueItems = queue.get();
if (queueItems.length) {
runQueue?.(queueItems);
queue.reset();
}
// Beacuse we're using reactive state to trigger this effect, we need to flip
// it back to false.
setShouldFlush(false);
}, [shouldFlush]);
return queue;
}

View File

@@ -0,0 +1,16 @@
import type { Queue, QueueItem } from './types';
export function createQueue<T>(cb: () => void): Queue<T> {
let queue: QueueItem<T>[] = [];
return {
get: () => queue,
reset: () => {
queue = [];
},
push: (item) => {
queue.push(item);
cb();
},
};
}

View File

@@ -2,6 +2,7 @@ import { useState, type ReactNode } from 'react';
import { Provider } from '../../contexts/StoreContext';
import { createStore } from '../../store';
import { BatchProvider } from '../BatchProvider';
import type { Node, Edge } from '../../types';
export type ReactFlowProviderProps = {
@@ -37,5 +38,9 @@ export function ReactFlowProvider({
})
);
return <Provider value={store}>{children}</Provider>;
return (
<Provider value={store}>
<BatchProvider>{children}</BatchProvider>
</Provider>
);
}

View File

@@ -1,4 +1,4 @@
import { useCallback, useMemo, useRef, useState } from 'react';
import { useCallback, useMemo } from 'react';
import {
evaluateAbsolutePosition,
getElementsToRemove,
@@ -10,9 +10,9 @@ import {
import useViewportHelper from './useViewportHelper';
import { useStoreApi } from './useStore';
import { useBatchContext } from '../components/BatchProvider';
import { isNode } from '../utils';
import type { ReactFlowInstance, Instance, Node, Edge, InternalNode } from '../types';
import { getElementsDiffChanges, isNode } from '../utils';
import { useIsomorphicLayoutEffect } from './useIsomorphicLayoutEffect';
/**
* Hook for accessing the ReactFlow instance.
@@ -26,6 +26,7 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
> {
const viewportHelper = useViewportHelper();
const store = useStoreApi();
const batchContext = useBatchContext();
const getNodes = useCallback<Instance.GetNodes<NodeType>>(
() => store.getState().nodes.map((n) => ({ ...n })) as NodeType[],
@@ -49,110 +50,39 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
const getEdge = useCallback<Instance.GetEdge<EdgeType>>((id) => store.getState().edgeLookup.get(id) as EdgeType, []);
type SetElementsQueue = {
nodes: (NodeType[] | ((nodes: NodeType[]) => NodeType[]))[];
edges: (EdgeType[] | ((edges: EdgeType[]) => EdgeType[]))[];
};
const setNodes = useCallback<Instance.SetNodes<NodeType>>(
(payload) => {
batchContext?.nodeQueue.push(payload as NodeType[]);
},
[batchContext]
);
// 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);
const setEdges = useCallback<Instance.SetEdges<EdgeType>>(
(payload) => {
batchContext?.edgeQueue.push(payload as EdgeType[]);
},
[batchContext]
);
// 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`).
useIsomorphicLayoutEffect(() => {
// 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 addNodes = useCallback<Instance.AddNodes<NodeType>>(
(payload) => {
const newNodes = Array.isArray(payload) ? payload : [payload];
if (setElementsQueue.current.nodes.length) {
const { nodes = [], setNodes, hasDefaultNodes, onNodesChange, nodeLookup } = store.getState();
// Queueing a functional update means that we won't worry about other calls
// to `setNodes` that might happen elsewhere.
batchContext?.nodeQueue.push((nodes) => [...nodes, ...newNodes]);
},
[batchContext]
);
// 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;
}
const addEdges = useCallback<Instance.AddEdges<EdgeType>>(
(payload) => {
const newEdges = Array.isArray(payload) ? payload : [payload];
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 newNodes = Array.isArray(payload) ? payload : [payload];
// 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 newEdges = Array.isArray(payload) ? payload : [payload];
setElementsQueue.current.edges.push((edges) => [...edges, ...newEdges]);
setShouldFlushQueue(true);
}, []);
batchContext?.edgeQueue.push((edges) => [...edges, ...newEdges]);
},
[batchContext]
);
const toObject = useCallback<Instance.ToObject<NodeType, EdgeType>>(() => {
const { nodes = [], edges = [], transform } = store.getState();

View File

@@ -236,7 +236,8 @@ export function getElementsDiffChanges({
const itemsLookup = new Map<string, any>(items.map((item) => [item.id, item]));
for (const item of items) {
const storeItem = lookup.get(item.id);
const lookupItem = lookup.get(item.id);
const storeItem = lookupItem?.internals?.userNode ?? lookupItem;
if (storeItem !== undefined && storeItem !== item) {
changes.push({ id: item.id, item: item, type: 'replace' });