refactor(zustand): use shallow when needed

src/container/EdgeRenderer/index.tsx

@@ -1,15 +1,14 @@
 import React, { memo, CSSProperties, useCallback } from 'react';
+import shallow from 'zustand/shallow';
 
 import { useStore } from '../../store';
 import ConnectionLine from '../../components/ConnectionLine/index';
-import { isEdge } from '../../utils/graph';
 import MarkerDefinitions from './MarkerDefinitions';
 import { getEdgePositions, getHandle } from './utils';
 import {
   Position,
   Edge,
   Node,
-  Elements,
   Connection,
   ConnectionLineType,
   ConnectionLineComponent,
@@ -55,7 +54,6 @@ interface EdgeWrapperProps {
   onEdgeUpdate?: OnEdgeUpdateFunc;
   targetNode?: Node;
   sourceNode?: Node;
-  selectedElements: Elements | null;
   elementsSelectable: boolean;
   connectionMode?: ConnectionMode;
 }
@@ -77,7 +75,6 @@ const Edge = memo(
     onEdgeUpdate,
     targetNode,
     sourceNode,
-    selectedElements,
     elementsSelectable,
     connectionMode,
   }: EdgeWrapperProps) => {
@@ -152,8 +149,6 @@ const Edge = memo(
     //   return null;
     // }
 
-    const isSelected = selectedElements?.some((elm) => isEdge(elm) && elm.id === edge.id) || false;
-
     return (
       <EdgeComponent
         key={edge.id}
@@ -162,7 +157,7 @@ const Edge = memo(
         type={edge.type}
         data={edge.data}
         onClick={onElementClick}
-        selected={isSelected}
+        selected={!!edge.selected}
         animated={edge.animated}
         label={edge.label}
         labelStyle={edge.labelStyle}
@@ -207,7 +202,6 @@ const selector = (s: ReactFlowState) => ({
   connectionHandleId: s.connectionHandleId,
   connectionHandleType: s.connectionHandleType,
   connectionPosition: s.connectionPosition,
-  selectedElements: s.selectedElements,
   nodesConnectable: s.nodesConnectable,
   elementsSelectable: s.elementsSelectable,
   width: s.width,
@@ -222,12 +216,11 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
     connectionHandleId,
     connectionHandleType,
     connectionPosition,
-    selectedElements,
     nodesConnectable,
     elementsSelectable,
     width,
     height,
-  } = useStore(selector);
+  } = useStore(selector, shallow);
 
   if (!width) {
     return null;
@@ -247,7 +240,6 @@ const EdgeRenderer = (props: EdgeRendererProps) => {
             edge={edge}
             sourceNode={edge.sourceNode}
             targetNode={edge.targetNode}
-            selectedElements={selectedElements}
             elementsSelectable={elementsSelectable}
             markerEndId={props.markerEndId}
             onEdgeContextMenu={props.onEdgeContextMenu}

src/container/EdgeRenderer/utils.ts

@@ -169,3 +169,16 @@ export const getSourceTargetNodes = (edge: Edge, nodes: Node[]): SourceTargetNod
     { sourceNode: null, targetNode: null } as SourceTargetNode
   );
 };
+
+export const extendEdgeWithSourceAndTarget = (edge: Edge, nodes: Node[]): Edge => {
+  const { sourceNode, targetNode } = getSourceTargetNodes(edge, nodes);
+
+  if (sourceNode) {
+    edge.sourceNode = sourceNode;
+  }
+  if (targetNode) {
+    edge.targetNode = targetNode;
+  }
+
+  return edge;
+};

src/container/FlowRenderer/index.tsx

@@ -1,4 +1,5 @@
 import React, { useCallback, memo, ReactNode, WheelEvent, MouseEvent } from 'react';
+import shallow from 'zustand/shallow';
 
 import { useStore } from '../../store';
 import useGlobalKeyHandler from '../../hooks/useGlobalKeyHandler';
@@ -60,7 +61,7 @@ const FlowRenderer = ({
   onSelectionDragStop,
   onSelectionContextMenu,
 }: FlowRendererProps) => {
-  const { unsetNodesSelection, resetSelectedElements, nodesSelectionActive } = useStore(selector);
+  const { unsetNodesSelection, resetSelectedElements, nodesSelectionActive } = useStore(selector, shallow);
 
   const selectionKeyPressed = useKeyPress(selectionKeyCode);
 

src/container/GraphView/index.tsx

@@ -1,4 +1,5 @@
 import React, { useEffect, useRef, memo } from 'react';
+import shallow from 'zustand/shallow';
 
 import { useStore, useStoreApi } from '../../store';
 import FlowRenderer from '../FlowRenderer';
@@ -136,7 +137,7 @@ const GraphView = ({
     setConnectionMode,
     setOnNodesChange,
     setOnEdgesChange,
-  } = useStore(selector);
+  } = useStore(selector, shallow);
 
   const { zoomIn, zoomOut, zoomTo, transform, fitView, initialized } = useZoomPanHelper();
 

src/container/NodeRenderer/index.tsx

@@ -1,4 +1,5 @@
 import React, { memo, useMemo, ComponentType, MouseEvent } from 'react';
+import shallow from 'zustand/shallow';
 
 import { useStore } from '../../store';
 import { Node, NodeTypesType, ReactFlowState, Edge, WrapNodeProps } from '../../types';
@@ -21,7 +22,6 @@ interface NodeRendererProps {
 
 const selector = (s: ReactFlowState) => ({
   transform: s.transform,
-  selectedElements: s.selectedElements,
   nodesDraggable: s.nodesDraggable,
   nodesConnectable: s.nodesConnectable,
   elementsSelectable: s.elementsSelectable,
@@ -30,15 +30,10 @@ const selector = (s: ReactFlowState) => ({
 });
 
 const NodeRenderer = (props: NodeRendererProps) => {
-  const {
-    transform,
-    selectedElements,
-    nodesDraggable,
-    nodesConnectable,
-    elementsSelectable,
-    nodes,
-    updateNodeDimensions,
-  } = useStore(selector);
+  const { transform, nodesDraggable, nodesConnectable, elementsSelectable, nodes, updateNodeDimensions } = useStore(
+    selector,
+    shallow
+  );
 
   // const visibleNodes = props.onlyRenderVisibleElements
   //   ? getNodesInside(nodes, { x: 0, y: 0, width, height }, transform, true)
@@ -108,7 +103,7 @@ const NodeRenderer = (props: NodeRendererProps) => {
             onNodeDrag={props.onNodeDrag}
             onNodeDragStop={props.onNodeDragStop}
             scale={transform[2]}
-            selected={selectedElements?.some(({ id }) => id === node.id) || false}
+            selected={!!node.selected}
             isDraggable={isDraggable}
             isSelectable={isSelectable}
             isConnectable={isConnectable}

src/container/ZoomPane/index.tsx

@@ -1,6 +1,7 @@
 import React, { useEffect, useRef, ReactNode } from 'react';
 import { zoom, zoomIdentity } from 'd3-zoom';
 import { select, pointer } from 'd3-selection';
+import shallow from 'zustand/shallow';
 
 import { clamp } from '../../utils';
 import useKeyPress from '../../hooks/useKeyPress';
@@ -74,7 +75,7 @@ const ZoomPane = ({
   const store = useStoreApi();
   const zoomPane = useRef<HTMLDivElement>(null);
   const prevTransform = useRef<FlowTransform>({ x: 0, y: 0, zoom: 0 });
-  const { d3Zoom, d3Selection, d3ZoomHandler, initD3Zoom, updateTransform } = useStore(selector);
+  const { d3Zoom, d3Selection, d3ZoomHandler, initD3Zoom, updateTransform } = useStore(selector, shallow);
   const zoomActivationKeyPressed = useKeyPress(zoomActivationKeyCode);
 
   useResizeHandler(zoomPane);