Save groups in the graph

packages/graph-editor/src/annotations/index.ts

@@ -19,6 +19,8 @@ export const hidden = 'ui.hidden';
 //Used on nodes to indicate meta from a ui
 export const xpos = 'ui.position.x';
 export const ypos = 'ui.position.y';
+export const width = 'ui.dimension.width';
+export const height = 'ui.dimension.height';
 
 //Used on nodes and graph
 export const title = 'ui.title';

packages/graph-editor/src/components/commandPalette/index.tsx

@@ -21,9 +21,9 @@ export interface ICommandMenu {
   items: DropPanelStore;
   handleSelectNewNodeType: (node: NodeRequest) =>
     | {
-        graphNode: Node;
-        flowNode: ReactFlowNode;
-      }
+      graphNode: Node;
+      flowNode: ReactFlowNode;
+    }
     | undefined;
 }
 
@@ -74,14 +74,16 @@ const CommandMenuGroup = observer(
       <Command.Group
         key={group.key}
         heading={
-          <Stack align="center">
+          <Stack align="center" gap={2} css={{ margin: '$3 0 $2' }}>
             {group.icon}
-            {group.title}
+            <Text size="small" css={{ fontWeight: '$sansMedium' }}>
+              {group.title}
+            </Text>
           </Stack>
         }
       >
         {group.items.map((item) => (
-          <CommandItem item={item} handleSelectItem={handleSelectItem} />
+          <CommandItem key={item.type} item={item} handleSelectItem={handleSelectItem} />
         ))}
       </Command.Group>
     );
