130 lines
3.3 KiB
TypeScript
130 lines
3.3 KiB
TypeScript
import type { Edge, Node } from "reactflow";
|
|
import { Position } from "reactflow";
|
|
import type { GraphData, GraphNodeData, MindmapNode } from "../types/mindmap";
|
|
|
|
const LEVEL_GAP = 320;
|
|
const SIBLING_GAP = 120;
|
|
const PADDING_X = 80;
|
|
const PADDING_Y = 100;
|
|
|
|
type PositionMap = Map<string, { x: number; y: number }>;
|
|
|
|
type TreeToGraphOptions = {
|
|
expandedNodeIds?: Set<string>;
|
|
onToggleNode?: (nodeId: string) => void;
|
|
onOpenNodeChat?: (nodeLabel: string) => void;
|
|
};
|
|
|
|
function getVisibleChildren(
|
|
node: MindmapNode,
|
|
expandedNodeIds: Set<string>,
|
|
): MindmapNode[] {
|
|
if (node.children.length === 0) {
|
|
return [];
|
|
}
|
|
|
|
return expandedNodeIds.has(node.id) ? node.children : [];
|
|
}
|
|
|
|
function assignPositions(
|
|
node: MindmapNode,
|
|
positions: PositionMap,
|
|
leafIndexRef: { value: number },
|
|
expandedNodeIds: Set<string>,
|
|
): number {
|
|
const visibleChildren = getVisibleChildren(node, expandedNodeIds);
|
|
|
|
if (visibleChildren.length === 0) {
|
|
const y = leafIndexRef.value * SIBLING_GAP;
|
|
leafIndexRef.value += 1;
|
|
positions.set(node.id, { x: node.level * LEVEL_GAP, y });
|
|
return y;
|
|
}
|
|
|
|
const childYs = visibleChildren.map((child) =>
|
|
assignPositions(child, positions, leafIndexRef, expandedNodeIds),
|
|
);
|
|
const y =
|
|
childYs.reduce((total, current) => total + current, 0) / childYs.length;
|
|
|
|
positions.set(node.id, { x: node.level * LEVEL_GAP, y });
|
|
return y;
|
|
}
|
|
|
|
function buildGraph(
|
|
node: MindmapNode,
|
|
positions: PositionMap,
|
|
nodes: Node<GraphNodeData>[],
|
|
edges: Edge[],
|
|
expandedNodeIds: Set<string>,
|
|
options: TreeToGraphOptions,
|
|
): void {
|
|
const position = positions.get(node.id);
|
|
|
|
if (!position) {
|
|
return;
|
|
}
|
|
|
|
const hasChildren = node.children.length > 0;
|
|
const isExpanded = expandedNodeIds.has(node.id);
|
|
const onToggleNode = options.onToggleNode;
|
|
const onOpenNodeChat = options.onOpenNodeChat;
|
|
|
|
nodes.push({
|
|
id: node.id,
|
|
type: "mindmapNode",
|
|
position: {
|
|
x: position.x + PADDING_X,
|
|
y: position.y + PADDING_Y,
|
|
},
|
|
style: {
|
|
cursor: "default",
|
|
},
|
|
data: {
|
|
label: node.label,
|
|
level: node.level,
|
|
isLeaf: node.is_leaf,
|
|
hasChildren,
|
|
isExpanded,
|
|
onToggle: hasChildren && onToggleNode ? () => onToggleNode(node.id) : undefined,
|
|
onOpenChat: onOpenNodeChat ? () => onOpenNodeChat(node.label) : undefined,
|
|
},
|
|
sourcePosition: Position.Right,
|
|
targetPosition: Position.Left,
|
|
draggable: false,
|
|
selectable: false,
|
|
});
|
|
|
|
getVisibleChildren(node, expandedNodeIds).forEach((child) => {
|
|
edges.push({
|
|
id: `${node.id}-${child.id}`,
|
|
source: node.id,
|
|
target: child.id,
|
|
type: "smoothstep",
|
|
animated: false,
|
|
style: {
|
|
stroke: "#60a5fa",
|
|
strokeWidth: 2,
|
|
},
|
|
});
|
|
|
|
buildGraph(child, positions, nodes, edges, expandedNodeIds, options);
|
|
});
|
|
}
|
|
|
|
export function treeToGraph(
|
|
tree: MindmapNode,
|
|
options: TreeToGraphOptions = {},
|
|
): GraphData {
|
|
const positions: PositionMap = new Map();
|
|
const leafIndexRef = { value: 0 };
|
|
const nodes: Node<GraphNodeData>[] = [];
|
|
const edges: Edge[] = [];
|
|
const expandedNodeIds = options.expandedNodeIds ?? new Set([tree.id]);
|
|
|
|
assignPositions(tree, positions, leafIndexRef, expandedNodeIds);
|
|
buildGraph(tree, positions, nodes, edges, expandedNodeIds, options);
|
|
|
|
return { nodes, edges };
|
|
}
|