implement collapse/expand

src/components/CustomNode/TextNode.tsx

@@ -1,7 +1,9 @@
 import React from "react";
 import { BiChevronLeft, BiChevronRight } from "react-icons/bi";
+import { NodeData, EdgeData } from "reaflow";
 import { ConditionalWrapper, CustomNodeProps } from "src/components/CustomNode";
 import useConfig from "src/hooks/store/useConfig";
+import useGraph from "src/hooks/store/useGraph";
 import styled from "styled-components";
 import * as Styled from "./styles";
 
@@ -23,7 +25,8 @@ const StyledExpand = styled.button`
   border-left: 1px solid ${({ theme }) => theme.BACKGROUND_MODIFIER_ACCENT};
 `;
 
-const TextNode: React.FC<CustomNodeProps<string>> = ({
+const TextNode: React.FC<CustomNodeProps<string> & { node: NodeData }> = ({
+  node,
   width,
   height,
   value,
@@ -34,6 +37,7 @@ const TextNode: React.FC<CustomNodeProps<string>> = ({
   const performanceMode = useConfig((state) => state.performanceMode);
   const expand = useConfig((state) => state.expand);
   const [isExpanded, setIsExpanded] = React.useState(!expand);
+  const toggleNode = useGraph((state) => state.toggleNode);
 
   React.useEffect(() => {
     setIsExpanded(expand);
@@ -41,6 +45,7 @@ const TextNode: React.FC<CustomNodeProps<string>> = ({
 
   const handleExpand = (e: React.MouseEvent<HTMLButtonElement>) => {
     e.stopPropagation();
+    toggleNode(node, isExpanded);
     setIsExpanded(!isExpanded);
   };
 

src/components/CustomNode/index.tsx

@@ -37,7 +37,7 @@ export const CustomNode = (nodeProps: NodeProps) => {
 
   return (
     <Node {...nodeProps} label={<Label style={baseLabelStyle} />}>
-      {({ width, height, x, y }) => {
+      {({ width, height, x, y, node }) => {
         if (properties.text instanceof Object) {
           const entries = Object.entries<string>(properties.text);
 
@@ -54,6 +54,7 @@ export const CustomNode = (nodeProps: NodeProps) => {
 
         return (
           <TextNode
+            node={node}
             isParent={properties.data.isParent}
             value={properties.text}
             width={width}

src/components/Graph/index.tsx

@@ -4,7 +4,7 @@ import {
   TransformComponent,
   TransformWrapper,
 } from "react-zoom-pan-pinch";
-import { Canvas, Edge, ElkRoot, useSelection } from "reaflow";
+import { Canvas, Edge, ElkRoot, NodeData, useSelection } from "reaflow";
 import { CustomNode } from "src/components/CustomNode";
 import { NodeModal } from "src/containers/Modals/NodeModal";
 import { getEdgeNodes } from "src/containers/Editor/LiveEditor/helpers";
@@ -12,6 +12,7 @@ import useConfig from "src/hooks/store/useConfig";
 import styled from "styled-components";
 import shallow from "zustand/shallow";
 import useNodeTools from "src/hooks/store/useNodeTools";
+import useGraph from "src/hooks/store/useGraph";
 
 interface LayoutProps {
   isWidget: boolean;
@@ -34,18 +35,9 @@ const StyledEditorWrapper = styled.div<{ isWidget: boolean }>`
 
 const MemoizedGraph = React.memo(function Layout({ isWidget }: LayoutProps) {
   const json = useConfig((state) => state.json);
-
-  const [nodes, edges, newNodes, newEdges, selectedNode] = useNodeTools(
-    (state) => [
-      state.nodes,
-      state.edges,
-      state.newNodes,
-      state.newEdges,
-      state.selectedNode,
-    ],
-    shallow
-  );
-  const setNodeTools = useNodeTools((state) => state.setNodeTools);
+  const nodes = useGraph((state) => state.nodes);
+  const edges = useGraph((state) => state.edges);
+  const setGraphValue = useGraph((state) => state.setGraphValue);
 
   const [size, setSize] = React.useState({
     width: 2000,
@@ -61,11 +53,11 @@ const MemoizedGraph = React.memo(function Layout({ isWidget }: LayoutProps) {
   React.useEffect(() => {
     const { nodes, edges } = getEdgeNodes(json, expand);
 
-    setNodeTools("nodes", nodes);
-    setNodeTools("edges", edges);
-    setNodeTools("newNodes", nodes);
-    setNodeTools("newEdges", edges);
-  }, [json, expand, setNodeTools]);
+    setGraphValue("nodes", nodes);
+    setGraphValue("edges", edges);
+    setGraphValue("initialNodes", nodes);
+    setGraphValue("initialEdges", edges);
+  }, [expand, json, setGraphValue]);
 
   const onInit = (ref: ReactZoomPanPinchRef) => {
     setConfig("zoomPanPinch", ref);
@@ -76,20 +68,6 @@ const MemoizedGraph = React.memo(function Layout({ isWidget }: LayoutProps) {
       setSize({ width: layout.width, height: layout.height });
   };
 
-  const { selections, onClick, removeSelection } = useSelection({
-    nodes,
-    edges,
-    onSelection: (s) => {
-      if (!isWidget) {
-        if (s[0] === selectedNode) {
-          removeSelection(selectedNode);
-        } else {
-          setNodeTools("selectedNode", s[0]);
-        }
-      }
-    },
-  });
-
   return (
     <StyledEditorWrapper isWidget={isWidget}>
       <TransformWrapper
@@ -119,11 +97,7 @@ const MemoizedGraph = React.memo(function Layout({ isWidget }: LayoutProps) {
             direction={layout}
             key={layout}
             onLayoutChange={onLayoutChange}
-            selections={selections}
             node={(props) => <CustomNode {...props} />}
-            edge={(props) =>
-              newEdges.find((e) => e.id === props.id) ? <Edge /> : <></>
-            }
             zoomable={false}
             readonly
           />

src/hooks/store/useGraph.tsx

@@ -0,0 +1,139 @@
+import create from "zustand";
+import { EdgeData, NodeData } from "reaflow";
+import { Graph } from "src/components/Graph";
+import { findEdgeChildren } from "src/utils/findEdgeChildren";
+import { findNodeChildren } from "src/utils/findNodeChildren";
+
+export interface Graph {
+  initialNodes: NodeData[];
+  initialEdges: EdgeData[];
+  nodes: NodeData[];
+  edges: EdgeData[];
+  collapsedNodes: { [key: string]: NodeData[] };
+  collapsedEdges: { [key: string]: EdgeData[] };
+}
+
+interface GraphActions {
+  setGraphValue: (key: keyof Graph, value: any) => void;
+  toggleNode: (node: NodeData, isExpanded: boolean) => void;
+}
+
+function isNode(element: NodeData | EdgeData) {
+  if ("text" in element) return true;
+  return false;
+}
+
+export const getIncomers = (
+  node: NodeData,
+  nodes: NodeData[],
+  edges: EdgeData[]
+): NodeData[] => {
+  if (!isNode(node)) {
+    return [];
+  }
+
+  const incomersIds = edges.filter((e) => e.to === node.id).map((e) => e.from);
+  return nodes.filter((n) => incomersIds.includes(n.id));
+};
+
+export const getOutgoingEdges = (
+  node: NodeData,
+  edges: EdgeData[]
+): EdgeData[] => {
+  if (!isNode(node)) {
+    return [];
+  }
+
+  return edges.filter((edge) => edge.from === node.id);
+};
+
+export const getIncomingEdges = (
+  node: NodeData,
+  edges: EdgeData[]
+): EdgeData[] => {
+  if (!isNode(node)) {
+    return [];
+  }
+
+  return edges.filter((edge) => edge.to === node.id);
+};
+
+export const getOutgoers = (
+  node: NodeData,
+  nodes: NodeData[],
+  edges: EdgeData[]
+) => {
+  const allOutgoers: NodeData[] = [];
+
+  if (!isNode(node)) {
+    return [];
+  }
+
+  const runner = (currentNode: NodeData) => {
+    const outgoerIds = edges
+      .filter((e) => e.from === currentNode.id)
+      .map((e) => e.to);
+    const nodeList = nodes.filter((n) => outgoerIds.includes(n.id));
+    allOutgoers.push(...nodeList);
+    nodeList.forEach((node) => runner(node));
+  };
+
+  return allOutgoers;
+};
+
+const initialStates: Graph = {
+  initialNodes: [],
+  initialEdges: [],
+  nodes: [],
+  edges: [],
+  collapsedNodes: {},
+  collapsedEdges: {},
+};
+
+const useGraph = create<Graph & GraphActions>((set, get) => ({
+  ...initialStates,
+  setGraphValue: (key, value) =>
+    set((state) => ({
+      ...state,
+      [key]: value,
+    })),
+  toggleNode: (node, isExpanded) =>
+    set((state) => {
+      const childrenNodes = findNodeChildren(
+        node.id,
+        state.initialNodes,
+        state.initialEdges
+      );
+      const childrenEdges = findEdgeChildren(
+        node.id,
+        childrenNodes,
+        state.initialEdges
+      );
+
+      const expand = () => {
+        if (isExpanded) {
+          return {
+            nodes: state.nodes.filter(
+              (sNode) => !childrenNodes.map((n) => n.id).includes(sNode.id)
+            ),
+            edges: state.edges.filter(
+              (sEdge) => !childrenEdges.map((n) => n.id).includes(sEdge.id)
+            ),
+          };
+        }
+
+        return {
+          nodes: state.nodes.concat(childrenNodes),
+          edges: state.edges.concat(childrenEdges),
+        };
+      };
+
+      return {
+        ...state,
+        nodes: expand().nodes,
+        edges: expand().edges,
+      };
+    }),
+}));
+
+export default useGraph;

src/utils/findEdgeChildren.ts

@@ -1,15 +1,17 @@
 import { NodeData, EdgeData } from "reaflow/dist/types";
 
+export const findEdgeChildren = (
+  selectedNode: string,
+  connections: NodeData[],
+  edges: EdgeData[]
+) => {
+  const nodeIds = connections.map((n) => n.id);
 
-export const findEdgeChildren = (selectedNode: string, connections: NodeData[], edges: EdgeData[]) => {
+  nodeIds.push(selectedNode);
+  const newEdges = edges.filter(
+    (e) =>
+      nodeIds.includes(e.from as string) && nodeIds.includes(e.to as string)
+  );
 
-    const nodeIds = connections.map((n) => n.id);
-    nodeIds.push(selectedNode);
-    const newEdges = edges.filter(
-      (e) =>
-        nodeIds.includes(e.from as string) && nodeIds.includes(e.to as string)
-    );
-
-    return newEdges;
-
-  };
+  return newEdges;
+};