refactor(nodes): simplify selection handling

example/src/Basic/index.tsx

@@ -66,6 +66,7 @@ const BasicFlow = () => {
       maxZoom={4}
       fitView
       defaultEdgeOptions={defaultEdgeOptions}
+      selectNodesOnDrag={false}
     >
       <Background variant={BackgroundVariant.Lines} />
 

src/components/Nodes/DefaultNode.tsx

@@ -4,14 +4,11 @@ import Handle from '../../components/Handle';
 import { NodeProps, Position } from '../../types';
 
 const DefaultNode = ({
-  id,
   data,
   isConnectable,
   targetPosition = Position.Top,
   sourcePosition = Position.Bottom,
 }: NodeProps) => {
-  console.log('render', id);
-
   return (
     <>
       <Handle type="target" position={targetPosition} isConnectable={isConnectable} />

src/components/Nodes/utils.ts

@@ -1,5 +1,5 @@
 import { MouseEvent } from 'react';
-import { GetState } from 'zustand';
+import { GetState, SetState } from 'zustand';
 
 import { HandleElement, Node, Position, ReactFlowState } from '../../types';
 import { getDimensions } from '../../utils';
@@ -55,3 +55,29 @@ export function getMouseHandler(
         handler(event, { ...node });
       };
 }
+
+// this handler is called by
+// 1. the click handler when node is not draggable or selectNodesOnDrag = false
+// or
+// 2. the on drag start handler when node is draggable and selectNodesOnDrag = true
+export function handleNodeClick({
+  id,
+  store,
+}: {
+  id: string;
+  store: {
+    getState: GetState<ReactFlowState>;
+    setState: SetState<ReactFlowState>;
+  };
+}) {
+  const { addSelectedNodes, unselectNodesAndEdges, multiSelectionActive, nodeInternals } = store.getState();
+  const node = nodeInternals.get(id)!;
+
+  store.setState({ nodesSelectionActive: false });
+
+  if (!node.selected) {
+    addSelectedNodes([id]);
+  } else if (node.selected && multiSelectionActive) {
+    unselectNodesAndEdges({ nodes: [node] });
+  }
+}

src/components/Nodes/wrapNode.tsx

@@ -1,17 +1,13 @@
 import React, { useEffect, useRef, memo, ComponentType, MouseEvent } from 'react';
 import cc from 'classcat';
-import shallow from 'zustand/shallow';
 
 import { useStore, useStoreApi } from '../../store';
 import { Provider } from '../../contexts/NodeIdContext';
 import { NodeProps, WrapNodeProps, ReactFlowState } from '../../types';
 import useDrag from '../../hooks/useDrag';
-import { getMouseHandler } from './utils';
+import { getMouseHandler, handleNodeClick } from './utils';
 
-const selector = (s: ReactFlowState) => ({
+const selector = (s: ReactFlowState) => s.updateNodeDimensions;
-  addSelectedNodes: s.addSelectedNodes,
-  updateNodeDimensions: s.updateNodeDimensions,
-});
 
 export default (NodeComponent: ComponentType<NodeProps>) => {
   const NodeWrapper = ({
@@ -27,9 +23,9 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
     onMouseLeave,
     onContextMenu,
     onNodeDoubleClick,
-    onNodeDragStart,
+    onDragStart,
-    onNodeDrag,
+    onDrag,
-    onNodeDragStop,
+    onDragStop,
     style,
     className,
     isDraggable,
@@ -47,8 +43,8 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
     noDragClassName,
   }: WrapNodeProps) => {
     const store = useStoreApi();
-    const { addSelectedNodes, updateNodeDimensions } = useStore(selector, shallow);
+    const updateNodeDimensions = useStore(selector);
-    const nodeElement = useRef<HTMLDivElement>(null);
+    const nodeRef = useRef<HTMLDivElement>(null);
     const prevSourcePosition = useRef(sourcePosition);
     const prevTargetPosition = useRef(targetPosition);
     const prevType = useRef(type);
@@ -60,12 +56,12 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
     const onContextMenuHandler = getMouseHandler(id, store.getState, onContextMenu);
     const onNodeDoubleClickHandler = getMouseHandler(id, store.getState, onNodeDoubleClick);
     const onSelectNodeHandler = (event: MouseEvent) => {
-      if (isSelectable) {
+      if (isSelectable && (!selectNodesOnDrag || !isDraggable)) {
-        store.setState({ nodesSelectionActive: false });
+        // this handler gets called within the drag start event when selectNodesOnDrag=true
-
+        handleNodeClick({
-        if (!selected) {
+          id,
-          addSelectedNodes([id]);
+          store,
-        }
+        });
       }
 
       if (onClick) {
@@ -75,8 +71,8 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
     };
 
     useEffect(() => {
-      if (nodeElement.current && !hidden) {
+      if (nodeRef.current && !hidden) {
-        const currNode = nodeElement.current;
+        const currNode = nodeRef.current;
         resizeObserver?.observe(currNode);
 
         return () => resizeObserver?.unobserve(currNode);
@@ -89,7 +85,7 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
       const sourcePosChanged = prevSourcePosition.current !== sourcePosition;
       const targetPosChanged = prevTargetPosition.current !== targetPosition;
 
-      if (nodeElement.current && (typeChanged || sourcePosChanged || targetPosChanged)) {
+      if (nodeRef.current && (typeChanged || sourcePosChanged || targetPosChanged)) {
         if (typeChanged) {
           prevType.current = type;
         }
@@ -99,15 +95,15 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
         if (targetPosChanged) {
           prevTargetPosition.current = targetPosition;
         }
-        updateNodeDimensions([{ id, nodeElement: nodeElement.current, forceUpdate: true }]);
+        updateNodeDimensions([{ id, nodeElement: nodeRef.current, forceUpdate: true }]);
       }
     }, [id, type, sourcePosition, targetPosition]);
 
     const dragging = useDrag({
-      onStart: onNodeDragStart,
+      onStart: onDragStart,
-      onDrag: onNodeDrag,
+      onDrag: onDrag,
-      onStop: onNodeDragStop,
+      onStop: onDragStop,
-      nodeRef: nodeElement,
+      nodeRef,
       disabled: !isDraggable,
       noDragClassName,
       handleSelector: dragHandle,
@@ -120,22 +116,20 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
       return null;
     }
 
-    const nodeClasses = cc([
-      'react-flow__node',
-      `react-flow__node-${type}`,
-      noPanClassName,
-      className,
-      {
-        selected,
-        selectable: isSelectable,
-        parent: isParent,
-      },
-    ]);
-
     return (
       <div
-        className={nodeClasses}
+        className={cc([
-        ref={nodeElement}
+          'react-flow__node',
+          `react-flow__node-${type}`,
+          noPanClassName,
+          className,
+          {
+            selected,
+            selectable: isSelectable,
+            parent: isParent,
+          },
+        ])}
+        ref={nodeRef}
         style={{
           zIndex,
           transform: `translate(${xPos}px,${yPos}px)`,

src/hooks/useDrag/index.ts

@@ -12,6 +12,7 @@ import {
   selectorExistsTargetToNode,
   updatePosition,
 } from './utils';
+import { handleNodeClick } from '../../components/Nodes/utils';
 
 export type UseDragEvent = D3DragEvent<HTMLDivElement, null, SubjectPosition>;
 export type UseDragData = { dx: number; dy: number };
@@ -76,30 +77,23 @@ function useDrag({
       } else {
         const dragHandler = drag()
           .on('start', (event: UseDragEvent) => {
-            const { nodeInternals, addSelectedNodes, unselectNodesAndEdges, multiSelectionActive } = store.getState();
+            const { nodeInternals, multiSelectionActive, unselectNodesAndEdges } = store.getState();
             parentPos.current = getParentNodePosition(nodeInternals, nodeId);
 
-            // this part is the regular drag handler for a single node
+            if (!selectNodesOnDrag && !multiSelectionActive && nodeId) {
-            // it selects the dragged node and deselects all other nodes if multiSelectionActive = false
+              if (!nodeInternals.get(nodeId)?.selected) {
-            if (nodeId && isSelectable) {
+                // we need to reset selected nodes when selectNodesOnDrag=false
-              const node = nodeInternals.get(nodeId)!;
+                unselectNodesAndEdges();
-
-              if (selectNodesOnDrag) {
-                store.setState({ nodesSelectionActive: false });
-
-                if (!node.selected) {
-                  addSelectedNodes([nodeId]);
-                }
-              } else if (!selectNodesOnDrag && !node.selected) {
-                if (multiSelectionActive) {
-                  addSelectedNodes([nodeId]);
-                } else {
-                  unselectNodesAndEdges();
-                  store.setState({ nodesSelectionActive: false });
-                }
               }
             }
 
+            if (nodeId && isSelectable && selectNodesOnDrag) {
+              handleNodeClick({
+                id: nodeId,
+                store,
+              });
+            }
+
             const mousePos = getMousePosition(event);
             dragItems.current = getDragItems(nodeInternals, mousePos, nodeId);
 

src/store/index.ts

@@ -14,6 +14,7 @@ import {
   NodeSelectionChange,
   NodePositionChange,
   NodeDragItem,
+  UnselectNodesAndEdgesParams,
 } from '../types';
 import { getHandleBounds } from '../components/Nodes/utils';
 import { createSelectionChange, getSelectionChanges } from '../utils/changes';
@@ -152,19 +153,22 @@ const createStore = () =>
         set,
       });
     },
-    unselectNodesAndEdges: () => {
+    unselectNodesAndEdges: ({ nodes, edges }: UnselectNodesAndEdgesParams = {}) => {
-      const { nodeInternals, edges } = get();
+      const { nodeInternals, edges: storeEdges } = get();
-      const nodes = Array.from(nodeInternals.values());
+      const nodesToUnselect = nodes ? nodes : Array.from(nodeInternals.values());
+      const edgesToUnselect = edges ? edges : storeEdges;
 
-      const nodesToUnselect = nodes.map((n) => {
+      const changedNodes = nodesToUnselect.map((n) => {
         n.selected = false;
         return createSelectionChange(n.id, false);
       }) as NodeSelectionChange[];
-      const edgesToUnselect = edges.map((edge) => createSelectionChange(edge.id, false)) as EdgeSelectionChange[];
+      const changedEdges = edgesToUnselect.map((edge) =>
+        createSelectionChange(edge.id, false)
+      ) as EdgeSelectionChange[];
 
       updateNodesAndEdgesSelections({
-        changedNodes: nodesToUnselect,
+        changedNodes,
-        changedEdges: edgesToUnselect,
+        changedEdges,
         get,
         set,
       });

src/types/general.ts

@@ -105,6 +105,11 @@ export type FitBoundsOptions = ViewportHelperFunctionOptions & {
   padding?: number;
 };
 
+export type UnselectNodesAndEdgesParams = {
+  nodes?: Node[];
+  edges?: Edge[];
+};
+
 export interface ViewportHelperFunctions {
   zoomIn: ZoomInOut;
   zoomOut: ZoomInOut;
@@ -187,7 +192,7 @@ export type ReactFlowActions = {
   updateNodeDimensions: (updates: NodeDimensionUpdate[]) => void;
   updateNodePositions: (nodeDragItems: NodeDragItem[]) => void;
   resetSelectedElements: () => void;
-  unselectNodesAndEdges: () => void;
+  unselectNodesAndEdges: (params?: UnselectNodesAndEdgesParams) => void;
   addSelectedNodes: (nodeIds: string[]) => void;
   addSelectedEdges: (edgeIds: string[]) => void;
   setMinZoom: (minZoom: number) => void;

src/types/nodes.ts

@@ -71,9 +71,9 @@ export interface WrapNodeProps<T = any> {
   onMouseMove?: NodeMouseHandler;
   onMouseLeave?: NodeMouseHandler;
   onContextMenu?: NodeMouseHandler;
-  onNodeDragStart?: NodeDragHandler;
+  onDragStart?: NodeDragHandler;
-  onNodeDrag?: NodeDragHandler;
+  onDrag?: NodeDragHandler;
-  onNodeDragStop?: NodeDragHandler;
+  onDragStop?: NodeDragHandler;
   style?: CSSProperties;
   className?: string;
   sourcePosition: Position;

src/utils/changes.ts

@@ -45,8 +45,7 @@ function handleParentExpand(res: any[], updateItem: any) {
 }
 
 function applyChanges(changes: any[], elements: any[]): any[] {
-  // unfortunately we need this hack to handle the setNodes and setEdges function of the
+  // we need this hack to handle the setNodes and setEdges function of the useReactFlow hook for controlled flows
-  // useReactFlow hook.
   if (changes.some((c) => c.type === 'reset')) {
     return changes.filter((c) => c.type === 'reset').map((c) => c.item);
   }