A bunch of improvements

crates/compiler/checkmate/www/src/components/Graph/VariableNode.tsx

@@ -0,0 +1,197 @@
+import clsx from "clsx";
+import { Handle, Position } from "reactflow";
+import { Variable } from "../../schema";
+import { assertExhaustive } from "../../utils/exhaustive";
+import { contentStyles } from "../Content";
+import { VariableElHelp } from "../Common/Variable";
+import { SubsSnapshot, TypeDescriptor } from "../../engine/subs";
+
+type AddSubVariableLink = (from: Variable, subVariable: Variable) => void;
+
+export interface VariableNodeProps {
+  data: {
+    subs: SubsSnapshot;
+    variable: Variable;
+    addSubVariableLink: AddSubVariableLink;
+  };
+}
+
+export default function VariableNode({ data }: VariableNodeProps): JSX.Element {
+  const { variable, subs, addSubVariableLink } = data;
+
+  const desc = subs.get_root(variable);
+  const styles = contentStyles(desc);
+  const basis: BasisProps = {
+    subs,
+    origin: variable,
+    addSubVariableLink,
+  };
+
+  const content = Object.entries(
+    VariableNodeContent(variable, desc, basis)
+  ).filter((el): el is [string, JSX.Element] => !!el[1]);
+
+  let expandedContent = <></>;
+  if (content.length > 0) {
+    expandedContent = (
+      <ul className="text-sm text-left mt-2 space-y-1">
+        {content.map(([key, value], i) => (
+          <li key={i} className="space-x-2">
+            {key}: {value}
+          </li>
+        ))}
+      </ul>
+    );
+  }
+
+  return (
+    <div
+      className={clsx(
+        styles.bg,
+        "bg-opacity-50 py-2 px-4 rounded-lg border",
+        "text-center font-mono"
+      )}
+    >
+      <Handle type="target" position={Position.Top} isConnectable={false} />
+      <div>
+        <VariableElHelp variable={variable} styles={styles} />
+      </div>
+      {expandedContent}
+      <Handle type="source" position={Position.Bottom} isConnectable={false} />
+    </div>
+  );
+}
+
+function VariableNodeContent(
+  variable: Variable,
+  desc: TypeDescriptor | undefined,
+  basis: BasisProps
+): Record<string, JSX.Element | null> {
+  if (!desc) return {};
+  const { content } = desc;
+
+  switch (content.type) {
+    case "Flex":
+    case "Rigid": {
+      const { name } = content;
+      return { name: name ? <>{name}</> : null };
+    }
+    case "FlexAble":
+    case "RigidAble": {
+      const { name, abilities } = content;
+      return {
+        name: <>{name}</>,
+        abilities: <>[{abilities.join(", ")}]</>,
+      };
+    }
+    case "Recursive": {
+      const { name, structure } = content;
+      return {
+        name: <>{name}</>,
+        structure: <SubVariable {...basis} variable={structure} />,
+      };
+    }
+    case "LambdaSet": {
+      const { ambient_function, solved, unspecialized, recursion_var } =
+        content;
+      return {
+        "^": <SubVariable {...basis} variable={ambient_function} />,
+        as: recursion_var ? (
+          <SubVariable {...basis} variable={recursion_var} />
+        ) : null,
+      };
+    }
+    case "ErasedLambda": {
+      return {};
+    }
+    case "Alias": {
+      const { name, real_variable, variables } = content;
+      return {
+        name: <>{name}</>,
+      };
+    }
+    case "Apply": {
+      const { name, variables } = content;
+      return {
+        name: <>{name}</>,
+      };
+    }
+    case "Function": {
+      const { arguments: args, lambda_type, ret } = content;
+      return {
+        args: (
+          <>
+            {args.map((arg, i) => (
+              <SubVariable key={i} {...basis} variable={arg} />
+            ))}
+          </>
+        ),
+        "||": <SubVariable {...basis} variable={lambda_type} />,
+        ret: <SubVariable {...basis} variable={ret} />,
+      };
+    }
+    case "FunctionOrTagUnion": {
+      const { tags, functions, extension } = content;
+      return {
+        tags: <>[{tags.join(", ")}]</>,
+        fns: <>[{functions.join(", ")}]</>,
+      };
+    }
+    case "TagUnion": {
+      const { tags, extension } = content;
+      return {};
+    }
+    case "RecursiveTagUnion": {
+      const { recursion_var, extension, tags } = content;
+      return {
+        as: <SubVariable {...basis} variable={recursion_var} />,
+      };
+    }
+    case "Record": {
+      const { fields, extension } = content;
+      return {};
+    }
+    case "Tuple": {
+      const { elements, extension } = content;
+      return {};
+    }
+    case "RangedNumber": {
+      const { range } = content;
+      return {};
+    }
+    case "EmptyRecord":
+    case "EmptyTuple":
+    case "EmptyTagUnion":
+    case "Error": {
+      return {};
+    }
+    default: {
+      return assertExhaustive(content);
+    }
+  }
+}
+
+interface BasisProps {
+  subs: SubsSnapshot;
+  origin: Variable;
+  addSubVariableLink: AddSubVariableLink;
+}
+
+function SubVariable({
+  subs,
+  origin,
+  variable,
+  addSubVariableLink,
+}: {
+  variable: Variable;
+} & BasisProps): JSX.Element {
+  const desc = subs.get_root(variable);
+  const styles = contentStyles(desc);
+  return (
+    <VariableElHelp
+      variable={variable}
+      styles={styles}
+      onClick={() => addSubVariableLink(origin, variable)}
+    />
+  );
+}