@@ -183,7 +185,7 @@ const CommandMenu = ({ items, handleSelectNewNodeType }: ICommandMenu) => {
           direction="row"
           css={{
             overflowY: 'scroll',
-            maxHeight: '450px',
+            maxHeight: '500px',
             scrollbarColor: 'var(--colors-bgSubtle) transparent',
             scrollbarWidth: 'thin',
           }}

packages/graph-editor/src/components/contextMenus/selectionContextMenu.tsx

@@ -1,81 +1,123 @@
-import { Edge, Graph } from '@tokens-studio/graph-engine';
+import { Edge, Graph, Node as GraphNode } from '@tokens-studio/graph-engine';
+import { GROUP } from '@/ids.js';
+import { GROUP_NODE_PADDING } from '@/constants.js';
 import { Item, Menu, Separator } from 'react-contexify';
-import { Node, getRectOfNodes, useReactFlow, useStoreApi } from 'reactflow';
-import { NodeTypes } from '../flow/types.js';
-import { getId } from '../flow/utils.js';
+import { Node, getNodesBounds, useReactFlow, useStoreApi } from 'reactflow';
+import { height, width, xpos, ypos } from '@/annotations/index.js';
 import { useAction } from '@/editor/actions/provider.js';
 import { useLocalGraph } from '@/hooks/index.js';
 import { v4 as uuid } from 'uuid';
-import React, { useCallback } from 'react';
+import React, { useCallback, useMemo } from 'react';
 
 export type INodeContextMenuProps = {
   id: string;
   nodes: Node[];
 };
 
-const padding = 25;
-
 export const SelectionContextMenu = ({ id, nodes }: INodeContextMenuProps) => {
   const reactFlowInstance = useReactFlow();
   const graph = useLocalGraph();
   const store = useStoreApi();
   const createNode = useAction('createNode');
   const duplicateNodes = useAction('duplicateNodes');
 
-  //Note that we use a filter here to prevent getting nodes that have a parent node, ie are part of a group
+  const reactFlowNodes = reactFlowInstance.getNodes();
+
+  // Note that we use a filter here to prevent getting nodes that have a parent node, ie are part of a group
   const selectedNodes = nodes.filter(
-    (node) => node.selected && !node.parentNode,
+    (node) => node.selected && !node.parentId,
   );
   const selectedNodeIds = selectedNodes.map((node) => node.id);
 
   const onGroup = useCallback(() => {
-    const rectOfNodes = getRectOfNodes(nodes);
-    const groupId = getId('group');
+    const bounds = getNodesBounds(nodes);
     const parentPosition = {
-      x: rectOfNodes.x,
-      y: rectOfNodes.y,
+      x: bounds.x,
+      y: bounds.y,
     };
-    const groupNode = {
-      id: groupId,
-      type: NodeTypes.GROUP,
-      position: parentPosition,
-      style: {
-        width: rectOfNodes.width + padding * 2,
-        height: rectOfNodes.height + padding * 2,
-      },
-      data: {
-        expandable: true,
-        expanded: true,
-      },
-    } as Node;
 
     store.getState().resetSelectedElements();
     store.setState({ nodesSelectionActive: false });
+
+    const newNodes = createNode({
+      type: GROUP,
+      position: parentPosition,
+    });
+
+    if (!newNodes) {
+      return;
+    }
+
+    const { flowNode } = newNodes;
+
     reactFlowInstance.setNodes((nodes) => {
-      //Note that group nodes should always occur before their parents
-      return [groupNode].concat(
-        nodes.map((node) => {
+      // Note that group nodes should always occur before their children
+      return [{
+        ...flowNode,
+        dragHandle: undefined,
+        style: {
+          width: bounds.width + GROUP_NODE_PADDING * 2,
+          height: bounds.height + GROUP_NODE_PADDING * 2,
+        },
+        data: {
+          expandable: true,
+          expanded: true,
+        }
+      } as Node]
+        .concat(nodes)
+        .map((node) => {
           if (selectedNodeIds.includes(node.id)) {
             return {
               ...node,
               position: {
-                x: node.position.x - parentPosition.x + padding,
-                y: node.position.y - parentPosition.y + padding,
+                x: node.position.x - parentPosition.x + GROUP_NODE_PADDING,
+                y: node.position.y - parentPosition.y + GROUP_NODE_PADDING,
               },
               extent: 'parent' as const,
-              parentNode: groupId,
+              parentId: flowNode.id,
             };
           }
 
           return node;
-        }),
-      );
+        });
+    });
+
+    const reactFlowNodesMap = new Map<string, Node>(
+      reactFlowNodes.map((node) => [node.id, node]),
+    );
+
+    // Set annotations for all items in the group
+    nodes.forEach((node) => {
+      const graphNode = graph.getNode(node.id);
+      if (graphNode) {
+        graphNode.annotations[xpos] = node.position.x - parentPosition.x + GROUP_NODE_PADDING;
+        graphNode.annotations[ypos] = node.position.y - parentPosition.y + GROUP_NODE_PADDING;
+        graphNode.annotations[width] = reactFlowNodesMap.get(node.id)?.width || 200;
+        graphNode.annotations[height] = reactFlowNodesMap.get(node.id)?.height || 100;
+        graphNode.annotations['parentId'] = flowNode.id;
+      }
     });
-  }, [nodes, reactFlowInstance, selectedNodeIds, store]);
+
+  }, [createNode, graph, nodes, reactFlowInstance, reactFlowNodes, selectedNodeIds, store]);
 
   const onCreateSubgraph = useCallback(() => {
-    //We need to work out which nodes do not have parents in the selection
+    // Get all selected node ids, including children of groups
+    const selectedNodeIds = selectedNodes
+      .reduce((acc, node) => {
+        if (node.type !== GROUP) {
+          return [...acc, node.id];
+        }
+
+        const children = reactFlowNodes
+          .filter((n) => n.parentId === node.id)
+          .map((x) => x.id);
 
+        if (children.length > 0) {
+          return [...acc, node.id, ...children];
+        }
+
+        return acc;
+      }, [] as string[]);
     const lookup = new Set(selectedNodeIds);
 
     //Lets create a new subgraph node
@@ -95,17 +137,17 @@ export const SelectionContextMenu = ({ id, nodes }: INodeContextMenuProps) => {
       y: position.y / selectedNodes.length,
     };
 
-    const nodes = createNode({
+    const newNodes = createNode({
       type: 'studio.tokens.generic.subgraph',
       position: finalPosition,
     });
 
     //Request failed in some way
-    if (!nodes) {
+    if (!newNodes) {
       return;
     }
 
-    const { graphNode, flowNode } = nodes;
+    const { graphNode, flowNode } = newNodes;
 
     //@ts-expect-error
     const internalGraph = graphNode._innerGraph as unknown as Graph;
@@ -294,9 +336,11 @@ export const SelectionContextMenu = ({ id, nodes }: INodeContextMenuProps) => {
     duplicateNodes(selectedNodeIds);
   };
 
+  const hasGroup = selectedNodes.some((node) => node.type === GROUP);
+
   return (
     <Menu id={id}>
-      <Item onClick={onGroup}>Create group</Item>
+      {!hasGroup && <Item onClick={onGroup}>Create group</Item>}
       <Item onClick={onCreateSubgraph}>Create Subgraph</Item>
       <Separator />
       <Item onClick={onDuplicate}>Duplicate</Item>

packages/graph-editor/src/components/flow/nodes/groupNode.tsx

@@ -1,34 +1,36 @@
 import { Button, Stack } from '@tokens-studio/ui';
+import { GROUP_NODE_PADDING } from '@/constants.js';
 import {
   NodeProps,
   NodeToolbar,
-  getRectOfNodes,
+  getNodesBounds,
   useReactFlow,
   useStore,
   useStoreApi,
 } from 'reactflow';
 import { NodeResizer } from '@reactflow/node-resizer';
 import { useCallback } from 'react';
+import { useLocalGraph } from '@/context/graph.js';
 import React from 'react';
 import useDetachNodes from '../../../hooks/useDetachNodes.js';
 
 const lineStyle = { borderColor: 'white' };
-const padding = 25;
 
 function GroupNode(props: NodeProps) {
   const { id, data } = props;
   const store = useStoreApi();
   const { deleteElements } = useReactFlow();
   const detachNodes = useDetachNodes();
+  const graph = useLocalGraph()
   const { minWidth, minHeight, hasChildNodes } = useStore((store) => {
     const childNodes = Array.from(store.nodeInternals.values()).filter(
-      (n) => n.parentNode === id,
+      (n) => n.parentId === id,
     );
-    const rect = getRectOfNodes(childNodes);
+    const bounds = getNodesBounds(childNodes);
 
     return {
-      minWidth: rect.width + padding * 2,
-      minHeight: rect.height + padding * 2,
+      minWidth: bounds.width + GROUP_NODE_PADDING * 2,
+      minHeight: bounds.height + GROUP_NODE_PADDING * 2,
       hasChildNodes: childNodes.length > 0,
     };
   }, isEqual);
@@ -39,11 +41,20 @@ function GroupNode(props: NodeProps) {
 
   const onDetach = useCallback(() => {
     const childNodeIds = Array.from(store.getState().nodeInternals.values())
-      .filter((n) => n.parentNode === id)
+      .filter((n) => n.parentId === id)
       .map((n) => n.id);
 
     detachNodes(childNodeIds, id);
-  }, [detachNodes, id, store]);
+
+    graph.removeNode(id);
+
+    childNodeIds.forEach((nodeId) => {
+      const graphNode = graph.getNode(nodeId);
+      if (graphNode) {
+        delete graphNode.annotations['parentId'];
+      }
+    });
+  }, [detachNodes, graph, id, store]);
 
   return (
     <div

packages/graph-editor/src/components/flow/utils.ts

@@ -1,12 +1,12 @@
+import { GROUP } from '@/ids.js';
 import { Node } from 'reactflow';
-import { NodeTypes } from './types.js';
 
 // we have to make sure that parent nodes are rendered before their children
 export const sortNodes = (a: Node, b: Node): number => {
   if (a.type === b.type) {
     return 0;
   }
-  return a.type === NodeTypes.GROUP && b.type !== NodeTypes.GROUP ? -1 : 1;
+  return a.type === GROUP && b.type !== GROUP ? -1 : 1;
 };
 
 export const getId = (prefix = 'node') => `${prefix}_${Math.random() * 10000}`;

packages/graph-editor/src/components/panels/dropPanel/data.tsx

@@ -1,3 +1,4 @@
+import { GROUP } from '@/ids.js';
 import { nodes } from '@tokens-studio/graph-engine';
 import { observable } from 'mobx';
 
@@ -99,36 +100,39 @@ function CapitalCase(string) {
   return string.charAt(0).toUpperCase() + string.slice(1);
 }
 
+const nodesToIgnoreInPanel = [GROUP];
+
 export const defaultPanelGroupsFactory = (): DropPanelStore => {
   const auto = Object.values<PanelGroup>(
-    nodes.reduce(
-      (acc, node) => {
-        const defaultGroup = node.type.split('.');
-        const groups = node.groups || [defaultGroup[defaultGroup.length - 2]];
+    nodes.filter(node => !nodesToIgnoreInPanel.includes(node.type))
+      .reduce(
+        (acc, node) => {
+          const defaultGroup = node.type.split('.');
+          const groups = node.groups || [defaultGroup[defaultGroup.length - 2]];
 
-        groups.forEach((group) => {
-          //If the group does not exist, create it
-          if (!acc[group]) {
-            acc[group] = new PanelGroup({
-              title: CapitalCase(group),
-              key: group,
-              items: [],
-            });
-          }
-          acc[group].items.push(
-            new PanelItem({
-              type: node.type,
-              text: CapitalCase(
-                node.title || defaultGroup[defaultGroup.length - 1],
-              ),
-              description: node.description,
-            }),
-          );
-        });
-        return acc;
-      },
-      {} as Record<string, PanelGroup>,
-    ),
+          groups.forEach((group) => {
+            //If the group does not exist, create it
+            if (!acc[group]) {
+              acc[group] = new PanelGroup({
+                title: CapitalCase(group),
+                key: group,
+                items: [],
+              });
+            }
+            acc[group].items.push(
+              new PanelItem({
+                type: node.type,
+                text: CapitalCase(
+                  node.title || defaultGroup[defaultGroup.length - 1],
+                ),
+                description: node.description,
+              }),
+            );
+          });
+          return acc;
+        },
+        {} as Record<string, PanelGroup>,
+      ),
   );
 
   return new DropPanelStore(auto);