feat: improve nodes and edges state update

packages/web/src/components/EditorNew/Edge/Edge.jsx

@@ -1,10 +1,10 @@
 import { EdgeLabelRenderer, getStraightPath } from 'reactflow';
 import IconButton from '@mui/material/IconButton';
 import AddIcon from '@mui/icons-material/Add';
-import { useMutation } from '@apollo/client';
-import { CREATE_STEP } from 'graphql/mutations/create-step';
-import { useQueryClient } from '@tanstack/react-query';
+
 import PropTypes from 'prop-types';
+import { useContext } from 'react';
+import { EdgesContext } from '../EditorNew';
 
 export default function Edge({
   sourceX,
@@ -12,11 +12,11 @@ export default function Edge({
   targetX,
   targetY,
   source,
-  data: { flowId, setCurrentStepId, flowActive, layouted },
+  data: { layouted },
 }) {
-  const [createStep, { loading: creationInProgress }] =
-    useMutation(CREATE_STEP);
-  const queryClient = useQueryClient();
+  const { stepCreationInProgress, flowActive, onAddStep } =
+    useContext(EdgesContext);
+
   const [edgePath, labelX, labelY] = getStraightPath({
     sourceX,
     sourceY,
@@ -24,30 +24,11 @@ export default function Edge({
     targetY,
   });
 
-  const addStep = async (previousStepId) => {
-    const mutationInput = {
-      previousStep: {
-        id: previousStepId,
-      },
-      flow: {
-        id: flowId,
-      },
-    };
-
-    const createdStep = await createStep({
-      variables: { input: mutationInput },
-    });
-
-    const createdStepId = createdStep.data.createStep.id;
-    setCurrentStepId(createdStepId);
-    await queryClient.invalidateQueries({ queryKey: ['flows', flowId] });
-  };
-
   return (
     <>
       <EdgeLabelRenderer>
         <IconButton
-          onClick={() => addStep(source)}
+          onClick={() => onAddStep(source)}
           color="primary"
           sx={{
             position: 'absolute',
@@ -55,7 +36,7 @@ export default function Edge({
             pointerEvents: 'all',
             visibility: layouted ? 'visible' : 'hidden',
           }}
-          disabled={creationInProgress || flowActive}
+          disabled={stepCreationInProgress || flowActive}
         >
           <AddIcon />
         </IconButton>
@@ -71,9 +52,6 @@ Edge.propTypes = {
   targetY: PropTypes.number.isRequired,
   source: PropTypes.string.isRequired,
   data: PropTypes.shape({
-    flowId: PropTypes.string.isRequired,
-    setCurrentStepId: PropTypes.func.isRequired,
-    flowActive: PropTypes.bool.isRequired,
     layouted: PropTypes.bool,
   }).isRequired,
 };

packages/web/src/components/EditorNew/EditorNew.jsx

@@ -1,10 +1,11 @@
-import { useEffect, useState, useCallback } from 'react';
+import { useEffect, useCallback, createContext, useRef } from 'react';
 import { useMutation } from '@apollo/client';
 import { useQueryClient } from '@tanstack/react-query';
 import { FlowPropType } from 'propTypes/propTypes';
-import ReactFlow, { useNodesState, useEdgesState, addEdge } from 'reactflow';
+import ReactFlow, { useNodesState, useEdgesState } from 'reactflow';
 import 'reactflow/dist/style.css';
 import { UPDATE_STEP } from 'graphql/mutations/update-step';
+import { CREATE_STEP } from 'graphql/mutations/create-step';
 
 import { useAutoLayout } from './useAutoLayout';
 import { useScrollBoundries } from './useScrollBoundries';
@@ -12,39 +13,69 @@ import FlowStepNode from './FlowStepNode/FlowStepNode';
 import Edge from './Edge/Edge';
 import InvisibleNode from './InvisibleNode/InvisibleNode';
 import { EditorWrapper } from './style';
+import {
+  generateEdgeId,
+  generateInitialEdges,
+  generateInitialNodes,
+  updatedCollapsedNodes,
+} from './utils';
+import { EDGE_TYPES, INVISIBLE_NODE_ID, NODE_TYPES } from './constants';
 
-const nodeTypes = { flowStep: FlowStepNode, invisible: InvisibleNode };
+export const EdgesContext = createContext();
+export const NodesContext = createContext();
 
-const edgeTypes = {
-  addNodeEdge: Edge,
+const nodeTypes = {
+  [NODE_TYPES.FLOW_STEP]: FlowStepNode,
+  [NODE_TYPES.INVISIBLE]: InvisibleNode,
 };
 
-const INVISIBLE_NODE_ID = 'invisible-node';
-
-const generateEdgeId = (sourceId, targetId) => `${sourceId}-${targetId}`;
+const edgeTypes = {
+  [EDGE_TYPES.ADD_NODE_EDGE]: Edge,
+};
 
 const EditorNew = ({ flow }) => {
-  const [triggerStep] = flow.steps;
-  const [currentStepId, setCurrentStepId] = useState(triggerStep.id);
-
   const [updateStep] = useMutation(UPDATE_STEP);
   const queryClient = useQueryClient();
+  const [createStep, { loading: stepCreationInProgress }] =
+    useMutation(CREATE_STEP);
+
+  const [nodes, setNodes, onNodesChange] = useNodesState(
+    generateInitialNodes(flow),
+  );
+  const [edges, setEdges, onEdgesChange] = useEdgesState(
+    generateInitialEdges(flow),
+  );
 
-  const [nodes, setNodes, onNodesChange] = useNodesState([]);
-  const [edges, setEdges, onEdgesChange] = useEdgesState([]);
   useAutoLayout();
   useScrollBoundries();
 
-  const onConnect = useCallback(
-    (params) => setEdges((eds) => addEdge(params, eds)),
-    [setEdges],
-  );
+  const createdStepIdRef = useRef(null);
 
   const openNextStep = useCallback(
-    (nextStep) => () => {
-      setCurrentStepId(nextStep?.id);
+    (currentStepId) => {
+      setNodes((nodes) => {
+        const currentStepIndex = nodes.findIndex(
+          (node) => node.id === currentStepId,
+        );
+        if (currentStepIndex >= 0) {
+          const nextStep = nodes[currentStepIndex + 1];
+          return updatedCollapsedNodes(nodes, nextStep.id);
+        }
+        return nodes;
+      });
     },
-    [],
+    [setNodes],
+  );
+
+  const onStepClose = useCallback(() => {
+    setNodes((nodes) => updatedCollapsedNodes(nodes));
+  }, [setNodes]);
+
+  const onStepOpen = useCallback(
+    (stepId) => {
+      setNodes((nodes) => updatedCollapsedNodes(nodes, stepId));
+    },
+    [setNodes],
   );
 
   const onStepChange = useCallback(
@@ -76,178 +107,166 @@ const EditorNew = ({ flow }) => {
     [flow.id, updateStep, queryClient],
   );
 
-  const generateEdges = useCallback((flow, prevEdges) => {
-    const newEdges =
-      flow.steps
-        .map((step, i) => {
-          const sourceId = step.id;
-          const targetId = flow.steps[i + 1]?.id;
-          const edge = prevEdges?.find(
-            (edge) => edge.id === generateEdgeId(sourceId, targetId),
-          );
-          if (targetId) {
-            return {
-              id: generateEdgeId(sourceId, targetId),
-              source: sourceId,
-              target: targetId,
-              type: 'addNodeEdge',
-              data: {
-                flowId: flow.id,
-                flowActive: flow.active,
-                setCurrentStepId,
-                layouted: !!edge,
-              },
-            };
-          }
-        })
-        .filter((edge) => !!edge) || [];
+  const onAddStep = useCallback(
+    async (previousStepId) => {
+      const mutationInput = {
+        previousStep: {
+          id: previousStepId,
+        },
+        flow: {
+          id: flow.id,
+        },
+      };
 
-    const lastStep = flow.steps[flow.steps.length - 1];
-
-    return lastStep
-      ? [
-          ...newEdges,
-          {
-            id: generateEdgeId(lastStep.id, INVISIBLE_NODE_ID),
-            source: lastStep.id,
-            target: INVISIBLE_NODE_ID,
-            type: 'addNodeEdge',
-            data: {
-              flowId: flow.id,
-              flowActive: flow.active,
-              setCurrentStepId,
-              layouted: false,
-            },
-          },
-        ]
-      : newEdges;
-  }, []);
-
-  const generateNodes = useCallback(
-    (flow, prevNodes) => {
-      const newNodes = flow.steps.map((step, index) => {
-        const node = prevNodes?.find(({ id }) => id === step.id);
-        const collapsed = currentStepId !== step.id;
-        return {
-          id: step.id,
-          type: 'flowStep',
-          position: {
-            x: node ? node.position.x : 0,
-            y: node ? node.position.y : 0,
-          },
-          zIndex: collapsed ? 0 : 1,
-          data: {
-            step,
-            index: index,
-            flowId: flow.id,
-            collapsed,
-            openNextStep: openNextStep(flow.steps[index + 1]),
-            onOpen: () => setCurrentStepId(step.id),
-            onClose: () => setCurrentStepId(null),
-            onChange: onStepChange,
-            layouted: !!node,
-          },
-        };
+      const {
+        data: { createStep: createdStep },
+      } = await createStep({
+        variables: { input: mutationInput },
       });
 
-      const prevInvisibleNode = nodes.find((node) => node.type === 'invisible');
-
-      return [
-        ...newNodes,
-        {
-          id: INVISIBLE_NODE_ID,
-          type: 'invisible',
-          position: {
-            x: prevInvisibleNode ? prevInvisibleNode.position.x : 0,
-            y: prevInvisibleNode ? prevInvisibleNode.position.y : 0,
-          },
-        },
-      ];
+      const createdStepId = createdStep.id;
+      await queryClient.invalidateQueries({ queryKey: ['flows', flow.id] });
+      createdStepIdRef.current = createdStepId;
     },
-    [currentStepId, nodes, onStepChange, openNextStep],
+    [flow.id, createStep, queryClient],
   );
 
-  const updateNodesData = useCallback(
-    (steps) => {
-      setNodes((nodes) =>
-        nodes.map((node) => {
-          const step = steps.find((step) => step.id === node.id);
-          if (step) {
-            return { ...node, data: { ...node.data, step: { ...step } } };
-          }
-          return node;
-        }),
-      );
-    },
-    [setNodes],
-  );
-
-  const updateEdgesData = useCallback(
-    (flow) => {
-      setEdges((edges) =>
-        edges.map((edge) => {
-          return {
-            ...edge,
-            data: { ...edge.data, flowId: flow.id, flowActive: flow.active },
-          };
-        }),
-      );
-    },
-    [setEdges],
-  );
-
-  useEffect(() => {
-    setNodes(
-      nodes.map((node) => {
-        if (node.type === 'flowStep') {
-          const collapsed = currentStepId !== node.data.step.id;
-          return {
-            ...node,
-            zIndex: collapsed ? 0 : 1,
-            data: {
-              ...node.data,
-              collapsed,
-            },
-          };
-        }
-        return node;
-      }),
-    );
-  }, [currentStepId]);
-
   useEffect(() => {
     if (flow.steps.length + 1 !== nodes.length) {
-      const newNodes = generateNodes(flow, nodes);
-      const newEdges = generateEdges(flow, edges);
+      setNodes((nodes) => {
+        const newNodes = flow.steps.map((step) => {
+          const createdStepId = createdStepIdRef.current;
+          const prevNode = nodes.find(({ id }) => id === step.id);
+          if (prevNode) {
+            return {
+              ...prevNode,
+              zIndex: createdStepId ? 0 : prevNode.zIndex,
+              data: {
+                ...prevNode.data,
+                collapsed: createdStepId ? true : prevNode.data.collapsed,
+              },
+            };
+          } else {
+            return {
+              id: step.id,
+              type: NODE_TYPES.FLOW_STEP,
+              position: {
+                x: 0,
+                y: 0,
+              },
+              zIndex: 1,
+              data: {
+                collapsed: false,
+                layouted: false,
+              },
+            };
+          }
+        });
 
-      setNodes(newNodes);
-      setEdges(newEdges);
-    } else {
-      updateNodesData(flow.steps);
-      updateEdgesData(flow);
+        const prevInvisible = nodes.find(({ id }) => id === INVISIBLE_NODE_ID);
+        return [
+          ...newNodes,
+          {
+            id: INVISIBLE_NODE_ID,
+            type: NODE_TYPES.INVISIBLE,
+            position: {
+              x: prevInvisible?.position.x || 0,
+              y: prevInvisible?.position.y || 0,
+            },
+          },
+        ];
+      });
+
+      setEdges((edges) => {
+        const newEdges = flow.steps
+          .map((step, i) => {
+            const sourceId = step.id;
+            const targetId = flow.steps[i + 1]?.id;
+            const edge = edges?.find(
+              (edge) => edge.id === generateEdgeId(sourceId, targetId),
+            );
+            if (targetId) {
+              return {
+                id: generateEdgeId(sourceId, targetId),
+                source: sourceId,
+                target: targetId,
+                type: 'addNodeEdge',
+                data: {
+                  layouted: edge ? edge?.data.layouted : false,
+                },
+              };
+            }
+            return null;
+          })
+          .filter((edge) => !!edge);
+
+        const lastStep = flow.steps[flow.steps.length - 1];
+        const lastEdge = edges[edges.length - 1];
+
+        return lastStep
+          ? [
+              ...newEdges,
+              {
+                id: generateEdgeId(lastStep.id, INVISIBLE_NODE_ID),
+                source: lastStep.id,
+                target: INVISIBLE_NODE_ID,
+                type: 'addNodeEdge',
+                data: {
+                  layouted:
+                    lastEdge?.id ===
+                    generateEdgeId(lastStep.id, INVISIBLE_NODE_ID)
+                      ? lastEdge?.data.layouted
+                      : false,
+                },
+              },
+            ]
+          : newEdges;
+      });
+
+      if (createdStepIdRef.current) {
+        createdStepIdRef.current = null;
+      }
     }
-  }, [flow]);
+  }, [flow.steps]);
 
   return (
-    <EditorWrapper direction="column">
-      <ReactFlow
-        nodes={nodes}
-        edges={edges}
-        onNodesChange={onNodesChange}
-        onEdgesChange={onEdgesChange}
-        onConnect={onConnect}
-        nodeTypes={nodeTypes}
-        edgeTypes={edgeTypes}
-        panOnScroll
-        panOnScrollMode="vertical"
-        panOnDrag={false}
-        zoomOnScroll={false}
-        zoomOnPinch={false}
-        zoomOnDoubleClick={false}
-        panActivationKeyCode={null}
-        proOptions={{ hideAttribution: true }}
-      />
-    </EditorWrapper>
+    <NodesContext.Provider
+      value={{
+        openNextStep,
+        onStepOpen,
+        onStepClose,
+        onStepChange,
+        flowId: flow.id,
+        steps: flow.steps,
+      }}
+    >
+      <EdgesContext.Provider
+        value={{
+          stepCreationInProgress,
+          onAddStep,
+          flowActive: flow.active,
+        }}
+      >
+        <EditorWrapper direction="column">
+          <ReactFlow
+            nodes={nodes}
+            edges={edges}
+            onNodesChange={onNodesChange}
+            onEdgesChange={onEdgesChange}
+            nodeTypes={nodeTypes}
+            edgeTypes={edgeTypes}
+            panOnScroll
+            panOnScrollMode="vertical"
+            panOnDrag={false}
+            zoomOnScroll={false}
+            zoomOnPinch={false}
+            zoomOnDoubleClick={false}
+            panActivationKeyCode={null}
+            proOptions={{ hideAttribution: true }}
+          />
+        </EditorWrapper>
+      </EdgesContext.Provider>
+    </NodesContext.Provider>
   );
 };
 

packages/web/src/components/EditorNew/FlowStepNode/FlowStepNode.jsx

@@ -1,25 +1,24 @@
 import { Handle, Position } from 'reactflow';
-import { Box } from '@mui/material';
 import PropTypes from 'prop-types';
 
 import FlowStep from 'components/FlowStep';
-import { StepPropType } from 'propTypes/propTypes';
 
 import { NodeWrapper, NodeInnerWrapper } from './style.js';
+import { useContext } from 'react';
+import { NodesContext } from '../EditorNew.jsx';
 
-function FlowStepNode({
-  data: {
-    step,
-    index,
-    flowId,
-    collapsed,
+function FlowStepNode({ data: { collapsed, layouted }, id }) {
+  const {
     openNextStep,
-    onOpen,
-    onClose,
-    onChange,
-    layouted,
-  },
-}) {
+    onStepOpen,
+    onStepClose,
+    onStepChange,
+    flowId,
+    steps,
+  } = useContext(NodesContext);
+
+  const step = steps.find(({ id: stepId }) => stepId === id);
+
   return (
     <NodeWrapper
       className="nodrag"
@@ -34,16 +33,17 @@ function FlowStepNode({
           isConnectable={false}
           style={{ visibility: 'hidden' }}
         />
-        <FlowStep
-          step={step}
-          index={index + 1}
-          collapsed={collapsed}
-          onOpen={onOpen}
-          onClose={onClose}
-          onChange={onChange}
-          flowId={flowId}
-          onContinue={openNextStep}
-        />
+        {step && (
+          <FlowStep
+            step={step}
+            collapsed={collapsed}
+            onOpen={() => onStepOpen(step.id)}
+            onClose={onStepClose}
+            onChange={onStepChange}
+            flowId={flowId}
+            onContinue={() => openNextStep(step.id)}
+          />
+        )}
         <Handle
           type="source"
           position={Position.Bottom}
@@ -56,15 +56,9 @@ function FlowStepNode({
 }
 
 FlowStepNode.propTypes = {
+  id: PropTypes.string,
   data: PropTypes.shape({
-    step: StepPropType.isRequired,
-    index: PropTypes.number.isRequired,
-    flowId: PropTypes.string.isRequired,
     collapsed: PropTypes.bool.isRequired,
-    openNextStep: PropTypes.func.isRequired,
-    onOpen: PropTypes.func.isRequired,
-    onClose: PropTypes.func.isRequired,
-    onChange: PropTypes.func.isRequired,
     layouted: PropTypes.bool.isRequired,
   }).isRequired,
 };

packages/web/src/components/EditorNew/constants.js

@@ -0,0 +1,10 @@
+export const INVISIBLE_NODE_ID = 'invisible-node';
+
+export const NODE_TYPES = {
+  FLOW_STEP: 'flowStep',
+  INVISIBLE: 'invisible',
+};
+
+export const EDGE_TYPES = {
+  ADD_NODE_EDGE: 'addNodeEdge',
+};

packages/web/src/components/EditorNew/utils.js

@@ -0,0 +1,88 @@
+import { INVISIBLE_NODE_ID, NODE_TYPES } from './constants';
+
+export const generateEdgeId = (sourceId, targetId) => `${sourceId}-${targetId}`;
+
+export const updatedCollapsedNodes = (nodes, openStepId) => {
+  return nodes.map((node) => {
+    if (node.type !== NODE_TYPES.FLOW_STEP) {
+      return node;
+    }
+
+    const collapsed = node.id !== openStepId;
+    return {
+      ...node,
+      zIndex: collapsed ? 0 : 1,
+      data: { ...node.data, collapsed },
+    };
+  });
+};
+
+export const generateInitialNodes = (flow) => {
+  const newNodes = flow.steps.map((step, index) => {
+    const collapsed = index !== 0;
+
+    return {
+      id: step.id,
+      type: NODE_TYPES.FLOW_STEP,
+      position: {
+        x: 0,
+        y: 0,
+      },
+      zIndex: collapsed ? 0 : 1,
+      data: {
+        collapsed,
+        layouted: false,
+      },
+    };
+  });
+
+  return [
+    ...newNodes,
+    {
+      id: INVISIBLE_NODE_ID,
+      type: NODE_TYPES.INVISIBLE,
+      position: {
+        x: 0,
+        y: 0,
+      },
+    },
+  ];
+};
+
+export const generateInitialEdges = (flow) => {
+  const newEdges = flow.steps
+    .map((step, i) => {
+      const sourceId = step.id;
+      const targetId = flow.steps[i + 1]?.id;
+      if (targetId) {
+        return {
+          id: generateEdgeId(sourceId, targetId),
+          source: sourceId,
+          target: targetId,
+          type: 'addNodeEdge',
+          data: {
+            layouted: false,
+          },
+        };
+      }
+      return null;
+    })
+    .filter((edge) => !!edge);
+
+  const lastStep = flow.steps[flow.steps.length - 1];
+
+  return lastStep
+    ? [
+        ...newEdges,
+        {
+          id: generateEdgeId(lastStep.id, INVISIBLE_NODE_ID),
+          source: lastStep.id,
+          target: INVISIBLE_NODE_ID,
+          type: 'addNodeEdge',
+          data: {
+            layouted: false,
+          },
+        },
+      ]
+    : newEdges;
+};

packages/web/src/components/FlowStep/index.jsx

@@ -360,7 +360,6 @@ function FlowStep(props) {
 FlowStep.propTypes = {
   collapsed: PropTypes.bool,
   step: StepPropType.isRequired,
-  index: PropTypes.number,
   onOpen: PropTypes.func,
   onClose: PropTypes.func,
   onChange: PropTypes.func.isRequired,