Use custom edges

Author: saulshanabrook

src/Visualizer.tsx

@@ -4,7 +4,8 @@ import "@xyflow/react/dist/style.css";
 
 import { scheme } from "vega-scale";
 import { ErrorBoundary } from "react-error-boundary";
-import ELK, { ElkNode, ElkPrimitiveEdge } from "elkjs/lib/elk.bundled.js";
+import type { EdgeProps, EdgeTypes, NodeProps } from "@xyflow/react";
+import ELK, { ElkExtendedEdge, ElkNode } from "elkjs/lib/elk.bundled.js";
 import { memo, startTransition, Suspense, use, useCallback, useEffect, useMemo, useState } from "react";
 import {
   ReactFlow,
@@ -13,13 +14,14 @@ import {
   Panel,
   NodeTypes,
   NodeMouseHandler,
-  Position,
-  Handle,
   Background,
   MarkerType,
   Edge,
   useNodesInitialized,
   useReactFlow,
+  BaseEdge,
+  Handle,
+  Position,
 } from "@xyflow/react";
 
 import "@xyflow/react/dist/style.css";
@@ -36,7 +38,7 @@ const layoutOptions = {
   "elk.direction": "DOWN",
   "elk.portConstraints": "FIXED_SIDE",
   "elk.hierarchyHandling": "INCLUDE_CHILDREN",
-  // "elk.layered.mergeEdges": "True",
+  "elk.layered.mergeEdges": "True",
   "elk.edgeRouting": "ORTHOGONAL",
   "elk.layered.nodePlacement.strategy": "NETWORK_SIMPLEX",
 
@@ -70,19 +72,26 @@ type Color = string;
 // https://vega.github.io/vega/docs/schemes/#categorical
 const colorScheme: Color[] = [...scheme("pastel1"), ...scheme("pastel2")];
 
+type FlowClass = Node<{ color: string | null; id: string }, "class">;
+type FlowNode = Node<{ label: string; id: string }, "node">;
+type FlowEdge = Edge<{ points: { x: number; y: number }[] }, "edge">;
+
+type MyELKEdge = ElkExtendedEdge & { data: { sourceNode: string } };
 /// ELK Node but with additional data added to be later used when converting to react flow nodes
-type MyELKNode = Omit<ElkNode, "children"> & {
+type MyELKNode = Omit<ElkNode, "children" | "edges"> & {
+  edges: MyELKEdge[];
   children: ({
-    type: string;
-    data: { color: string | null; port: string; id: string };
+    type: NonNullable<FlowClass["type"]>;
+    data: FlowClass["data"];
+    edges: MyELKEdge[];
     children: {
-      type: string;
+      type: NonNullable<FlowNode["type"]>;
       width: number;
       height: number;
-      data: { label: string; ports: { id: string }[]; id: string };
+      data: FlowNode["data"];
     }[] &
       ElkNode[];
-  } & Omit<ElkNode, "children" | "position">)[];
+  } & Omit<ElkNode, "children" | "position" | "edges">)[];
 };
 
 /// ELK node but with layout information added
@@ -157,87 +166,45 @@ function toELKNode(
     return {
       id: `class-${id}`,
       data: { color: type_to_color.get(egraph.class_data[id]?.type) || null, port: `port-${id}`, id },
-      ports: [
-        // {
-        //   id: `port-${id}`,
-        //   layoutOptions: {
-        //     "port.side": "NORTH",
-        //   },
-        // },
-      ],
-      type: "class",
+      type: "class" as const,
       children: nodes.map(([id, node]) => {
-        const ports = Object.keys(node.children).map((index) => ({
-          id: `port-${id}-${index}`,
-          layoutOptions: {
-            "port.side": "SOUTH",
-          },
-        }));
         // compute the size of the text by setting a dummy node element then measureing it
         innerElem.innerText = node.op;
         const size = outerElem.getBoundingClientRect();
         return {
           id: `node-${id}`,
-          type: "node",
-          data: { label: node.op, ports, id },
+          type: "node" as const,
+          data: { label: node.op, id },
           width: size.width,
           height: size.height,
-          // one port for every index
-          ports,
+          ports: Object.keys(node.children).map((index) => ({
+            id: `port-${id}-${index}`,
+            layoutOptions: {
+              "port.side": "SOUTH",
+            },
+          })),
         };
       }),
       edges: nodes.flatMap(([id, node]) =>
-        [...node.children.entries()].flatMap(([index, childNode]) => {
-          const sourcePort = `port-${id}-${index}`;
-          const class_ = nodeToClass.get(childNode)!;
-          // only include if this is a self loop
-          if (class_ != id) {
-            return [];
-          }
-          // If the target or source class is not in the selected nodes, don't draw the edge
-          const targetPort = `port-${class_}`;
-          return [
-            {
-              id: `edge-${id}-${index}`,
-              source: `node-${id}`,
-              sourcePort,
-              sourceHandle: sourcePort,
-              target: `class-${class_}`,
-              targetPort,
-              targetHandle: targetPort,
-            },
-          ];
-        })
+        [...node.children.entries()].map(([index, childNode]) => ({
+          id: `edge-${id}-${index}`,
+          data: { sourceNode: `node-${id}` },
+          sources: [`port-${id}-${index}`],
+          targets: [`class-${nodeToClass.get(childNode)!}`],
+        }))
       ),
     };
   });
 
