Use different colors per node type

egraphs-good · Aug 19, 2024 · d8985a2 · d8985a2
1 parent d1953d1
commit d8985a2
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 18 deletions.
diff --git a/package.json b/package.json
@@ -21,7 +21,8 @@
     "react": "19.0.0-rc-fa6eab58-20240815",
     "react-dom": "19.0.0-rc-fa6eab58-20240815",
     "react-error-boundary": "4.0.13",
-    "react-monaco-editor": "0.56.1"
+    "react-monaco-editor": "0.56.1",
+    "vega-scale": "7.4.1"
   },
   "devDependencies": {
     "@eslint/js": "^9.9.0",

diff --git a/src/Visualizer.tsx b/src/Visualizer.tsx
@@ -2,10 +2,11 @@
 
 import "@xyflow/react/dist/style.css";
 
+import { scheme } from "vega-scale";
 import { ErrorBoundary } from "react-error-boundary";
 import ELK, { ElkExtendedEdge, ElkNode, ElkPrimitiveEdge } from "elkjs/lib/elk.bundled.js";
 import { memo, Suspense, use, useMemo, useRef, useState } from "react";
-import { ReactFlow, ReactFlowProvider, Node, Panel, Edge, NodeTypes, Position, Handle, Background } from "@xyflow/react";
+import { ReactFlow, ReactFlowProvider, Node, Panel, Edge, NodeTypes, Position, Handle, Background, MarkerType } from "@xyflow/react";
 
 import "@xyflow/react/dist/style.css";
 
@@ -46,15 +47,10 @@ type EGraph = {
   class_data: { [id: EGraphClassID]: EGraphClassData };
 };
 
-// https://github.com/vega/editor/blob/8bcc0fc997dbf3bdcecf4584ffd742fc33ddb042/src/features/dataflow/utils/measureText.ts
-
-const ctx = document.createElement("canvas").getContext("2d");
-const fontFamily = "monospace";
-const fontSize = "12px";
-
-ctx!.font = `${fontSize} ${fontFamily}`;
-
-export type Size = { width: number; height: number };
+type Color = string;
+// Use these color schemes for the nodes
+// https://vega.github.io/vega/docs/schemes/#categorical
+const colorScheme: Color[] = [...scheme("pastel1"), ...scheme("pastel2")];
 
 // We wil convert this to a graph where the id of the nodes are class-{class_id} and node-{node_id}
 // the ID of the edges will be edge-{source_id}-{port-index} and the ports will be port-{source_id}-{port-index}
@@ -69,11 +65,18 @@ function toELKNode(egraph: EGraph, outerElem: HTMLDivElement, innerElem: HTMLDiv
     classToNodes.get(node.eclass)!.push([id, node]);
   }
 
+  const type_to_color = new Map<string | undefined, Color>();
+  for (const { type } of Object.values(egraph.class_data)) {
+    if (!type_to_color.has(type)) {
+      type_to_color.set(type, colorScheme[type_to_color.size % colorScheme.length]);
+    }
+  }
+
   const children = [...classToNodes.entries()].map(([id, nodes]) => {
     const parentID = `class-${id}`;
     return {
       id: parentID,
-      data: { type: egraph.class_data[id].type || null, port: `port-${id}` },
+      data: { color: type_to_color.get(egraph.class_data[id]?.type) || null, port: `port-${id}` },
       type: "class",
       children: nodes.map(([id, node]) => {
         const ports = Object.keys(node.children).map((index) => ({
@@ -138,10 +141,10 @@ function toFlowNodes(layout: ElkNode): Node[] {
   ]);
 }
 
-export function EClassNode({ data }: { data: { port: string; type: string | null } }) {
+export function EClassNode({ data }: { data: { port: string; color: string } }) {
   return (
-    <div className="rounded-md border border-dotted border-stone-400 h-full w-full">
-      <Handle type="target" id={data.port} position={Position.Top} className="invisible" />
+    <div className="rounded-md border border-dotted border-stone-400 h-full w-full" style={{ backgroundColor: data.color }}>
+      <Handle type="target" id={data.port} position={Position.Top} />
     </div>
   );
 }
@@ -160,7 +163,7 @@ export function ENode({
         {data.label}
       </div>
       {data.ports.map(({ id }) => (
-        <Handle key={id} type="source" position={Position.Bottom} id={id} style={{ top: 10, background: "#555" }} />
+        <Handle key={id} type="source" position={Position.Bottom} id={id} />
       ))}
     </div>
   );
@@ -177,9 +180,17 @@ function LayoutFlow({ egraph, outerElem, innerElem }: { egraph: string; outerEle
   const edges = useMemo(() => elkNode.edges!.map((e) => ({ ...e })), [elkNode]);
   const layoutPromise = useMemo(() => elk.layout(elkNode), [elkNode]);
   const layout = use(layoutPromise);
+  console.log(layout);
   const nodes = useMemo(() => toFlowNodes(layout), [layout]);
   return (
-    <ReactFlow nodes={nodes} nodeTypes={nodeTypes} edges={edges} fitView minZoom={0.05}>
+    <ReactFlow
+      nodes={nodes}
+      nodeTypes={nodeTypes}
+      edges={edges}
+      fitView
+      minZoom={0.05}
+      defaultEdgeOptions={{ type: "straight", markerEnd: { type: MarkerType.ArrowClosed } }}
+    >
       <Background />
     </ReactFlow>
   );

diff --git a/src/index.css b/src/index.css
@@ -1,3 +1,6 @@
 @tailwind base;
 @tailwind components;
 @tailwind utilities;
+.react-flow__node {
+  z-index: -1 !important;
+}
diff --git a/yarn.lock b/yarn.lock
@@ -1340,6 +1340,13 @@ csstype@^3.0.2:
   resolved "https://registry.yarnpkg.com/csstype/-/csstype-3.1.3.tgz#d80ff294d114fb0e6ac500fbf85b60137d7eff81"
   integrity sha512-M1uQkMl8rQK/szD0LNhtqxIPLpimGm8sOBwU7lLnCpSbTyY3yeU1Vc7l4KT5zT4s/yOxHH5O7tIuuLOCnLADRw==
 
+"d3-array@2 - 3", "d3-array@2.10.0 - 3", d3-array@^3.2.2:
+  version "3.2.4"
+  resolved "https://registry.yarnpkg.com/d3-array/-/d3-array-3.2.4.tgz#15fec33b237f97ac5d7c986dc77da273a8ed0bb5"
+  integrity sha512-tdQAmyA18i4J7wprpYq8ClcxZy3SC31QMeByyCFyRt7BVHdREQZ5lpzoe5mFEYZUWe+oq8HBvk9JjpibyEV4Jg==
+  dependencies:
+    internmap "1 - 2"
+
 "d3-color@1 - 3":
   version "3.1.0"
   resolved "https://registry.yarnpkg.com/d3-color/-/d3-color-3.1.0.tgz#395b2833dfac71507f12ac2f7af23bf819de24e2"
@@ -1363,18 +1370,56 @@ csstype@^3.0.2:
   resolved "https://registry.yarnpkg.com/d3-ease/-/d3-ease-3.0.1.tgz#9658ac38a2140d59d346160f1f6c30fda0bd12f4"
   integrity sha512-wR/XK3D3XcLIZwpbvQwQ5fK+8Ykds1ip7A2Txe0yxncXSdq1L9skcG7blcedkOX+ZcgxGAmLX1FrRGbADwzi0w==
 
-"d3-interpolate@1 - 3":
+"d3-format@1 - 3":
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/d3-format/-/d3-format-3.1.0.tgz#9260e23a28ea5cb109e93b21a06e24e2ebd55641"
+  integrity sha512-YyUI6AEuY/Wpt8KWLgZHsIU86atmikuoOmCfommt0LYHiQSPjvX2AcFc38PX0CBpr2RCyZhjex+NS/LPOv6YqA==
+
+"d3-interpolate@1 - 3", "d3-interpolate@1.2.0 - 3", d3-interpolate@^3.0.1:
   version "3.0.1"
   resolved "https://registry.yarnpkg.com/d3-interpolate/-/d3-interpolate-3.0.1.tgz#3c47aa5b32c5b3dfb56ef3fd4342078a632b400d"
   integrity sha512-3bYs1rOD33uo8aqJfKP3JWPAibgw8Zm2+L9vBKEHJ2Rg+viTR7o5Mmv5mZcieN+FRYaAOWX5SJATX6k1PWz72g==
   dependencies:
     d3-color "1 - 3"
 
+d3-scale-chromatic@^3.1.0:
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/d3-scale-chromatic/-/d3-scale-chromatic-3.1.0.tgz#34c39da298b23c20e02f1a4b239bd0f22e7f1314"
+  integrity sha512-A3s5PWiZ9YCXFye1o246KoscMWqf8BsD9eRiJ3He7C9OBaxKhAd5TFCdEx/7VbKtxxTsu//1mMJFrEt572cEyQ==
+  dependencies:
+    d3-color "1 - 3"
+    d3-interpolate "1 - 3"
+
+d3-scale@^4.0.2:
+  version "4.0.2"
+  resolved "https://registry.yarnpkg.com/d3-scale/-/d3-scale-4.0.2.tgz#82b38e8e8ff7080764f8dcec77bd4be393689396"
+  integrity sha512-GZW464g1SH7ag3Y7hXjf8RoUuAFIqklOAq3MRl4OaWabTFJY9PN/E1YklhXLh+OQ3fM9yS2nOkCoS+WLZ6kvxQ==
+  dependencies:
+    d3-array "2.10.0 - 3"
+    d3-format "1 - 3"
+    d3-interpolate "1.2.0 - 3"
+    d3-time "2.1.1 - 3"
+    d3-time-format "2 - 4"
+
 "d3-selection@2 - 3", d3-selection@3, d3-selection@^3.0.0:
   version "3.0.0"
   resolved "https://registry.yarnpkg.com/d3-selection/-/d3-selection-3.0.0.tgz#c25338207efa72cc5b9bd1458a1a41901f1e1b31"
   integrity sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==
 
+"d3-time-format@2 - 4":
+  version "4.1.0"
+  resolved "https://registry.yarnpkg.com/d3-time-format/-/d3-time-format-4.1.0.tgz#7ab5257a5041d11ecb4fe70a5c7d16a195bb408a"
+  integrity sha512-dJxPBlzC7NugB2PDLwo9Q8JiTR3M3e4/XANkreKSUxF8vvXKqm1Yfq4Q5dl8budlunRVlUUaDUgFt7eA8D6NLg==
+  dependencies:
+    d3-time "1 - 3"
+
+"d3-time@1 - 3", "d3-time@2.1.1 - 3", d3-time@^3.1.0:
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/d3-time/-/d3-time-3.1.0.tgz#9310db56e992e3c0175e1ef385e545e48a9bb5c7"
+  integrity sha512-VqKjzBLejbSMT4IgbmVgDjpkYrNWUYJnbCGo874u7MMKIWsILRX+OpX/gTk8MqjpT1A/c6HY2dCA77ZN0lkQ2Q==
+  dependencies:
+    d3-array "2 - 3"
+
 "d3-timer@1 - 3":
   version "3.0.1"
   resolved "https://registry.yarnpkg.com/d3-timer/-/d3-timer-3.0.1.tgz#6284d2a2708285b1abb7e201eda4380af35e63b0"
@@ -1790,6 +1835,11 @@ imurmurhash@^0.1.4:
   resolved "https://registry.yarnpkg.com/imurmurhash/-/imurmurhash-0.1.4.tgz#9218b9b2b928a238b13dc4fb6b6d576f231453ea"
   integrity sha512-JmXMZ6wuvDmLiHEml9ykzqO6lwFbof0GG4IkcGaENdCRDDmMVnny7s5HsIgHCbaq0w2MyPhDqkhTUgS2LU2PHA==
 
+"internmap@1 - 2":
+  version "2.0.3"
+  resolved "https://registry.yarnpkg.com/internmap/-/internmap-2.0.3.tgz#6685f23755e43c524e251d29cbc97248e3061009"
+  integrity sha512-5Hh7Y1wQbvY5ooGgPbDaL5iYLAPzMTUrjMulskHLH6wnv/A+1q5rgEaiuqEjB+oxGXIVZs1FF+R/KPN3ZSQYYg==
+
 is-binary-path@~2.1.0:
   version "2.1.0"
   resolved "https://registry.yarnpkg.com/is-binary-path/-/is-binary-path-2.1.0.tgz#ea1f7f3b80f064236e83470f86c09c254fb45b09"
@@ -2620,6 +2670,32 @@ util-deprecate@^1.0.2:
   resolved "https://registry.yarnpkg.com/util-deprecate/-/util-deprecate-1.0.2.tgz#450d4dc9fa70de732762fbd2d4a28981419a0ccf"
   integrity sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==
 
+vega-scale@7.4.1:
+  version "7.4.1"
+  resolved "https://registry.yarnpkg.com/vega-scale/-/vega-scale-7.4.1.tgz#2dcd3e39ebb00269b03a8be86e44c7b48c67442a"
+  integrity sha512-dArA28DbV/M92O2QvswnzCmQ4bq9WwLKUoyhqFYWCltmDwkmvX7yhqiFLFMWPItIm7mi4Qyoygby6r4DKd1X2A==
+  dependencies:
+    d3-array "^3.2.2"
+    d3-interpolate "^3.0.1"
+    d3-scale "^4.0.2"
+    d3-scale-chromatic "^3.1.0"
+    vega-time "^2.1.2"
+    vega-util "^1.17.2"
+
+vega-time@^2.1.2:
+  version "2.1.2"
+  resolved "https://registry.yarnpkg.com/vega-time/-/vega-time-2.1.2.tgz#0c414e74780613d6d3234fb97f19b50c0ebd9f49"
+  integrity sha512-6rXc6JdDt8MnCRy6UzUCsa6EeFycPDmvioMddLfKw38OYCV8pRQC5nw44gyddOwXgUTJLiCtn/sp53P0iA542A==
+  dependencies:
+    d3-array "^3.2.2"
+    d3-time "^3.1.0"
+    vega-util "^1.17.2"
+
+vega-util@^1.17.2:
+  version "1.17.2"
+  resolved "https://registry.yarnpkg.com/vega-util/-/vega-util-1.17.2.tgz#f69aa09fd5d6110c19c4a0f0af9e35945b99987d"
+  integrity sha512-omNmGiZBdjm/jnHjZlywyYqafscDdHaELHx1q96n5UOz/FlO9JO99P4B3jZg391EFG8dqhWjQilSf2JH6F1mIw==
+
 vite@^5.4.1:
   version "5.4.1"
   resolved "https://registry.yarnpkg.com/vite/-/vite-5.4.1.tgz#2aa72370de824d23f53658affd807e4c9905b058"