refactor(zustand): use shallow when needed

src/components/Edges/wrapEdge.tsx

@@ -1,5 +1,6 @@
 import React, { memo, ComponentType, useCallback, useState, useMemo } from 'react';
 import cc from 'classcat';
+import shallow from 'zustand/shallow';
 
 import { useStore, useStoreApi } from '../../store';
 import { Edge, EdgeProps, WrapEdgeProps, ReactFlowState } from '../../types';
@@ -56,8 +57,10 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
     onEdgeUpdateEnd,
   }: WrapEdgeProps): JSX.Element | null => {
     const store = useStoreApi();
-    const { addSelectedElements, setConnectionNodeId, unsetNodesSelection, setPosition, connectionMode } =
-      useStore(selector);
+    const { addSelectedElements, setConnectionNodeId, unsetNodesSelection, setPosition, connectionMode } = useStore(
+      selector,
+      shallow
+    );
 
     const [updating, setUpdating] = useState<boolean>(false);
 

src/components/ElementUpdater/index.tsx

@@ -1,4 +1,5 @@
 import { useEffect } from 'react';
+import shallow from 'zustand/shallow';
 
 import { useStore } from '../../store';
 import { Node, Edge, ReactFlowState } from '../../types';
@@ -14,7 +15,7 @@ const selector = (s: ReactFlowState) => ({
 });
 
 const ElementUpdater = ({ nodes, edges }: ElementUpdaterProps) => {
-  const { setNodes, setEdges } = useStore(selector);
+  const { setNodes, setEdges } = useStore(selector, shallow);
 
   useEffect(() => {
     setNodes(nodes);

src/components/Handle/index.tsx

@@ -1,5 +1,6 @@
 import React, { memo, useContext, useCallback, HTMLAttributes, forwardRef } from 'react';
 import cc from 'classcat';
+import shallow from 'zustand/shallow';
 
 import { useStore } from '../../store';
 import NodeIdContext from '../../contexts/NodeIdContext';
@@ -45,7 +46,7 @@ const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
       onConnectStop,
       onConnectEnd,
       connectionMode,
-    } = useStore(selector);
+    } = useStore(selector, shallow);
 
     const handleId = id || null;
     const isTarget = type === 'target';

src/components/Nodes/wrapNode.tsx

@@ -1,6 +1,7 @@
 import React, { useEffect, useRef, memo, ComponentType, CSSProperties, useMemo, MouseEvent, useCallback } from 'react';
 import { DraggableCore, DraggableData, DraggableEvent } from 'react-draggable';
 import cc from 'classcat';
+import shallow from 'zustand/shallow';
 
 import { useStore } from '../../store';
 import { Provider } from '../../contexts/NodeIdContext';
@@ -10,6 +11,7 @@ const selector = (s: ReactFlowState) => ({
   addSelectedElements: s.addSelectedElements,
   onNodesChange: s.onNodesChange,
   unsetNodesSelection: s.unsetNodesSelection,
+  updateNodePosDiff: s.updateNodePosDiff,
 });
 
 export default (NodeComponent: ComponentType<NodeComponentProps>) => {
@@ -47,7 +49,7 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
     dragHandle,
   }: WrapNodeProps) => {
     // const updateNodeDimensions = useStoreActions((actions) => actions.updateNodeDimensions);
-    const { addSelectedElements, onNodesChange, unsetNodesSelection } = useStore(selector);
+    const { addSelectedElements, onNodesChange, unsetNodesSelection, updateNodePosDiff } = useStore(selector, shallow);
     const nodeElement = useRef<HTMLDivElement>(null);
 
     const node = useMemo(() => ({ id, type, position: { x: xPos, y: yPos }, data }), [id, type, xPos, yPos, data]);
@@ -154,18 +156,7 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
           onNodeDrag(event as MouseEvent, node);
         }
 
-        onNodesChange?.([
-          {
-            id,
-            change: {
-              position: {
-                x: node.position.x,
-                y: node.position.y,
-              },
-              isDragging: true,
-            },
-          },
-        ]);
+        updateNodePosDiff({ id, isDragging: true, diff: { x: draggableData.deltaX, y: draggableData.deltaY } });
       },
       [id, node, onNodeDrag]
     );

