color code buffer handles; join node no longer allows data edges

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

Author: koonpeng

diagram-editor/frontend/app.css

@@ -23,8 +23,15 @@
   top: -100%;
 }
 .react-flow__handle.connectingto {
-  background: var(--mui-palette-error-main);
+  background-color: var(--mui-palette-error-main);
 }
 .react-flow__handle.valid {
-  background: var(--mui-palette-success-main);
+  background-color: var(--mui-palette-success-main);
 }
+.react-flow__handle.handle-buffer {
+  background-color: var(--mui-palette-secondary-main);
+}
+/* Some operations can output either Data or Stream so we cannot color code Stream handles.  */
+/* .react-flow__handle.handle-stream {
+  background-color: var(--mui-palette-warning-main);
+} */

diagram-editor/frontend/nodes/base-node.tsx

@@ -1,20 +1,32 @@
 import { Box, Button, type ButtonProps, Paper } from '@mui/material';
 import { Handle, type NodeProps, Position } from '@xyflow/react';
-import React from 'react';
+import { memo, useCallback } from 'react';
+import { EdgeCategory } from '../edges';
+import { exhaustiveCheck } from '../utils/exhaustive-check';
 import { LAYOUT_OPTIONS } from '../utils/layout';
 
 export interface BaseNodeProps extends NodeProps {
   color?: ButtonProps['color'];
   icon?: React.JSX.Element | string;
   label: string;
   variant: 'input' | 'output' | 'inputOutput';
+  /**
+   * defaults to `EdgeCateogry.Data`.
+   */
+  inputHandleType?: EdgeCategory;
+  /**
+   * defaults to `EdgeCateogry.Data`.
+   */
+  outputHandleType?: EdgeCategory;
 }
 
 function BaseNode({
   color,
   icon: materialIconOrSymbol,
   label,
   variant,
+  inputHandleType = EdgeCategory.Data,
+  outputHandleType = EdgeCategory.Data,
   isConnectable,
   selected,
   sourcePosition = Position.Bottom,
@@ -26,13 +38,38 @@ function BaseNode({
     ) : (
       materialIconOrSymbol
     );
+
+  const handleClassName = useCallback((handleType?: EdgeCategory) => {
+    if (handleType === undefined) {
+      return undefined;
+    }
+
+    switch (handleType) {
+      case EdgeCategory.Data: {
+        // use the default style
+        return undefined;
+      }
+      case EdgeCategory.Buffer: {
+        return 'handle-buffer';
+      }
+      case EdgeCategory.Stream: {
+        return undefined;
+      }
+      default: {
+        exhaustiveCheck(handleType);
+        throw new Error('unknown edge category');
+      }
+    }
+  }, []);
+
   return (
     <Paper>
       {(variant === 'input' || variant === 'inputOutput') && (
         <Handle
           type="target"
           position={targetPosition}
           isConnectable={isConnectable}
+          className={handleClassName(inputHandleType)}
         />
       )}
       <Button
@@ -64,10 +101,11 @@ function BaseNode({
           type="source"
           position={sourcePosition}
           isConnectable={isConnectable}
+          className={handleClassName(outputHandleType)}
         />
       )}
     </Paper>
   );
 }
 
-export default React.memo(BaseNode);
+export default memo(BaseNode);

diagram-editor/frontend/nodes/buffer-node.tsx

@@ -1,4 +1,5 @@
 import type { NodeProps } from '@xyflow/react';
+import { EdgeCategory } from '../edges';
 import type { OperationNode } from '.';
 import BaseNode from './base-node';
 import { BufferIcon } from './icons';
@@ -10,6 +11,7 @@ function BufferNodeComp(props: NodeProps<OperationNode<'buffer'>>) {
       icon={<BufferIcon />}
       label="Buffer"
       variant="inputOutput"
+      outputHandleType={EdgeCategory.Buffer}
     />
   );
 }

diagram-editor/frontend/nodes/join-node.tsx

@@ -1,4 +1,5 @@
 import type { NodeProps } from '@xyflow/react';
+import { EdgeCategory } from '../edges';
 import type { OperationNode } from '.';
 import BaseNode from './base-node';
 import { JoinIcon } from './icons';
@@ -10,6 +11,7 @@ function JoinNodeComp(props: NodeProps<OperationNode<'join'>>) {
       icon={<JoinIcon />}
       label="Join"
       variant="inputOutput"
+      inputHandleType={EdgeCategory.Buffer}
     />
   );
 }

diagram-editor/frontend/utils/connection.test.ts

@@ -157,7 +157,7 @@ describe('validate edges', () => {
     }
   });
 
-  test('"join" node accepts both data and buffer edges', () => {
+  test('"join" node only accepts buffer edges', () => {
     const nodeNode = createOperationNode(
       ROOT_NAMESPACE,
       undefined,
@@ -190,8 +190,7 @@ describe('validate edges', () => {
     for (const targetNode of [joinNode, serializedJoinNode]) {
       {
         const validEdges = getValidEdgeTypes(nodeNode, targetNode);
-        expect(validEdges.length).toBe(1);
-        expect(validEdges).toContain('default');
+        expect(validEdges.length).toBe(0);
       }
       {
         const validEdges = getValidEdgeTypes(bufferNode, targetNode);

diagram-editor/frontend/utils/connection.ts

@@ -35,15 +35,15 @@ const ALLOWED_INPUT_EDGE_CATEGORIES: Record<NodeTypes, EdgeCategory[]> = {
   buffer_access: [EdgeCategory.Data, EdgeCategory.Buffer],
   fork_clone: [EdgeCategory.Data],
   fork_result: [EdgeCategory.Data],
-  join: [EdgeCategory.Buffer, EdgeCategory.Data],
+  join: [EdgeCategory.Buffer],
   listen: [EdgeCategory.Buffer],
   node: [EdgeCategory.Data],
   scope: [EdgeCategory.Data],
   section: [EdgeCategory.Data, EdgeCategory.Buffer],
   sectionInput: [],
   sectionOutput: [EdgeCategory.Data],
   sectionBuffer: [],
-  serialized_join: [EdgeCategory.Data, EdgeCategory.Buffer],
+  serialized_join: [EdgeCategory.Buffer],
   split: [EdgeCategory.Data],
   start: [],
   stream_out: [EdgeCategory.Stream],