Feat/layout nodes and edges

dnd-ui/webview-ui/src/App.tsx

@@ -6,6 +6,7 @@ import { Provider } from './context/store';
 import InitGlobalState from './InitGlobalState';
 import { ToastContainer } from 'react-toastify';
 import 'vite/modulepreload-polyfill';
+import GroupNodeProvider from './context/CollapsibleContext';
 
 declare global {
   interface Window {
@@ -18,12 +19,14 @@ declare global {
 function App() {
   return (
     <Provider>
-      <InitGlobalState>
-        <ReactFlowProvider>
-          <Flow />
-        </ReactFlowProvider>
-        <ToastContainer />
-      </InitGlobalState>
+      <GroupNodeProvider>
+        <InitGlobalState>
+          <ReactFlowProvider>
+            <Flow />
+          </ReactFlowProvider>
+          <ToastContainer />
+        </InitGlobalState>
+      </GroupNodeProvider>
     </Provider>
   );
 }

dnd-ui/webview-ui/src/components/NodeMenuPanel.tsx

@@ -1,12 +1,13 @@
 // @ts-nocheck
-import { Dispatch, memo, SetStateAction, useRef } from 'react';
+import { faAngleDown, faPlus } from '@fortawesome/free-solid-svg-icons';
 import { FontAwesomeIcon } from '@fortawesome/react-fontawesome';
+import { Dispatch, memo, SetStateAction, useRef } from 'react';
 import { Panel } from 'reactflow';
-import defaultNodes from '../nodes/defaultNode.json';
-import { RFNode } from '../models/nodeFactory';
+import { useCollapsibleNodes } from '../context/CollapsibleContext';
 import { useStore } from '../context/store';
 import { useStaticClickAway } from '../hooks/useClickOutside';
-import { faAngleDown, faPlus } from '@fortawesome/free-solid-svg-icons';
+import { RFNode } from '../models/nodeFactory';
+import defaultNodes from '../nodes/defaultNode.json';
 
 interface INodeMenuPanel {
   setShowMenu: Dispatch<SetStateAction<boolean>>;
@@ -23,6 +24,7 @@ function NodeMenuPanel({
 }: INodeMenuPanel) {
   const nodeMenuRef = useRef(null);
   const addNodeRef = useRef(null);
+  const { setGroupNodes } = useCollapsibleNodes();
   const [takeSnapshot] = useStore((store) => store.takeSnapshot);
 
   useStaticClickAway(nodeMenuRef, () => setShowMenu(false), addNodeRef);
@@ -38,6 +40,15 @@ function NodeMenuPanel({
       },
       position: { x: event.clientX - curX + 20, y: event.clientY - curY + 20 },
     });
+
+    if (newNode.type === 'group') {
+      const { width, height } = newNode;
+
+      setGroupNodes((prev) => {
+        return { ...prev, [newNode.id]: { width, height } };
+      });
+    }
+
     await setNodes((prev) => [...prev, newNode]);
     setShowMenu(false);
     takeSnapshot();
@@ -78,6 +89,7 @@ function NodeMenuPanel({
                 >
                   <button
                     id="anyNode"
+                    className="w-full text-left"
                     onClick={(event) => handleClick(event, node.type)}
                   >
                     {name.replace('Node', '')} Node
@@ -91,6 +103,15 @@ function NodeMenuPanel({
                 CTFlow Recorder
               </button>
             </div>
+            <div className="hover:bg-slate-200 p-2 rounded">
+              <button
+                id="group"
+                className="w-full text-left"
+                onClick={(event) => handleClick(event, 'group')}
+              >
+                Group
+              </button>
+            </div>
 
             {/* <div className="hover:bg-slate-200 p-2 rounded">
               <button id="codeInjectionNode" onClick={handleClick}>

dnd-ui/webview-ui/src/context/CollapsibleContext.tsx

@@ -0,0 +1,45 @@
+import { createContext, useContext } from 'react';
+
+interface IGroupNodeContext {
+  groupNodes: Record<string, any>;
+  setGroupNodes: (prev: any) => void;
+}
+
+const GroupNodeContext = createContext<IGroupNodeContext>({
+  groupNodes: {},
+  setGroupNodes: () => {},
+});
+
+const GroupNodeProvider = (props: { children: React.ReactNode }) => {
+  const { children } = props;
+  const value = {
+    groupNodes: {},
+    setGroupNodes,
+  };
+
+  function setGroupNodes(callback: (prev: any) => any) {
+    const newValue = callback(value.groupNodes);
+    console.log('newValue:', newValue);
+    value.groupNodes = newValue;
+  }
+
+  return (
+    <GroupNodeContext.Provider value={value}>
+      {children}
+    </GroupNodeContext.Provider>
+  );
+};
+
+export const useCollapsibleNodes = () => {
+  const context = useContext(GroupNodeContext);
+
+  if (!context) {
+    throw new Error(
+      'useCollapsibleNodes must be used within GroupNodeProvider'
+    );
+  }
+
+  return context;
+};
+
+export default GroupNodeProvider;

dnd-ui/webview-ui/src/hooks/useDetachNodes.tsx

@@ -0,0 +1,36 @@
+import { useCallback } from 'react';
+import { useReactFlow, useStoreApi } from 'reactflow';
+
+function useDetachNodes() {
+  const { setNodes } = useReactFlow();
+  const store = useStoreApi();
+
+  const detachNodes = useCallback(
+    (ids: string[], removeParentId?: string) => {
+      const { nodeInternals } = store.getState();
+      const nextNodes = Array.from(nodeInternals.values()).map((n) => {
+        if (ids.includes(n.id) && n.parentNode) {
+          const parentNode = nodeInternals.get(n.parentNode);
+          return {
+            ...n,
+            position: {
+              x: n.position.x + (parentNode?.positionAbsolute?.x ?? 0),
+              y: n.position.y + (parentNode?.positionAbsolute?.y ?? 0),
+            },
+            extent: undefined,
+            parentNode: undefined,
+          };
+        }
+        return n;
+      });
+      setNodes(
+        nextNodes.filter((n) => !removeParentId || n.id !== removeParentId)
+      );
+    },
+    [setNodes, store]
+  );
+
+  return detachNodes;
+}
+
+export default useDetachNodes;

dnd-ui/webview-ui/src/hooks/useExpandCollapse.ts

@@ -0,0 +1,99 @@
+import { useMemo, useRef } from 'react';
+import { Node, Edge } from 'reactflow';
+import Dagre from '@dagrejs/dagre';
+// import { useStore } from '../context/store';
+
+type NodeData = {
+  collapse: boolean;
+};
+
+export type UseExpandCollapseOptions = {
+  layoutNodes?: boolean;
+  treeWidth?: number;
+  treeHeight?: number;
+};
+
+function filterCollapsedChildren(
+  dagre: Dagre.graphlib.Graph,
+  node: Node<NodeData>,
+  nodes: Node[]
+) {
+  if (!node.data.collapse) return;
+  const children = nodes.filter((n) => n.parentNode === node.id);
+
+  for (const child of children) {
+    dagre.removeNode(child.id);
+  }
+}
+
+function useExpandCollapse(
+  nodes: Node[],
+  edges: Edge[],
+  {
+    layoutNodes = true,
+    treeWidth = 220,
+    treeHeight = 100,
+  }: UseExpandCollapseOptions = {}
+): { nodes: Node[]; edges: Edge[] } {
+  return useMemo(() => {
+    if (nodes.length === 0 && edges.length === 0) return { nodes, edges };
+
+    if (!layoutNodes) return { nodes, edges };
+
+    // 1. Create a new instance of `Dagre.graphlib.Graph` and set some default
+    // properties.
+    const dagre = new Dagre.graphlib.Graph()
+      .setDefaultEdgeLabel(() => ({}))
+      .setGraph({ rankdir: 'TB' });
+
+    // 2. Add each node and edge to the dagre graph. Instead of using each node's
+    // intrinsic width and height, we tell dagre to use the `treeWidth` and
+    // `treeHeight` values. This lets you control the space between nodes.
+    for (const node of nodes) {
+      dagre.setNode(node.id, {
+        width: treeWidth,
+        height: treeHeight,
+        data: node.data,
+      });
+    }
+
+    for (const edge of edges) {
+      dagre.setEdge(edge.source, edge.target);
+    }
+
+    // 3. Iterate over the nodes *again* to determine which ones should be hidden
+    // based on expand/collapse state. Hidden nodes are removed from the dagre
+    // graph entirely.
+    for (const node of nodes) {
+      filterCollapsedChildren(dagre, node, nodes);
+    }
+
+    const newNodes = nodes.flatMap((node) => {
+      if (!dagre.hasNode(node.id)) return [];
+
+      if (node.data.collapse) {
+        return [{ ...node, style: { width: 130, height: 50 } }];
+      }
+
+      const data = { ...node.data };
+
+      return [{ ...node, data }];
+    });
+
+    return {
+      //   nodes,
+      nodes: newNodes,
+      edges,
+    };
+  }, [
+    nodes,
+    edges,
+    layoutNodes,
+    treeWidth,
+    treeHeight,
+    nodes.length,
+    edges.length,
+  ]);
+}
+
+export default useExpandCollapse;