-  const edges: ElkPrimitiveEdge[] = Object.entries(egraph.nodes).flatMap(([id, node]) =>
-    [...node.children.entries()].flatMap(([index, childNode]) => {
-      const sourcePort = `port-${id}-${index}`;
-      const class_ = nodeToClass.get(childNode)!;
-      // If the target or source class is not in the selected nodes, don't draw the edge
-      if (
-        !classToNodes.get(node.eclass)?.find(([sourceID]) => sourceID == id) ||
-        !classToNodes.has(class_) ||
-        class_ == nodeToClass.get(id)
-      ) {
-        return [];
-      }
-      const targetPort = `port-${class_}`;
-      return [
-        {
-          id: `edge-${id}-${index}`,
-          source: `node-${id}`,
-          sourcePort,
-          sourceHandle: sourcePort,
-          target: `class-${class_}`,
-          targetPort,
-          targetHandle: targetPort,
-        },
-      ];
-    })
-  );
+  // move all edges that aren't self loops to the root
+  // https://github.com/eclipse/elk/issues/1068
+  const edges = [];
+  for (const child of children) {
+    const { loop, not } = Object.groupBy(child.edges, ({ targets }) => (targets[0] === child.id ? "loop" : "not"));
+    child.edges = loop || [];
+    edges.push(...(not || []));
+  }
+
   return {
     id: "--eclipse-layout-kernel-root",
     layoutOptions,
@@ -247,7 +214,7 @@ function toELKNode(
 }
 
 // This function takes an EGraph and returns an ELK node that can be used to layout the graph.
-function toFlowNodes(layout: MyELKNodeLayedOut): Node[] {
+function toFlowNodes(layout: MyELKNodeLayedOut): (FlowClass | FlowNode)[] {
   return layout.children.flatMap(({ children, x, y, data, id: parentId, type, height, width }) => [
     { position: { x, y }, data, id: parentId, type, height, width },
     ...children!.map(({ x, y, height, width, data, id, type }) => ({
@@ -262,68 +229,81 @@ function toFlowNodes(layout: MyELKNodeLayedOut): Node[] {
   ]);
 }
 
-export function EClassNode({ data }: { data: { port: string; color: string; id: string } }) {
+function toFlowEdges(layout: MyELKNodeLayedOut): FlowEdge[] {
+  const allEdges = [...layout.edges!, ...layout.children.flatMap(({ edges }) => edges!)];
+  return allEdges.map(({ id, sections, data: { sourceNode } }) => {
+    const [section] = sections!;
+    return {
+      type: "edge",
+      id,
+      source: sourceNode,
+      target: section.outgoingShape!,
+      data: {
+        points: [section.startPoint, ...(section.bendPoints || []), section.endPoint],
+      },
+    };
+  });
+}
+
+export function EClassNode({ data }: NodeProps<FlowClass>) {
   return (
-    <div className="rounded-md border border-dotted border-stone-400 h-full w-full" style={{ backgroundColor: data.color }}>
-      <Handle className="top-0 bottom-0 opacity-0  translate-0" type="target" id={data.port} position={Position.Top} />
+    <div className="rounded-md border border-dotted border-stone-400 h-full w-full" style={{ backgroundColor: data.color! }}>
+      <Handle type="target" position={Position.Top} className="invisible" />
     </div>
   );
 }
 
-export function ENode({
-  data,
-  ...rest
-}: {
-  data: { label: string; ports: { id: string }[]; id: string };
-  outerRef?: React.Ref<HTMLDivElement>;
-  innerRef?: React.Ref<HTMLDivElement>;
-}) {
+export function ENode(
+  props: Partial<
+    NodeProps<FlowNode> & {
+      outerRef: React.Ref<HTMLDivElement>;
+      innerRef: React.Ref<HTMLDivElement>;
+    }
+  >
+) {
   return (
-    <div className="p-1 rounded-md border bg-white border-stone-400 h-full w-full" ref={rest?.outerRef}>
-      <div className="font-mono truncate max-w-96" ref={rest?.innerRef}>
-        {data.label}
-      </div>
-      {/* Place at bottom of parent so that handles just touch the node */}
-      <div className="flex justify-around absolute bottom-0 w-full">
-        {data.ports.map(({ id }) => (
-          <Handle
-            className="relative left-auto transform-none opacity-0 translate-0"
-            key={id}
-            type="source"
-            position={Position.Bottom}
-            id={id}
-          />
-        ))}
+    <div className="p-1 rounded-md border bg-white border-stone-400 h-full w-full" ref={props?.outerRef}>
+      <div className="font-mono truncate max-w-96" ref={props?.innerRef}>
+        {props?.data?.label}
       </div>
+      {/* Only show handle if we aren't rendering this to calculate size */}
+      {props?.outerRef ? <></> : <Handle type="source" position={Position.Bottom} className="invisible" />}
     </div>
   );
 }
 
+export function CustomEdge({ markerEnd, data }: EdgeProps<FlowEdge>) {
+  const { points } = data!;
+  const edgePath = points.map(({ x, y }, index) => `${index === 0 ? "M" : "L"} ${x} ${y}`).join(" ");
+  return <BaseEdge path={edgePath} markerEnd={markerEnd} />;
+}
+
 const nodeTypes: NodeTypes = {
   class: memo(EClassNode),
   node: memo(ENode),
 };
 
+const edgeTypes: EdgeTypes = {
+  edge: memo(CustomEdge),
+};
+
 function LayoutFlow({ egraph, outerElem, innerElem }: { egraph: string; outerElem: HTMLDivElement; innerElem: HTMLDivElement }) {
   const parsedEGraph: EGraph = useMemo(() => JSON.parse(egraph), [egraph]);
   /// e-class ID we have currently selected
   const [selectedNode, setSelectedNode] = useState<{ type: "class" | "node"; id: string } | null>(null);
-  const elkNode = useMemo(() => {
-    const r = toELKNode(parsedEGraph, outerElem, innerElem, selectedNode);
-    // console.log(JSON.parse(JSON.stringify(r, null, 2)));
-    return r;
-  }, [parsedEGraph, outerElem, innerElem, selectedNode]);
-  const edges = useMemo(
-    () => [...elkNode.children.flatMap((c) => c.edges!.map((e) => ({ ...e }))), ...elkNode.edges!.map((e) => ({ ...e }))],
-    [elkNode]
+  const elkNode = useMemo(
+    () => toELKNode(parsedEGraph, outerElem, innerElem, selectedNode),
+    [parsedEGraph, outerElem, innerElem, selectedNode]
   );
   const layoutPromise = useMemo(() => elk.layout(elkNode) as Promise<MyELKNodeLayedOut>, [elkNode]);
   const layout = use(layoutPromise);
+  const edges = useMemo(() => toFlowEdges(layout), [layout]);
   const nodes = useMemo(() => toFlowNodes(layout), [layout]);
   const onNodeClick = useCallback(
     ((_, node) => {
-      startTransition(() => setSelectedNode({ type: node.type! as "class" | "node", id: node.data!.id! as string }));
-    }) as NodeMouseHandler,
+      // Use start transition so that the whole component doesn't re-render
+      startTransition(() => setSelectedNode({ type: node.type!, id: node.data.id }));
+    }) as NodeMouseHandler<FlowClass | FlowNode>,
     [setSelectedNode]
   );
 
@@ -340,7 +320,8 @@ function LayoutFlow({ egraph, outerElem, innerElem }: { egraph: string; outerEle
     <ReactFlow
       nodes={nodes}
       nodeTypes={nodeTypes}
-      edges={edges as unknown as Edge[]}
+      edgeTypes={edgeTypes}
+      edges={edges}
       minZoom={0.05}
       maxZoom={10}
       defaultEdgeOptions={{ markerEnd: { type: MarkerType.ArrowClosed } }}
@@ -360,7 +341,7 @@ function Visualizer({ egraph }: { egraph: string }) {
     <>
       {/* Hidden node to measure text size  */}
       <div className="invisible absolute">
-        <ENode data={{ label: "test", ports: [] }} outerRef={setOuterElem} innerRef={setInnerElem} />
+        <ENode outerRef={setOuterElem} innerRef={setInnerElem} />
       </div>
       <ReactFlowProvider>
         <ErrorBoundary fallback={<p>⚠️Something went wrong</p>}>

src/anywidget.tsx

@@ -1,4 +1,4 @@
-import { StrictMode, useSyncExternalStore } from "react";
+import { useSyncExternalStore } from "react";
 import { createRoot } from "react-dom/client";
 import "./index.css";
 import { DOMWidgetModel } from "@jupyter-widgets/base";
@@ -18,11 +18,7 @@ function ModelApp({ model }: { model: DOMWidgetModel }) {
 
 function render({ model, el }: { el: HTMLElement; model: DOMWidgetModel }) {
   const root = createRoot(el);
-  root.render(
-    <StrictMode>
-      <ModelApp model={model} />
-    </StrictMode>
-  );
+  root.render(<ModelApp model={model} />);
   return () => root.unmount();
 }
 

tsconfig.app.json

@@ -2,23 +2,28 @@
   "compilerOptions": {
     "target": "ES2020",
     "useDefineForClassFields": true,
-    "lib": ["ES2020", "DOM", "DOM.Iterable"],
+    "lib": [
+      "ES2020",
+      "DOM",
+      "DOM.Iterable",
+      "ESNext"
+    ],
     "module": "ESNext",
     "skipLibCheck": true,
-
     /* Bundler mode */
     "moduleResolution": "bundler",
     "allowImportingTsExtensions": true,
     "isolatedModules": true,
     "moduleDetection": "force",
     "noEmit": true,
     "jsx": "react-jsx",
-
     /* Linting */
     "strict": true,
     "noUnusedLocals": true,
     "noUnusedParameters": true,
     "noFallthroughCasesInSwitch": true
   },
-  "include": ["src"]
+  "include": [
+    "src"
+  ]
 }