281 lines
8.5 KiB
TypeScript
281 lines
8.5 KiB
TypeScript
import { Layout } from 'webcola';
|
|
import { BaseCollection } from '../collections';
|
|
import { Editor, TLArrowShape, TLGeoShape, TLShape, TLShapeId } from '@tldraw/tldraw';
|
|
|
|
type ColaNode = {
|
|
id: TLShapeId;
|
|
x: number;
|
|
y: number;
|
|
width: number;
|
|
height: number;
|
|
rotation: number;
|
|
color?: string;
|
|
};
|
|
type ColaIdLink = {
|
|
source: TLShapeId
|
|
target: TLShapeId
|
|
};
|
|
type ColaNodeLink = {
|
|
source: ColaNode
|
|
target: ColaNode
|
|
};
|
|
|
|
type AlignmentConstraint = {
|
|
type: 'alignment',
|
|
axis: 'x' | 'y',
|
|
offsets: { node: TLShapeId, offset: number }[]
|
|
}
|
|
|
|
type ColaConstraint = AlignmentConstraint
|
|
|
|
export class GraphLayoutCollection extends BaseCollection {
|
|
override id = 'graph';
|
|
graphSim: Layout;
|
|
animFrame = -1;
|
|
colaNodes: Map<TLShapeId, ColaNode> = new Map();
|
|
colaLinks: Map<TLShapeId, ColaIdLink> = new Map();
|
|
colaConstraints: ColaConstraint[] = [];
|
|
|
|
constructor(editor: Editor) {
|
|
super(editor)
|
|
this.graphSim = new Layout();
|
|
const simLoop = () => {
|
|
this.step();
|
|
this.animFrame = requestAnimationFrame(simLoop);
|
|
};
|
|
simLoop();
|
|
}
|
|
|
|
override onAdd(shapes: TLShape[]) {
|
|
for (const shape of shapes) {
|
|
if (shape.type !== "arrow") {
|
|
this.addGeo(shape);
|
|
}
|
|
else {
|
|
this.addArrow(shape as TLArrowShape);
|
|
}
|
|
}
|
|
this.refreshGraph();
|
|
}
|
|
|
|
override onRemove(shapes: TLShape[]) {
|
|
const removedShapeIds = new Set(shapes.map(shape => shape.id));
|
|
|
|
for (const shape of shapes) {
|
|
this.colaNodes.delete(shape.id);
|
|
this.colaLinks.delete(shape.id);
|
|
}
|
|
|
|
// Filter out links where either source or target has been removed
|
|
for (const [key, link] of this.colaLinks) {
|
|
if (removedShapeIds.has(link.source) || removedShapeIds.has(link.target)) {
|
|
this.colaLinks.delete(key);
|
|
}
|
|
}
|
|
|
|
this.refreshGraph();
|
|
}
|
|
|
|
override onShapeChange(prev: TLShape, next: TLShape) {
|
|
if (prev.type === 'geo' && next.type === 'geo') {
|
|
const prevShape = prev as TLGeoShape
|
|
const nextShape = next as TLGeoShape
|
|
// update color if its changed and refresh constraints which use this
|
|
if (prevShape.props.color !== nextShape.props.color) {
|
|
const existingNode = this.colaNodes.get(next.id);
|
|
if (existingNode) {
|
|
this.colaNodes.set(next.id, {
|
|
...existingNode,
|
|
color: nextShape.props.color,
|
|
});
|
|
}
|
|
this.refreshGraph();
|
|
}
|
|
}
|
|
}
|
|
|
|
step = () => {
|
|
this.graphSim.start(1, 0, 0, 0, true, false);
|
|
for (const node of this.graphSim.nodes() as ColaNode[]) {
|
|
|
|
const shape = this.editor.getShape(node.id);
|
|
const { w, h } = this.editor.getShapeGeometry(node.id).bounds
|
|
if (!shape) continue;
|
|
|
|
const { x, y } = getCornerToCenterOffset(w, h, shape.rotation);
|
|
|
|
// Fix positions if we're dragging them
|
|
if (this.editor.getSelectedShapeIds().includes(node.id)) {
|
|
node.x = shape.x + x;
|
|
node.y = shape.y + y;
|
|
}
|
|
|
|
// Update shape props
|
|
node.width = w;
|
|
node.height = h;
|
|
node.rotation = shape.rotation;
|
|
|
|
this.editor.updateShape({
|
|
id: node.id,
|
|
type: "geo",
|
|
x: node.x - x,
|
|
y: node.y - y,
|
|
props: {
|
|
...shape.props,
|
|
richText: (shape.props as any)?.richText || [] as any, // Ensure richText exists
|
|
},
|
|
});
|
|
}
|
|
};
|
|
|
|
addArrow = (arrow: TLArrowShape) => {
|
|
const bindings = this.editor.getBindingsInvolvingShape(arrow.id);
|
|
if (bindings.length !== 2) return;
|
|
|
|
const startBinding = bindings.find(binding => (binding.props as any).terminal === 'start');
|
|
const endBinding = bindings.find(binding => (binding.props as any).terminal === 'end');
|
|
|
|
if (startBinding && endBinding) {
|
|
const source = this.editor.getShape(startBinding.toId);
|
|
const target = this.editor.getShape(endBinding.toId);
|
|
|
|
if (source && target) {
|
|
const link: ColaIdLink = {
|
|
source: source.id,
|
|
target: target.id
|
|
};
|
|
this.colaLinks.set(arrow.id, link);
|
|
}
|
|
}
|
|
}
|
|
|
|
addGeo = (shape: TLShape) => {
|
|
const { w, h } = this.editor.getShapeGeometry(shape).bounds
|
|
const { x, y } = getCornerToCenterOffset(w, h, shape.rotation)
|
|
const node: ColaNode = {
|
|
id: shape.id,
|
|
x: shape.x + x,
|
|
y: shape.y + y,
|
|
width: w,
|
|
height: h,
|
|
rotation: shape.rotation,
|
|
color: (shape.props as any).color
|
|
};
|
|
this.colaNodes.set(shape.id, node);
|
|
}
|
|
|
|
refreshGraph() {
|
|
// TODO: remove this hardcoded behaviour
|
|
this.editor.selectNone()
|
|
this.refreshConstraints();
|
|
const nodes = [...this.colaNodes.values()];
|
|
const nodeIdToIndex = new Map(nodes.map((n, i) => [n.id, i]));
|
|
// Convert the Map values to an array for processing
|
|
const links = Array.from(this.colaLinks.values()).map(l => ({
|
|
source: nodeIdToIndex.get(l.source),
|
|
target: nodeIdToIndex.get(l.target)
|
|
}));
|
|
|
|
const constraints = this.colaConstraints.map(constraint => {
|
|
if (constraint.type === 'alignment') {
|
|
return {
|
|
...constraint,
|
|
offsets: constraint.offsets.map(offset => ({
|
|
node: nodeIdToIndex.get(offset.node),
|
|
offset: offset.offset
|
|
}))
|
|
};
|
|
}
|
|
return constraint;
|
|
});
|
|
|
|
this.graphSim
|
|
.nodes(nodes)
|
|
// @ts-ignore
|
|
.links(links)
|
|
.constraints(constraints)
|
|
// you could use .linkDistance(250) too, which is stable but does not handle size/rotation
|
|
.linkDistance((edge) => calcEdgeDistance(edge as ColaNodeLink))
|
|
.avoidOverlaps(true)
|
|
.handleDisconnected(true)
|
|
}
|
|
|
|
refreshConstraints() {
|
|
const alignmentConstraintX: AlignmentConstraint = {
|
|
type: 'alignment',
|
|
axis: 'x',
|
|
offsets: [],
|
|
};
|
|
const alignmentConstraintY: AlignmentConstraint = {
|
|
type: 'alignment',
|
|
axis: 'y',
|
|
offsets: [],
|
|
};
|
|
|
|
// Iterate over shapes and generate constraints based on conditions
|
|
for (const node of this.colaNodes.values()) {
|
|
if (node.color === "red") {
|
|
// Add alignment offset for red shapes
|
|
alignmentConstraintX.offsets.push({ node: node.id, offset: 0 });
|
|
}
|
|
if (node.color === "blue") {
|
|
// Add alignment offset for red shapes
|
|
alignmentConstraintY.offsets.push({ node: node.id, offset: 0 });
|
|
}
|
|
}
|
|
|
|
const constraints = [];
|
|
if (alignmentConstraintX.offsets.length > 0) {
|
|
constraints.push(alignmentConstraintX);
|
|
}
|
|
if (alignmentConstraintY.offsets.length > 0) {
|
|
constraints.push(alignmentConstraintY);
|
|
}
|
|
this.colaConstraints = constraints;
|
|
}
|
|
}
|
|
|
|
function getCornerToCenterOffset(w: number, h: number, rotation: number) {
|
|
|
|
// Calculate the center coordinates relative to the top-left corner
|
|
const centerX = w / 2;
|
|
const centerY = h / 2;
|
|
|
|
// Apply rotation to the center coordinates
|
|
const rotatedCenterX = centerX * Math.cos(rotation) - centerY * Math.sin(rotation);
|
|
const rotatedCenterY = centerX * Math.sin(rotation) + centerY * Math.cos(rotation);
|
|
|
|
return { x: rotatedCenterX, y: rotatedCenterY };
|
|
}
|
|
|
|
function calcEdgeDistance(edge: ColaNodeLink) {
|
|
const LINK_DISTANCE = 100;
|
|
|
|
// horizontal and vertical distances between centers
|
|
const dx = edge.target.x - edge.source.x;
|
|
const dy = edge.target.y - edge.source.y;
|
|
|
|
// the angles of the nodes in radians
|
|
const sourceAngle = edge.source.rotation;
|
|
const targetAngle = edge.target.rotation;
|
|
|
|
// Calculate the rotated dimensions of the nodes
|
|
const sourceWidth = Math.abs(edge.source.width * Math.cos(sourceAngle)) + Math.abs(edge.source.height * Math.sin(sourceAngle));
|
|
const sourceHeight = Math.abs(edge.source.width * Math.sin(sourceAngle)) + Math.abs(edge.source.height * Math.cos(sourceAngle));
|
|
const targetWidth = Math.abs(edge.target.width * Math.cos(targetAngle)) + Math.abs(edge.target.height * Math.sin(targetAngle));
|
|
const targetHeight = Math.abs(edge.target.width * Math.sin(targetAngle)) + Math.abs(edge.target.height * Math.cos(targetAngle));
|
|
|
|
// Calculate edge-to-edge distances
|
|
const horizontalGap = Math.max(0, Math.abs(dx) - (sourceWidth + targetWidth) / 2);
|
|
const verticalGap = Math.max(0, Math.abs(dy) - (sourceHeight + targetHeight) / 2);
|
|
|
|
// Calculate straight-line distance between the centers of the nodes
|
|
const centerToCenterDistance = Math.sqrt(dx * dx + dy * dy);
|
|
|
|
// Adjust the distance by subtracting the edge-to-edge distance and adding the desired travel distance
|
|
const adjustedDistance = centerToCenterDistance -
|
|
Math.sqrt(horizontalGap * horizontalGap + verticalGap * verticalGap) +
|
|
LINK_DISTANCE;
|
|
|
|
return adjustedDistance;
|
|
}; |