crates/compiler/checkmate/www/src/components/Graph/VariablesGraph.tsx

@@ -0,0 +1,299 @@
+import Dagre from "@dagrejs/dagre";
+import ReactFlow, {
+  Node,
+  Edge,
+  Background,
+  BackgroundVariant,
+  useReactFlow,
+  ReactFlowProvider,
+  NodeChange,
+  applyNodeChanges,
+  EdgeChange,
+  applyEdgeChanges,
+  Panel,
+  NodeTypes,
+} from "reactflow";
+import { useCallback, useState } from "react";
+import { Variable } from "../../schema";
+
+import "reactflow/dist/style.css";
+import clsx from "clsx";
+import VariableNode, { VariableNodeProps } from "./VariableNode";
+import { SubsSnapshot } from "../../engine/subs";
+
+export interface VariablesGraphProps {
+  subs: SubsSnapshot;
+  onVariable: (handler: (variable: Variable) => void) => void;
+}
+
+type GraphDirection = "TB" | "BT" | "LR" | "RL";
+
+const DEFAULT_DIRECTION: GraphDirection = "TB";
+
+interface LayoutedElements {
+  nodes: Node[];
+  edges: Edge[];
+}
+
+interface ComputeLayoutedElementsProps extends LayoutedElements {
+  direction: GraphDirection;
+}
+
+function computeLayoutedElements({
+  nodes,
+  edges,
+  direction,
+}: ComputeLayoutedElementsProps): LayoutedElements {
+  if (nodes.length === 0) {
+    return {
+      nodes: [],
+      edges: [],
+    };
+  }
+
+  const g = new Dagre.graphlib.Graph().setDefaultEdgeLabel(() => ({}));
+  g.setGraph({ rankdir: direction });
+
+  edges.forEach((edge) => g.setEdge(edge.source, edge.target));
+  nodes.forEach((node) => g.setNode(node.id, node));
+
+  Dagre.layout(g);
+
+  const result = {
+    nodes: nodes.map((node) => {
+      const { x, y } = g.node(node.id);
+
+      return { ...node, position: { x, y } };
+    }),
+    edges,
+  };
+  return result;
+}
+
+const NODE_TYPES: NodeTypes = {
+  variable: VariableNode,
+};
+
+function newVariable(id: string, data: VariableNodeProps["data"]): Node {
+  return {
+    id,
+    position: { x: 0, y: 0 },
+    type: "variable",
+    data,
+  };
+}
+
+function addNodeChange(node: Node, existingNodes: Node[]): NodeChange | null {
+  if (existingNodes.some((n) => n.id === node.id)) {
+    return null;
+  }
+  return {
+    type: "add",
+    item: node,
+  };
+}
+
+function addEdgeChange(edge: Edge, existingEdges: Edge[]): EdgeChange | null {
+  if (existingEdges.some((e) => e.id === edge.id)) {
+    return null;
+  }
+  return {
+    type: "add",
+    item: edge,
+  };
+}
+
+function Graph({ subs, onVariable }: VariablesGraphProps): JSX.Element {
+  const instance = useReactFlow();
+  const initialNodes: Node[] = [];
+  const initialEdges: Edge[] = [];
+
+  const [direction, setDirection] = useState(DEFAULT_DIRECTION);
+  const [elements, setElements] = useState<LayoutedElements>({
+    nodes: initialNodes,
+    edges: initialEdges,
+  });
+
+  const refit = useCallback(() => {
+    window.requestAnimationFrame(() => {
+      instance.fitView();
+    });
+  }, [instance]);
+
+  const onLayoutChange = useCallback(
+    (newDirection: GraphDirection) => {
+      setDirection(newDirection);
+
+      setElements(({ nodes, edges }) => {
+        return computeLayoutedElements({
+          direction: newDirection,
+          nodes,
+          edges,
+        });
+      });
+      refit();
+    },
+    [refit]
+  );
+
+  const onNodesChange = useCallback((changes: NodeChange[]) => {
+    setElements(({ nodes, edges }) => {
+      return {
+        nodes: applyNodeChanges(changes, nodes),
+        edges,
+      };
+    });
+  }, []);
+
+  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
+    setElements(({ nodes, edges }) => {
+      return {
+        nodes,
+        edges: applyEdgeChanges(changes, edges),
+      };
+    });
+  }, []);
+
+  const addSubVariableLink = useCallback(
+    (fromN: Variable, subLinkN: Variable) => {
+      fromN = subs.get_root_key(fromN);
+      subLinkN = subs.get_root_key(subLinkN);
+      const from = fromN.toString();
+      const to = subLinkN.toString();
+
+      setElements(({ nodes, edges }) => {
+        const optNewNode = addNodeChange(
+          newVariable(to, {
+            subs,
+            variable: subLinkN,
+            addSubVariableLink,
+          }),
+          nodes
+        );
+        const newNodes = optNewNode
+          ? applyNodeChanges([optNewNode], nodes)
+          : nodes;
+
+        const optNewEdge = addEdgeChange(
+          { id: `${from}->${to}`, source: from, target: to },
+          edges
+        );
+        const newEdges = optNewEdge
+          ? applyEdgeChanges([optNewEdge], edges)
+          : edges;
+
+        return computeLayoutedElements({
+          direction: "TB",
+          nodes: newNodes,
+          edges: newEdges,
+        });
+      });
+
+      refit();
+    },
+    [subs, refit]
+  );
+
+  const addNode = useCallback(
+    (variableN: Variable) => {
+      variableN = subs.get_root_key(variableN);
+      const variable = variableN.toString();
+
+      setElements(({ nodes, edges }) => {
+        const optNewNode = addNodeChange(
+          newVariable(variable, {
+            subs,
+            variable: variableN,
+            addSubVariableLink,
+          }),
+          nodes
+        );
+        const newNodes = optNewNode
+          ? applyNodeChanges([optNewNode], nodes)
+          : nodes;
+
+        return computeLayoutedElements({
+          direction: "TB",
+          nodes: newNodes,
+          edges,
+        });
+      });
+
+      refit();
+    },
+    [subs, refit, addSubVariableLink]
+  );
+
+  onVariable(addNode);
+
+  return (
+    <ReactFlow
+      nodes={elements.nodes}
+      edges={elements.edges}
+      onNodesChange={(e) => onNodesChange(e)}
+      onEdgesChange={(e) => onEdgesChange(e)}
+      fitView
+      nodesDraggable
+      nodesConnectable={false}
+      nodeTypes={NODE_TYPES}
+      proOptions={{
+        // https://reactflow.dev/docs/guides/remove-attribution/
+        hideAttribution: true,
+      }}
+    >
+      <Panel position="top-right">
+        <DirectionPanel
+          direction={direction}
+          onChange={(e) => onLayoutChange(e)}
+        />
+      </Panel>
+
+      <Background variant={BackgroundVariant.Dots} />
+    </ReactFlow>
+  );
+}
+
+interface DirectionPanelProps {
+  direction: GraphDirection;
+  onChange: (direction: GraphDirection) => void;
+}
+
+function DirectionPanel({
+  direction,
+  onChange,
+}: DirectionPanelProps): JSX.Element {
+  const commonStyle = "rounded cursor-pointer text-2xl select-none";
+
+  const dirs: { dir: GraphDirection; text: string }[] = [
+    { dir: "TB", text: "⬇️" },
+    //{ dir: "LR", text: "➡️" },
+  ];
+
+  return (
+    <>
+      {dirs.map(({ dir, text }, i) => (
+        <span
+          key={i}
+          className={clsx(
+            commonStyle,
+            i !== 0 ? "ml-2" : "",
+            dir !== direction ? "opacity-50" : ""
+          )}
+          onClick={() => {
+            onChange(dir);
+          }}
+        >
+          {text}
+        </span>
+      ))}
+    </>
+  );
+}
+
+export default function VariablesGraph(props: VariablesGraphProps) {
+  return (
+    <ReactFlowProvider>
+      <Graph {...props} />
+    </ReactFlowProvider>
+  );
+}