more fixes for templates

koonpeng · koonpeng · commit 0a624fb83310 · 2025-07-28T04:52:08.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
@@ -1,7 +1,7 @@
 import { Button, ButtonGroup, styled } from '@mui/material';
 import {
-  useReactFlow,
   type NodeAddChange,
+  useReactFlow,
   type XYPosition,
 } from '@xyflow/react';
 import React from 'react';
@@ -46,15 +46,16 @@ export interface AddOperationProps {
 
 function createSectionInterfaceChange(
   type: SectionInterfaceNodeTypes,
-  connectKey: string,
+  remappedId: string,
+  targetId: string,
   position: XYPosition,
 ): NodeAddChange<SectionInterfaceNode> {
   return {
     type: 'add',
     item: {
       id: uuidv4(),
       type,
-      data: { namespace: ROOT_NAMESPACE, connectKey },
+      data: { namespace: ROOT_NAMESPACE, remappedId, targetId },
       position,
     },
   };
@@ -172,50 +173,57 @@ function AddOperation({ parentId, newNodePosition, onAdd }: AddOperationProps) {
     >
       {editorMode.mode === EditorMode.Template &&
         namespace === ROOT_NAMESPACE && (
-          <>
-            <StyledOperationButton
-              startIcon={<SectionInputIcon />}
-              onClick={() => {
-                onAdd?.([
-                  createSectionInterfaceChange(
-                    'sectionInput',
-                    'input',
-                    newNodePosition,
-                  ),
-                ]);
-              }}
-            >
-              Section Input
-            </StyledOperationButton>
-            <StyledOperationButton
-              startIcon={<SectionOutputIcon />}
-              onClick={() => {
-                onAdd?.([
-                  createSectionInterfaceChange(
-                    'sectionOutput',
-                    'output',
-                    newNodePosition,
-                  ),
-                ]);
-              }}
-            >
-              Section Output
-            </StyledOperationButton>
-            <StyledOperationButton
-              startIcon={<SectionBufferIcon />}
-              onClick={() => {
-                onAdd?.([
-                  createSectionInterfaceChange(
-                    'sectionBuffer',
-                    'buffer',
-                    newNodePosition,
-                  ),
-                ]);
-              }}
-            >
-              Section Buffer
-            </StyledOperationButton>
-          </>
+          <StyledOperationButton
+            startIcon={<SectionInputIcon />}
+            onClick={() => {
+              onAdd?.([
+                createSectionInterfaceChange(
+                  'sectionInput',
+                  'input',
+                  'input',
+                  newNodePosition,
+                ),
+              ]);
+            }}
+          >
+            Section Input
+          </StyledOperationButton>
+        )}
+      {editorMode.mode === EditorMode.Template &&
+        namespace === ROOT_NAMESPACE && (
+          <StyledOperationButton
+            startIcon={<SectionOutputIcon />}
+            onClick={() => {
+              onAdd?.([
+                createSectionInterfaceChange(
+                  'sectionOutput',
+                  'output',
+                  'output',
+                  newNodePosition,
+                ),
+              ]);
+            }}
+          >
+            Section Output
+          </StyledOperationButton>
+        )}
+      {editorMode.mode === EditorMode.Template &&
+        namespace === ROOT_NAMESPACE && (
+          <StyledOperationButton
+            startIcon={<SectionBufferIcon />}
+            onClick={() => {
+              onAdd?.([
+                createSectionInterfaceChange(
+                  'sectionBuffer',
+                  'buffer',
+                  'buffer',
+                  newNodePosition,
+                ),
+              ]);
+            }}
+          >
+            Section Buffer
+          </StyledOperationButton>
         )}
       <StyledOperationButton
         startIcon={<NodeIcon />}
diff --git a/diagram-editor/frontend/diagram-editor.tsx b/diagram-editor/frontend/diagram-editor.tsx
@@ -1,3 +1,4 @@
+import { v4 as uuidv4 } from 'uuid';
 import {
   Alert,
   alpha,
@@ -43,7 +44,11 @@ import {
   type OperationNode,
 } from './nodes';
 import { useTemplates } from './templates-provider';
-import { exhaustiveCheck, allowEdges as getAllowEdges } from './utils';
+import {
+  exhaustiveCheck,
+  allowEdges as getAllowEdges,
+  ROOT_NAMESPACE,
+} from './utils';
 import { autoLayout } from './utils/auto-layout';
 import { calculateScopeBounds, LAYOUT_OPTIONS } from './utils/layout';
 import { loadDiagramJson, loadEmpty } from './utils/load-diagram';
@@ -92,6 +97,35 @@ function getChangeParentIdAndPosition(
   }
 }
 
+function defaultTemplateNodes(): DiagramEditorNode[] {
+  return [
+    {
+      id: uuidv4(),
+      type: 'sectionInput',
+      position: { x: 0, y: 0 },
+      data: {
+        namespace: ROOT_NAMESPACE,
+        remappedId: 'input',
+        targetId: 'input',
+      },
+      width: LAYOUT_OPTIONS.nodeWidth,
+      height: LAYOUT_OPTIONS.nodeHeight,
+    },
+    {
+      id: uuidv4(),
+      type: 'sectionOutput',
+      position: { x: 0, y: 400 },
+      data: {
+        namespace: ROOT_NAMESPACE,
+        remappedId: 'output',
+        targetId: 'output',
+      },
+      width: LAYOUT_OPTIONS.nodeWidth,
+      height: LAYOUT_OPTIONS.nodeHeight,
+    },
+  ];
+}
+
 const DiagramEditor = () => {
   const reactFlowInstance = React.useRef<ReactFlowInstance<
     DiagramEditorNode,
@@ -103,8 +137,9 @@ const DiagramEditor = () => {
   const [nodes, setNodes] = React.useState<DiagramEditorNode[]>(
     () => loadEmpty().nodes,
   );
-  const [templateNodes, _setTemplateNodes] =
-    React.useState<DiagramEditorNode[]>();
+  const [templateNodes, _setTemplateNodes] = React.useState<
+    DiagramEditorNode[]
+  >(() => defaultTemplateNodes());
   const renderedNodes = React.useMemo(() => {
     switch (editorMode.mode) {
       case EditorMode.Normal:
diff --git a/diagram-editor/frontend/nodes/section-node.tsx b/diagram-editor/frontend/nodes/section-node.tsx
@@ -13,7 +13,8 @@ import {
 export type SectionInterfaceData = {
   // section interfaces is always in the root namespace
   namespace: typeof ROOT_NAMESPACE;
-  connectKey: string;
+  remappedId: string;
+  targetId: string;
 };
 
 export function SectionNode(props: NodeProps<OperationNode<'section'>>) {
@@ -49,7 +50,7 @@ export function SectionInputNode(
       {...props}
       color="secondary"
       icon={<SectionInputIcon />}
-      label={props.data.connectKey}
+      label={props.data.remappedId}
       variant="output"
     />
   );
@@ -63,7 +64,7 @@ export function SectionOutputNode(
       {...props}
       color="secondary"
       icon={<SectionOutputIcon />}
-      label={props.data.connectKey}
+      label={props.data.remappedId}
       variant="input"
     />
   );
@@ -77,7 +78,7 @@ export function SectionBufferNode(
       {...props}
       color="secondary"
       icon={<SectionBufferIcon />}
-      label={props.data.connectKey}
+      label={props.data.remappedId}
       variant="output"
     />
   );
diff --git a/diagram-editor/frontend/utils/export-diagram.ts b/diagram-editor/frontend/utils/export-diagram.ts
@@ -1,7 +1,7 @@
 import equal from 'fast-deep-equal';
 import type { DiagramEditorEdge, StreamOutEdge } from '../edges';
 import type { NodeManager } from '../node-manager';
-import { isBuiltinNode, isOperationNode } from '../nodes';
+import { isOperationNode } from '../nodes';
 import type {
   BufferSelection,
   Diagram,
@@ -143,7 +143,7 @@ function syncBufferSelection(
  */
 function syncEdge(
   nodeManager: NodeManager,
-  diagram: Diagram,
+  root: SubOperations,
   edge: DiagramEditorEdge,
 ): void {
   if (edge.type === 'bufferKey' || edge.type === 'bufferSeq') {
@@ -156,7 +156,7 @@ function syncEdge(
   if (sourceNode.type === 'start') {
     const subOps: SubOperations = (() => {
       if (!sourceNode.parentId) {
-        return diagram;
+        return root;
       }
       const scopeNode = nodeManager.getNode(sourceNode.parentId);
       if (!isScopeNode(scopeNode)) {
@@ -285,15 +285,15 @@ function syncEdge(
 /**
  * Update the operation connections from the edges.
  *
- * @param diagram only used to update the `start` field, does not actually populate the operations.
+ * @param root only used to update the `start` field, does not actually populate the operations.
  */
 function syncEdges(
   nodeManager: NodeManager,
-  diagram: Diagram,
+  root: SubOperations,
   edges: DiagramEditorEdge[],
 ): void {
   // first clear all the connections
-  diagram.start = { builtin: 'dispose' };
+  root.start = { builtin: 'dispose' };
   for (const node of nodeManager.iterNodes()) {
     if (!isOperationNode(node)) {
       continue;
@@ -356,7 +356,7 @@ function syncEdges(
   }
 
   for (const edge of edges) {
-    syncEdge(nodeManager, diagram, edge);
+    syncEdge(nodeManager, root, edge);
   }
 }
 
@@ -395,56 +395,55 @@ export function exportDiagram(
     }
   }
 
-  for (const edge of edges) {
-    const source = nodeManager.getNode(edge.source);
-    if (source.type === 'start') {
-      const subOps = getSubOperations(diagram, source.data.namespace);
-      const target = nodeManager.getNode(edge.target);
-      if (isOperationNode(target)) {
-        subOps.start = target.data.opId;
-      } else if (isBuiltinNode(target)) {
-        subOps.start = { builtin: target.type };
-      } else {
-        throw new Error('unknown node');
-      }
-    }
-  }
-
   return diagram;
 }
 
-// export function exportTemplate(
-//   nodeManager: NodeManager,
-//   edges: DiagramEditorEdge[],
-// ): Record<string, DiagramOperation> {
-//   const root: SubOperations = {
-//     start: { builtin: 'dispose' },
-//     ops: {},
-//   };
+export function exportTemplate(
+  nodeManager: NodeManager,
+  edges: DiagramEditorEdge[],
+): SectionTemplate {
+  const template: Required<SectionTemplate> = {
+    inputs: {},
+    outputs: [],
+    buffers: {},
+    ops: {},
+  };
+  const fakeRoot: SubOperations = {
+    start: { builtin: 'dispose' },
+    ops: {},
+  };
 
-//   for (const node of nodeManager.iterNodes()) {
-//     const subOps = getSubOperations(root, node.data.namespace);
+  syncEdges(nodeManager, fakeRoot, edges);
 
-//     if (isOperationNode(node)) {
-//       subOps.ops[node.data.opId] = node.data.op;
-//     }
-//   }
+  // visit the nodes breath first to ensure that the parents exist in the diagram before
+  // populating the children.
+  const sortedNodes = Array.from(nodeManager.iterNodes()).sort(
+    (a, b) =>
+      splitNamespaces(a.data.namespace).length -
+      splitNamespaces(b.data.namespace).length,
+  );
+  for (const node of sortedNodes) {
+    const subOps = getSubOperations(fakeRoot, node.data.namespace);
 
-//   for (const edge of edges) {
-//     const source = nodeManager.getNode(edge.source);
-//     if (source.type === 'start') {
-//       const subOps = getSubOperations(root, source.data.namespace);
-//       const target = nodeManager.getNode(edge.target);
-//       if (isOperationNode(target)) {
-//         subOps.start = target.data.opId;
-//       } else if (isBuiltinNode(target)) {
-//         subOps.start = { builtin: target.type };
-//       } else {
-//         throw new Error('unknown node');
-//       }
-//     }
-//     nodeManager.syncEdge(edge);
-//   }
+    if (isOperationNode(node)) {
+      subOps.ops[node.data.opId] = { ...node.data.op };
+      if (node.data.op.type === 'scope') {
+        // do not carry over stale ops from the node data
+        subOps.ops[node.data.opId].ops = {};
+      }
+    } else if (node.type === 'sectionInput') {
+      template.inputs[node.data.remappedId] = node.data.targetId;
+    } else if (node.type === 'sectionOutput') {
+      template.outputs.push(node.data.remappedId);
+    } else if (node.type === 'sectionBuffer') {
+      template.buffers[node.data.remappedId] = node.data.targetId;
+    }
+  }
 
-//   return root.ops;
-// }
+  return {
+    inputs: {},
+    outputs: [],
+    buffers: {},
+    ops: fakeRoot.ops,
+  };
+}
