93 lines
2.9 KiB
TypeScript
93 lines
2.9 KiB
TypeScript
import type { Edge, Node, Position } from '@vue-flow/core'
|
|
import dagre from '@dagrejs/dagre'
|
|
|
|
const DEFAULT_NODE_WIDTH = 380
|
|
const DEFAULT_HEADER_HEIGHT = 62
|
|
const MIN_NODE_HEIGHT = 170
|
|
|
|
type LayoutDirection = 'LR' | 'TB'
|
|
|
|
type LayoutOptions = {
|
|
direction?: LayoutDirection
|
|
useSavedPositions?: boolean
|
|
dimensions?: Record<string, { width: number; height: number }>
|
|
}
|
|
|
|
function estimatePortHeight(port: any, direction: 'in' | 'out') {
|
|
let height = 54
|
|
if (port?.type === 'Image' && port?.previewUrl) height += 108
|
|
if (direction === 'in') height += port?.output ? 32 : 36
|
|
return height
|
|
}
|
|
|
|
function estimateColumnHeight(ports: any[] | undefined, direction: 'in' | 'out') {
|
|
if (!Array.isArray(ports) || ports.length === 0) return 44
|
|
return ports.reduce((total, port) => total + estimatePortHeight(port, direction), 0) + Math.max(0, ports.length - 1) * 7
|
|
}
|
|
|
|
function estimateNodeHeight(node: Node) {
|
|
const inputHeight = estimateColumnHeight(node.data?.inputs, 'in')
|
|
const outputHeight = estimateColumnHeight(node.data?.outputs, 'out')
|
|
return Math.max(MIN_NODE_HEIGHT, DEFAULT_HEADER_HEIGHT + Math.max(inputHeight, outputHeight) + 26)
|
|
}
|
|
|
|
function targetPosition(direction: LayoutDirection): Position {
|
|
return direction === 'LR' ? 'left' as Position : 'top' as Position
|
|
}
|
|
|
|
function sourcePosition(direction: LayoutDirection): Position {
|
|
return direction === 'LR' ? 'right' as Position : 'bottom' as Position
|
|
}
|
|
|
|
export function layoutGraph(nodes: Node[], edges: Edge[], options: LayoutOptions = {}) {
|
|
const direction = options.direction || 'LR'
|
|
|
|
if (options.useSavedPositions && nodes.every((node) => node.data?.hasSavedPosition)) {
|
|
return nodes.map((node) => ({
|
|
...node,
|
|
targetPosition: targetPosition(direction),
|
|
sourcePosition: sourcePosition(direction),
|
|
}))
|
|
}
|
|
|
|
const graph = new dagre.graphlib.Graph()
|
|
graph.setDefaultEdgeLabel(() => ({}))
|
|
graph.setGraph({
|
|
rankdir: direction,
|
|
nodesep: 190,
|
|
ranksep: 260,
|
|
marginx: 100,
|
|
marginy: 100,
|
|
})
|
|
|
|
for (const node of nodes) {
|
|
const measured = options.dimensions?.[node.id]
|
|
graph.setNode(node.id, {
|
|
width: Math.max(DEFAULT_NODE_WIDTH, measured?.width || DEFAULT_NODE_WIDTH),
|
|
height: Math.max(estimateNodeHeight(node), measured?.height || 0),
|
|
})
|
|
}
|
|
|
|
for (const edge of edges) {
|
|
graph.setEdge(edge.source, edge.target)
|
|
}
|
|
|
|
dagre.layout(graph)
|
|
|
|
return nodes.map((node) => {
|
|
const graphNode = graph.node(node.id)
|
|
const measured = options.dimensions?.[node.id]
|
|
const width = Math.max(DEFAULT_NODE_WIDTH, measured?.width || DEFAULT_NODE_WIDTH)
|
|
const height = Math.max(estimateNodeHeight(node), measured?.height || 0)
|
|
|
|
return {
|
|
...node,
|
|
targetPosition: targetPosition(direction),
|
|
sourcePosition: sourcePosition(direction),
|
|
position: graphNode
|
|
? { x: graphNode.x - width / 2, y: graphNode.y - height / 2 }
|
|
: node.position,
|
|
}
|
|
})
|
|
}
|