more checks when connecting edges

Signed-off-by: Teo Koon Peng <[email protected]>

Author: koonpeng

diagram-editor/frontend/diagram-editor.tsx

@@ -30,7 +30,7 @@ import React from 'react';
 import AddOperation from './add-operation';
 import CommandPanel from './command-panel';
 import type { DiagramEditorEdge } from './edges';
-import { EDGE_TYPES } from './edges';
+import { createBaseEdge, EDGE_TYPES } from './edges';
 import {
   EditorMode,
   type EditorModeContext,
@@ -49,7 +49,7 @@ import {
 } from './nodes';
 import { useTemplates } from './templates-provider';
 import { autoLayout } from './utils/auto-layout';
-import { getValidEdgeTypes } from './utils/connection';
+import { checkValidEdgeSimple, getValidEdgeTypes } from './utils/connection';
 import { exhaustiveCheck } from './utils/exhaustive-check';
 import { exportTemplate } from './utils/export-diagram';
 import { calculateScopeBounds, LAYOUT_OPTIONS } from './utils/layout';
@@ -527,6 +527,10 @@ function DiagramEditor() {
           closeAllPopovers();
         }}
         onConnect={(conn) => {
+          if (!reactFlowInstance.current) {
+            return;
+          }
+
           const sourceNode = reactFlowInstance.current?.getNode(conn.source);
           const targetNode = reactFlowInstance.current?.getNode(conn.target);
           if (!sourceNode || !targetNode) {
@@ -541,16 +545,22 @@ function DiagramEditor() {
             return;
           }
 
-          setEdges((prev) =>
-            addEdge(
-              {
-                ...conn,
-                type: validEdges[0],
-                data: defaultEdgeData(validEdges[0]),
-              },
-              prev,
-            ),
+          const newEdge = {
+            ...createBaseEdge(conn.source, conn.target),
+            type: validEdges[0],
+            data: defaultEdgeData(validEdges[0]),
+          } as DiagramEditorEdge;
+
+          const validationResult = checkValidEdgeSimple(
+            newEdge,
+            reactFlowInstance.current,
           );
+          if (!validationResult.valid) {
+            showErrorToast(validationResult.error);
+            return;
+          }
+
+          setEdges((prev) => addEdge(newEdge, prev));
         }}
         isValidConnection={(conn) => {
           const sourceNode = reactFlowInstance.current?.getNode(conn.source);

diagram-editor/frontend/edges/create-edge.ts

@@ -2,7 +2,7 @@ import { MarkerType } from '@xyflow/react';
 import { v4 as uuidv4 } from 'uuid';
 import type { DiagramEditorEdge, EdgeData } from '.';
 
-function createBaseEdge(
+export function createBaseEdge(
   source: string,
   target: string,
 ): Pick<DiagramEditorEdge, 'id' | 'source' | 'target' | 'markerEnd'> {

diagram-editor/frontend/forms/edit-edge-form.tsx

@@ -46,9 +46,9 @@ const EDGE_DEFAULT_DATA = {
   splitRemaining: {},
   streamOut: { name: '' },
   unzip: { seq: 0 },
-} satisfies { [k in EdgeTypes]: EdgeData[k] };
+} satisfies Record<EdgeTypes, EdgeData<EdgeTypes>>;
 
-export function defaultEdgeData(type: EdgeTypes): EdgeData[EdgeTypes] {
+export function defaultEdgeData(type: EdgeTypes): EdgeData<EdgeTypes> {
   return { ...EDGE_DEFAULT_DATA[type] };
 }
 

diagram-editor/frontend/utils/connection.ts

@@ -1,5 +1,7 @@
-import type { EdgeTypes } from '../edges';
+import type { ReactFlowInstance } from '@xyflow/react';
+import type { DiagramEditorEdge, EdgeTypes } from '../edges';
 import type { DiagramEditorNode, NodeTypes } from '../nodes';
+import { exhaustiveCheck } from './exhaustive-check';
 
 const ALLOWED_OUTPUT_EDGES: Record<NodeTypes, Set<EdgeTypes>> = {
   buffer: new Set<EdgeTypes>(['bufferKey', 'bufferSeq']),
@@ -67,3 +69,147 @@ export function getValidEdgeTypes(
     : new Set([]);
   return Array.from(setIntersection(allowedOutputEdges, allowedInputEdges));
 }
+
+enum CardinalityType {
+  Single,
+  Pair,
+  Many,
+}
+
+function getOutputCardinality(type: NodeTypes): CardinalityType {
+  switch (type) {
+    case 'fork_clone':
+    case 'unzip':
+    case 'buffer':
+    case 'section':
+    case 'split': {
+      return CardinalityType.Many;
+    }
+    case 'fork_result': {
+      return CardinalityType.Pair;
+    }
+    case 'node':
+    case 'buffer_access':
+    case 'join':
+    case 'serialized_join':
+    case 'listen':
+    case 'scope':
+    case 'stream_out':
+    case 'transform':
+    case 'start':
+    case 'terminate':
+    case 'sectionBuffer':
+    case 'sectionInput':
+    case 'sectionOutput': {
+      return CardinalityType.Single;
+    }
+    default: {
+      exhaustiveCheck(type);
+      throw new Error('unknown op type');
+    }
+  }
+}
+
+export type EdgeValidationResult =
+  | { valid: true; validEdgeTypes: EdgeTypes[] }
+  | { valid: false; error: string };
+
+function createValidationError(error: string): EdgeValidationResult {
+  return { valid: false, error };
+}
+
+/**
+ * Perform a quick check if an edge is valid.
+ * This only checks if the edge type is valid, does not check for conflicting edges, data correctness etc.
+ * Complexity is O(1).
+ */
+export function checkValidEdgeQuick(
+  edge: DiagramEditorEdge,
+  reactFlow: ReactFlowInstance<DiagramEditorNode, DiagramEditorEdge>,
+): EdgeValidationResult {
+  const sourceNode = reactFlow.getNode(edge.source);
+  const targetNode = reactFlow.getNode(edge.target);
+
+  if (!sourceNode || !targetNode) {
+    return createValidationError('cannot find source or target node');
+  }
+
+  const validEdgeTypes = getValidEdgeTypes(sourceNode, targetNode);
+  if (!validEdgeTypes.includes(edge.type)) {
+    return createValidationError('invalid edge type');
+  }
+
+  return { valid: true, validEdgeTypes };
+}
+
+/**
+ * Perform a simple check of the validity of edges.
+ * A more complete check than `checkValidEdgeQuick`, but still does not do complicated checks like type compatibility.
+ * Complexity is O(numOfEdges), so it is not recommended to call this very frequently.
+ */
+export function checkValidEdgeSimple(
+  edge: DiagramEditorEdge,
+  reactFlow: ReactFlowInstance<DiagramEditorNode, DiagramEditorEdge>,
+): EdgeValidationResult {
+  const quickCheck = checkValidEdgeQuick(edge, reactFlow);
+  if (!quickCheck.valid) {
+    return quickCheck;
+  }
+
+  const sourceNode = reactFlow.getNode(edge.source);
+  if (!sourceNode) {
+    return createValidationError('cannot find source or target node');
+  }
+
+  // Check if the source supports emitting multiple outputs.
+  // NOTE: All nodes supports "Many" inputs so we don't need to check that.
+  const outputCardinality = getOutputCardinality(sourceNode.type);
+  switch (outputCardinality) {
+    case CardinalityType.Single: {
+      if (reactFlow.getEdges().some((edge) => edge.source === sourceNode.id)) {
+        return createValidationError('source node already has an edge');
+      }
+      break;
+    }
+    case CardinalityType.Pair: {
+      let count = 0;
+      for (const edge of reactFlow.getEdges()) {
+        if (edge.source === sourceNode.id) {
+          count++;
+        }
+        if (count > 1) {
+          return createValidationError('source node already has two edges');
+        }
+      }
+      break;
+    }
+    case CardinalityType.Many: {
+      break;
+    }
+    default: {
+      exhaustiveCheck(outputCardinality);
+      throw new Error('unknown output cardinality');
+    }
+  }
+
+  return { valid: true, validEdgeTypes: quickCheck.validEdgeTypes };
+}
+
+/**
+ * Perform a full check of the validity of edges.
+ * This can be slow so it is not recommended to call this frequently.
+ */
+export async function checkValidEdgeFull(
+  edge: DiagramEditorEdge,
+  reactFlow: ReactFlowInstance<DiagramEditorNode, DiagramEditorEdge>,
+): Promise<EdgeValidationResult> {
+  const simpleCheck = checkValidEdgeSimple(edge, reactFlow);
+  if (!simpleCheck.valid) {
+    return simpleCheck;
+  }
+
+  // TODO: check message type compatibility. Writing the same logic as `bevy_impulse` is hard, it might
+  // be better to introduce a validation endpoint.
+
+  return { valid: true, validEdgeTypes: simpleCheck.validEdgeTypes };
+}