src/components/NodesSelection/index.tsx

@@ -7,7 +7,6 @@ import React, { useMemo, useCallback, useRef, MouseEvent } from 'react';
 import ReactDraggable, { DraggableData } from 'react-draggable';
 
 import { useStore } from '../../store';
-import { isNode } from '../../utils/graph';
 import { Node, ReactFlowState } from '../../types';
 
 export interface NodesSelectionProps {
@@ -21,10 +20,9 @@ const selector = (s: ReactFlowState) => ({
   transform: s.transform,
   selectedNodesBbox: s.selectedNodesBbox,
   selectionActive: s.selectionActive,
-  selectedElements: s.selectedElements,
+  selectedNodes: s.nodes.filter((n) => n.selected),
   snapToGrid: s.snapToGrid,
   snapGrid: s.snapGrid,
-  nodes: s.nodes,
   updateNodePosDiff: s.updateNodePosDiff,
 });
 
@@ -34,37 +32,13 @@ export default ({
   onSelectionDragStop,
   onSelectionContextMenu,
 }: NodesSelectionProps) => {
-  const {
-    transform,
-    selectedNodesBbox,
-    selectionActive,
-    selectedElements,
-    snapToGrid,
-    snapGrid,
-    nodes,
-    updateNodePosDiff,
-  } = useStore(selector);
+  const { transform, selectedNodesBbox, selectionActive, selectedNodes, snapToGrid, snapGrid, updateNodePosDiff } =
+    useStore(selector);
   const [tX, tY, tScale] = transform;
-
   const nodeRef = useRef(null);
 
   const grid = useMemo(() => (snapToGrid ? snapGrid : [1, 1])! as [number, number], [snapToGrid, snapGrid]);
 
-  const selectedNodes = useMemo(
-    () =>
-      selectedElements
-        ? selectedElements.filter(isNode).map((selectedNode) => {
-            const matchingNode = nodes.find((node) => node.id === selectedNode.id);
-
-            return {
-              ...matchingNode,
-              position: matchingNode?.position,
-            } as Node;
-          })
-        : [],
-    [selectedElements, nodes]
-  );
-
   const style = useMemo(
     () => ({
       transform: `translate(${tX}px,${tY}px) scale(${tScale})`,
@@ -119,16 +93,12 @@ export default ({
 
   const onContextMenu = useCallback(
     (event: MouseEvent) => {
-      const selectedNodes = selectedElements
-        ? selectedElements.filter(isNode).map((selectedNode) => nodes.find((node) => node.id === selectedNode.id)!)
-        : [];
-
       onSelectionContextMenu?.(event, selectedNodes);
     },
-    [onSelectionContextMenu]
+    [onSelectionContextMenu, selectedNodes]
   );
 
-  if (!selectedElements || selectionActive) {
+  if (!selectedNodes || selectionActive) {
     return null;
   }
 

src/components/SelectionListener/index.tsx

@@ -1,4 +1,5 @@
 import { useEffect } from 'react';
+import shallow from 'zustand/shallow';
 
 import { Elements, ReactFlowState } from '../../types';
 import { useStore } from '../../store';
@@ -7,12 +8,15 @@ interface SelectionListenerProps {
   onSelectionChange: (elements: Elements | null) => void;
 }
 
-const selectedElementsSelector = (s: ReactFlowState) => s.selectedElements;
+const selectedElementsSelector = (s: ReactFlowState) => [
+  ...s.nodes.filter((n) => n.selected),
+  ...s.edges.filter((e) => e.selected),
+];
 
 // This is just a helper component for calling the onSelectionChange listener.
 
 export default ({ onSelectionChange }: SelectionListenerProps) => {
-  const selectedElements = useStore(selectedElementsSelector);
+  const selectedElements = useStore(selectedElementsSelector, shallow);
 
   useEffect(() => {
     onSelectionChange(selectedElements);