wip exporting template

koonpeng · koonpeng · commit c7902faf2c37 · 2025-07-29T09:30:38.000Z
Signed-off-by: Teo Koon Peng &lt;teokoonpeng@gmail.com&gt;
diff --git a/diagram-editor/frontend/add-operation.tsx b/diagram-editor/frontend/add-operation.tsx
@@ -5,7 +5,6 @@ import {
   type XYPosition,
 } from '@xyflow/react';
 import React from 'react';
-import { v4 as uuidv4 } from 'uuid';
 import { EditorMode, useEditorMode } from './editor-mode';
 import type { DiagramEditorNode } from './nodes';
 import {
@@ -30,7 +29,13 @@ import {
   UnzipIcon,
 } from './nodes';
 import type { DiagramOperation } from './types/api';
-import { calculateScopeBounds, LAYOUT_OPTIONS } from './utils/layout';
+import {
+  createOperationNode,
+  createScopeNode,
+  createSectionBufferNode,
+  createSectionInputNode,
+  createSectionOutputNode,
+} from './utils/create-node';
 import { joinNamespaces, ROOT_NAMESPACE } from './utils/namespace';
 
 const StyledOperationButton = styled(Button)({
@@ -44,18 +49,12 @@ export interface AddOperationProps {
 }
 
 function createSectionInputChange(
-  remappedId: string,
   targetId: string,
   position: XYPosition,
 ): NodeAddChange<SectionInputNode> {
   return {
     type: 'add',
-    item: {
-      id: uuidv4(),
-      type: 'sectionInput',
-      data: { remappedId, targetId },
-      position,
-    },
+    item: createSectionInputNode(targetId, position),
   };
 }
 
@@ -65,28 +64,17 @@ function createSectionOutputChange(
 ): NodeAddChange<SectionOutputNode> {
   return {
     type: 'add',
-    item: {
-      id: uuidv4(),
-      type: 'sectionOutput',
-      data: { outputId },
-      position,
-    },
+    item: createSectionOutputNode(outputId, position),
   };
 }
 
 function createSectionBufferChange(
-  remappedId: string,
   targetId: string,
   position: XYPosition,
 ): NodeAddChange<SectionBufferNode> {
   return {
     type: 'add',
-    item: {
-      id: uuidv4(),
-      type: 'sectionBuffer',
-      data: { remappedId, targetId },
-      position,
-    },
+    item: createSectionBufferNode(targetId, position),
   };
 }
 
@@ -97,82 +85,15 @@ function createNodeChange(
   op: DiagramOperation,
 ): NodeAddChange<DiagramEditorNode>[] {
   if (op.type === 'scope') {
-    const scopeId = uuidv4();
-    const children: NodeAddChange<DiagramEditorNode>[] = [
-      {
-        type: 'add',
-        item: {
-          id: uuidv4(),
-          type: 'start',
-          position: {
-            x: LAYOUT_OPTIONS.scopePadding.leftRight,
-            y: LAYOUT_OPTIONS.scopePadding.topBottom,
-          },
-          data: {
-            namespace: joinNamespaces(namespace, scopeId),
-          },
-          parentId: scopeId,
-        },
-      },
-      {
-        type: 'add',
-        item: {
-          id: uuidv4(),
-          type: 'terminate',
-          position: {
-            x: LAYOUT_OPTIONS.scopePadding.leftRight,
-            y: LAYOUT_OPTIONS.scopePadding.topBottom * 5,
-          },
-          data: {
-            namespace: joinNamespaces(namespace, scopeId),
-          },
-          parentId: scopeId,
-        },
-      },
-    ];
-    const scopeBounds = calculateScopeBounds(
-      children.map((c) => c.item.position),
+    return createScopeNode(namespace, parentId, newNodePosition, op).map(
+      (node) => ({ type: 'add', item: node }),
     );
-    return [
-      {
-        type: 'add',
-        item: {
-          id: scopeId,
-          type: 'scope',
-          position: {
-            x: newNodePosition.x + scopeBounds.x,
-            y: newNodePosition.y + scopeBounds.y,
-          },
-          data: {
-            namespace,
-            opId: uuidv4(),
-            op,
-          },
-          width: scopeBounds.width,
-          height: scopeBounds.height,
-          zIndex: -1,
-          parentId,
-        },
-      },
-      ...children,
-    ];
   }
 
   return [
     {
       type: 'add',
-      item: {
-        id: uuidv4(),
-        type: op.type,
-        position: newNodePosition,
-        data: {
-          namespace,
-          opId: uuidv4(),
-          op,
-        },
-
-        parentId,
-      },
+      item: createOperationNode(namespace, parentId, newNodePosition, op),
     },
   ];
 }
@@ -200,13 +121,7 @@ function AddOperation({ parentId, newNodePosition, onAdd }: AddOperationProps) {
           <StyledOperationButton
             startIcon={<SectionInputIcon />}
             onClick={() => {
-              onAdd?.([
-                createSectionInputChange(
-                  'new_input',
-                  'new_input',
-                  newNodePosition,
-                ),
-              ]);
+              onAdd?.([createSectionInputChange('new_input', newNodePosition)]);
             }}
           >
             Section Input
@@ -231,11 +146,7 @@ function AddOperation({ parentId, newNodePosition, onAdd }: AddOperationProps) {
             startIcon={<SectionBufferIcon />}
             onClick={() => {
               onAdd?.([
-                createSectionBufferChange(
-                  'new_buffer',
-                  'new_buffer',
-                  newNodePosition,
-                ),
+                createSectionBufferChange('new_buffer', newNodePosition),
               ]);
             }}
           >
diff --git a/diagram-editor/frontend/utils/create-node.ts b/diagram-editor/frontend/utils/create-node.ts
@@ -0,0 +1,133 @@
+import type { XYPosition } from '@xyflow/react';
+import { v4 as uuidv4 } from 'uuid';
+import type {
+  DiagramEditorNode,
+  OperationNode,
+  SectionBufferNode,
+  SectionInputNode,
+  SectionOutputNode,
+} from '../nodes';
+import type { DiagramOperation } from '../types/api';
+import { calculateScopeBounds, LAYOUT_OPTIONS } from './layout';
+import { joinNamespaces } from './namespace';
+
+/**
+ * Create a section input node. For simplicity, remapping the input is not supported, `targetId`
+ * must point to the id of an operation in the template.
+ */
+export function createSectionInputNode(
+  targetId: string,
+  position: XYPosition,
+): SectionInputNode {
+  return {
+    id: uuidv4(),
+    type: 'sectionInput',
+    data: { remappedId: targetId, targetId },
+    position,
+  };
+}
+
+export function createSectionOutputNode(
+  outputId: string,
+  position: XYPosition,
+): SectionOutputNode {
+  return {
+    id: uuidv4(),
+    type: 'sectionOutput',
+    data: { outputId },
+    position,
+  };
+}
+
+/**
+ * Create a section input node. For simplicity, remapping the input is not supported, `targetId`
+ * must point to the id of a `buffer` operation in the template.
+ */
+export function createSectionBufferNode(
+  targetId: string,
+  position: XYPosition,
+): SectionBufferNode {
+  return {
+    id: uuidv4(),
+    type: 'sectionBuffer',
+    data: { remappedId: targetId, targetId },
+    position,
+  };
+}
+
+export function createOperationNode(
+  namespace: string,
+  parentId: string | undefined,
+  position: XYPosition,
+  op: Exclude<DiagramOperation, { type: 'scope' }>,
+): OperationNode {
+  return {
+    id: uuidv4(),
+    type: op.type,
+    position,
+    data: {
+      namespace,
+      opId: uuidv4(),
+      op,
+    },
+    ...(parentId && { parentId }),
+  };
+}
+
+export function createScopeNode(
+  namespace: string,
+  parentId: string | undefined,
+  position: XYPosition,
+  op: DiagramOperation & { type: 'scope' },
+): DiagramEditorNode[] {
+  const scopeId = uuidv4();
+  const children: DiagramEditorNode[] = [
+    {
+      id: uuidv4(),
+      type: 'start',
+      position: {
+        x: LAYOUT_OPTIONS.scopePadding.leftRight,
+        y: LAYOUT_OPTIONS.scopePadding.topBottom,
+      },
+      data: {
+        namespace: joinNamespaces(namespace, scopeId),
+      },
+      parentId: scopeId,
+    },
+    {
+      id: uuidv4(),
+      type: 'terminate',
+      position: {
+        x: LAYOUT_OPTIONS.scopePadding.leftRight,
+        y: LAYOUT_OPTIONS.scopePadding.topBottom * 5,
+      },
+      data: {
+        namespace: joinNamespaces(namespace, scopeId),
+      },
+      parentId: scopeId,
+    },
+  ];
+  const scopeBounds = calculateScopeBounds(
+    children.map((child) => child.position),
+  );
+  return [
+    {
+      id: scopeId,
+      type: 'scope',
+      position: {
+        x: position.x + scopeBounds.x,
+        y: position.y + scopeBounds.y,
+      },
+      data: {
+        namespace,
+        opId: uuidv4(),
+        op,
+      },
+      width: scopeBounds.width,
+      height: scopeBounds.height,
+      zIndex: -1,
+      ...(parentId && { parentId }),
+    },
+    ...children,
+  ];
+}
diff --git a/diagram-editor/frontend/utils/export-diagram.test.ts b/diagram-editor/frontend/utils/export-diagram.test.ts
@@ -1,6 +1,12 @@
 import { NodeManager } from '../node-manager';
 import { START_ID } from '../nodes';
-import { exportDiagram } from './export-diagram';
+import {
+  createOperationNode,
+  createSectionBufferNode,
+  createSectionInputNode,
+  createSectionOutputNode,
+} from './create-node';
+import { exportDiagram, exportTemplate } from './export-diagram';
 import { loadDiagramJson } from './load-diagram';
 import { joinNamespaces, ROOT_NAMESPACE } from './namespace';
 import testDiagram from './test-data/test-diagram.json';
@@ -43,3 +49,36 @@ test('export diagram with scope', () => {
   diagram = exportDiagram(nodeManager, edges, {});
   expect(diagram.ops.scope.start).toBe('mul4');
 });
+
+test('export diagram with templates', () => {
+  const nodes = [
+    createSectionInputNode('test_input', { x: 0, y: 0 }),
+    createSectionOutputNode('test_output', { x: 0, y: 0 }),
+    createSectionBufferNode('test_buffer', { x: 0, y: 0 }),
+    createOperationNode(
+      ROOT_NAMESPACE,
+      undefined,
+      { x: 0, y: 0 },
+      { type: 'node', builder: 'test', next: 'test_output' },
+    ),
+    createOperationNode(
+      ROOT_NAMESPACE,
+      undefined,
+      { x: 0, y: 0 },
+      { type: 'buffer' },
+    ),
+  ];
+  const template = exportTemplate(new NodeManager(nodes), []);
+
+  if (typeof template.inputs !== 'object' || Array.isArray(template.inputs)) {
+    throw new Error('expected template inputs to be a mapping');
+  }
+  expect(template.inputs.test_input).toBe('test_input');
+
+  expect(template.outputs?.[0]).toBe('test_output');
+
+  if (typeof template.buffers !== 'object' || Array.isArray(template.buffers)) {
+    throw new Error('expected template buffers to be a mapping');
+  }
+  expect(template.buffers.test_buffer).toBe('test_buffer');
+});
