import dagre from 'dagre'; const edgeType = 'smoothstep'; export const flattenTree = (tree) => { let stack = [tree]; let nodes = []; let edges = []; while (stack.length > 0) { let node = stack.pop(); nodes.push({ id: node.id, data: { ...node, label: `${node.name}` }, position: node.position }); if (node.children && node.children.length > 0) { node.children.map(item => { edges.push({ source: node.id, target: item.id, id: item.id, type: edgeType, animated: true }); stack.unshift(item); }); }; }; return { nodes, edges }; }; export const getLayoutedElements = (nodes, edges, direction = 'TB') => { const dagreGraph = new dagre.graphlib.Graph(); dagreGraph.setDefaultEdgeLabel(() => ({})); const nodeWidth = 172; const nodeHeight = 36; const isHorizontal = direction === 'LR'; dagreGraph.setGraph({ rankdir: direction }); nodes.forEach((node: any) => { dagreGraph.setNode(node.id, { width: nodeWidth, height: nodeHeight }); }); edges.forEach((edge: any) => { dagreGraph.setEdge(edge.source, edge.target); }); dagre.layout(dagreGraph); nodes.forEach((node: any) => { const nodeWithPosition = dagreGraph.node(node.id); node.targetPosition = isHorizontal ? 'left' : 'top'; node.sourcePosition = isHorizontal ? 'right' : 'bottom'; node.position = { x: nodeWithPosition.x - nodeWidth / 2, y: nodeWithPosition.y - nodeHeight / 2, }; return node; }); return { nodes, edges }; };