Pin (bindings)

import {
	BindingOnShapeChangeOptions,
	BindingOnShapeDeleteOptions,
	BindingUtil,
	Box,
	DefaultFillStyle,
	DefaultToolbar,
	DefaultToolbarContent,
	RecordProps,
	Rectangle2d,
	ShapeUtil,
	StateNode,
	TLBinding,
	TLEditorComponents,
	TLPointerEventInfo,
	TLShape,
	TLShapeId,
	TLShapePartial,
	TLShapeUtilCanBindOpts,
	TLUiComponents,
	TLUiOverrides,
	Tldraw,
	TldrawUiMenuItem,
	Vec,
	VecModel,
	createShapeId,
	invLerp,
	lerp,
	useIsToolSelected,
	useTools,
} from 'tldraw'
import 'tldraw/tldraw.css'

const PIN_TYPE = 'pin'

declare module 'tldraw' {
	export interface TLGlobalShapePropsMap {
		[PIN_TYPE]: Record<string, never>
	}
}

type PinShape = TLShape<typeof PIN_TYPE>

const offsetX = -16
const offsetY = -26
class PinShapeUtil extends ShapeUtil<PinShape> {
	static override type = PIN_TYPE
	static override props: RecordProps<PinShape> = {}

	override getDefaultProps() {
		return {}
	}

	override canBind({ toShapeType, bindingType }: TLShapeUtilCanBindOpts<PinShape>) {
		if (bindingType === PIN_TYPE) {
			// pins cannot bind to other pins!
			return toShapeType !== PIN_TYPE
		}
		// Allow pins to participate in other bindings, e.g. arrows
		return true
	}
	override canEdit() {
		return false
	}
	override canResize() {
		return false
	}
	override hideRotateHandle() {
		return true
	}
	override isAspectRatioLocked() {
		return true
	}

	override getGeometry() {
		return new Rectangle2d({
			width: 32,
			height: 32,
			x: offsetX,
			y: offsetY,
			isFilled: true,
		})
	}

	override component() {
		return (
			<div
				style={{
					width: '100%',
					height: '100%',
					marginLeft: offsetX,
					marginTop: offsetY,
					fontSize: '26px',
					textAlign: 'center',
				}}
			>
				📍
			</div>
		)
	}

	override indicator() {
		return <rect width={32} height={32} x={offsetX} y={offsetY} />
	}

	override onTranslateStart(shape: PinShape) {
		const bindings = this.editor.getBindingsFromShape(shape, PIN_TYPE)
		this.editor.deleteBindings(bindings)
	}

	override onTranslateEnd(_initial: PinShape, pin: PinShape) {
		const pageAnchor = this.editor.getShapePageTransform(pin).applyToPoint({ x: 0, y: 0 })

		const targets = this.editor
			.getShapesAtPoint(pageAnchor, { hitInside: true })
			.filter(
				(shape) =>
					this.editor.canBindShapes({ fromShape: pin, toShape: shape, binding: PIN_TYPE }) &&
					shape.parentId === pin.parentId &&
					shape.index < pin.index
			)

		for (const target of targets) {
			const targetBounds = Box.ZeroFix(this.editor.getShapeGeometry(target)!.bounds)
			const pointInTargetSpace = this.editor.getPointInShapeSpace(target, pageAnchor)

			const anchor = {
				x: invLerp(targetBounds.minX, targetBounds.maxX, pointInTargetSpace.x),
				y: invLerp(targetBounds.minY, targetBounds.maxY, pointInTargetSpace.y),
			}

			this.editor.createBinding({
				type: PIN_TYPE,
				fromId: pin.id,
				toId: target.id,
				props: {
					anchor,
				},
			})
		}
	}
}

declare module 'tldraw' {
	export interface TLGlobalBindingPropsMap {
		[PIN_TYPE]: {
			anchor: VecModel
		}
	}
}

type PinBinding = TLBinding<typeof PIN_TYPE>

class PinBindingUtil extends BindingUtil<PinBinding> {
	static override type = PIN_TYPE

	override getDefaultProps() {
		return {
			anchor: { x: 0.5, y: 0.5 },
		}
	}

	private changedToShapes = new Set<TLShapeId>()

	override onOperationComplete(): void {
		if (this.changedToShapes.size === 0) return

		const fixedShapes = this.changedToShapes
		const toCheck = [...this.changedToShapes]

		const initialPositions = new Map<TLShapeId, VecModel>()
		const targetDeltas = new Map<TLShapeId, Map<TLShapeId, VecModel>>()

		const addTargetDelta = (fromId: TLShapeId, toId: TLShapeId, delta: VecModel) => {
			if (!targetDeltas.has(fromId)) targetDeltas.set(fromId, new Map())
			targetDeltas.get(fromId)!.set(toId, delta)

			if (!targetDeltas.has(toId)) targetDeltas.set(toId, new Map())
			targetDeltas.get(toId)!.set(fromId, { x: -delta.x, y: -delta.y })
		}

		const allShapes = new Set<TLShapeId>()
		while (toCheck.length) {
			const shapeId = toCheck.pop()!

			const shape = this.editor.getShape(shapeId)
			if (!shape) continue

			if (allShapes.has(shapeId)) continue
			allShapes.add(shapeId)

			const bindings = this.editor.getBindingsToShape(shape, PIN_TYPE)
			for (const binding of bindings) {
				if (allShapes.has(binding.fromId)) continue
				allShapes.add(binding.fromId)

				const pin = this.editor.getShape<PinShape>(binding.fromId)
				if (!pin) continue

				const pinPosition = this.editor.getShapePageTransform(pin).applyToPoint({ x: 0, y: 0 })
				initialPositions.set(pin.id, pinPosition)

				for (const binding of this.editor.getBindingsFromShape(pin.id, PIN_TYPE)) {
					const shapeBounds = this.editor.getShapeGeometry(binding.toId)!.bounds
					const shapeAnchor = {
						x: lerp(shapeBounds.minX, shapeBounds.maxX, binding.props.anchor.x),
						y: lerp(shapeBounds.minY, shapeBounds.maxY, binding.props.anchor.y),
					}
					const currentPageAnchor = this.editor
						.getShapePageTransform(binding.toId)
						.applyToPoint(shapeAnchor)

					const shapeOrigin = this.editor
						.getShapePageTransform(binding.toId)
						.applyToPoint({ x: 0, y: 0 })
					initialPositions.set(binding.toId, shapeOrigin)

					addTargetDelta(pin.id, binding.toId, {
						x: currentPageAnchor.x - shapeOrigin.x,
						y: currentPageAnchor.y - shapeOrigin.y,
					})

					if (!allShapes.has(binding.toId)) toCheck.push(binding.toId)
				}
			}
		}

		const currentPositions = new Map(initialPositions)

		const iterations = 30
		for (let i = 0; i < iterations; i++) {
			const movements = new Map<TLShapeId, VecModel[]>()
			for (const [aId, deltas] of targetDeltas) {
				if (fixedShapes.has(aId)) continue
				const aPosition = currentPositions.get(aId)!
				for (const [bId, targetDelta] of deltas) {
					const bPosition = currentPositions.get(bId)!

					const adjustmentDelta = {
						x: targetDelta.x - (aPosition.x - bPosition.x),
						y: targetDelta.y - (aPosition.y - bPosition.y),
					}

					if (!movements.has(aId)) movements.set(aId, [])
					movements.get(aId)!.push(adjustmentDelta)
				}
			}

			for (const [shapeId, deltas] of movements) {
				const currentPosition = currentPositions.get(shapeId)!
				currentPositions.set(shapeId, Vec.Average(deltas).add(currentPosition))
			}
		}

		const updates: TLShapePartial[] = []
		for (const [shapeId, position] of currentPositions) {
			const delta = Vec.Sub(position, initialPositions.get(shapeId)!)
			if (delta.len2() <= 0.01) continue

			const newPosition = this.editor.getPointInParentSpace(shapeId, position)
			updates.push({
				...this.editor.getShape(shapeId)!,
				id: shapeId,
				x: newPosition.x,
				y: newPosition.y,
			})
		}

		if (updates.length === 0) {
			this.changedToShapes.clear()
		} else {
			this.editor.updateShapes(updates)
		}
	}

	// when the shape we're stuck to changes, update the pin's position
	override onAfterChangeToShape({
		binding,
		shapeAfter,
	}: BindingOnShapeChangeOptions<PinBinding>): void {
		this.changedToShapes.add(binding.toId)
		const pin = this.editor.getShape(binding.fromId)
		if (!pin) return

		// If the bound shape changed parents, reparent the pin to follow
		if (pin.parentId !== shapeAfter.parentId) {
			this.editor.reparentShapes([pin.id], shapeAfter.parentId)
		}
	}

	// when the thing we're stuck to is deleted, delete the pin too
	override onBeforeDeleteToShape({ binding }: BindingOnShapeDeleteOptions<PinBinding>): void {
		this.editor.deleteShape(binding.fromId)
	}
}

class PinTool extends StateNode {
	static override id = PIN_TYPE

	override onEnter() {
		this.editor.setCursor({ type: 'cross', rotation: 0 })
	}

	override onPointerDown(info: TLPointerEventInfo) {
		const currentPagePoint = this.editor.inputs.getCurrentPagePoint()
		const pinId = createShapeId()
		this.editor.markHistoryStoppingPoint()
		this.editor.createShape({
			id: pinId,
			type: PIN_TYPE,
			x: currentPagePoint.x,
			y: currentPagePoint.y,
		})
		this.editor.setSelectedShapes([pinId])
		this.editor.setCurrentTool('select.translating', {
			...info,
			target: 'shape',
			shape: this.editor.getShape(pinId),
			isCreating: true,
			onInteractionEnd: 'pin',
			onCreate: () => {
				this.editor.setCurrentTool('pin')
			},
		})
	}
}

const overrides: TLUiOverrides = {
	tools(editor, schema) {
		schema['pin'] = {
			id: 'pin',
			label: 'Pin',
			icon: 'heart-icon',
			kbd: 'p',
			onSelect: () => {
				editor.setCurrentTool('pin')
			},
		}
		return schema
	},
}

const components: TLUiComponents & TLEditorComponents = {
	Toolbar: (...props) => {
		const pin = useTools().pin
		const isPinSelected = useIsToolSelected(pin)
		return (
			<DefaultToolbar {...props}>
				<TldrawUiMenuItem {...pin} isSelected={isPinSelected} />
				<DefaultToolbarContent />
			</DefaultToolbar>
		)
	},
}

export default function PinExample() {
	return (
		<div className="tldraw__editor">
			<Tldraw
				persistenceKey="pin-example"
				onMount={(editor) => {
					;(window as any).editor = editor
					editor.setStyleForNextShapes(DefaultFillStyle, 'semi')
				}}
				shapeUtils={[PinShapeUtil]}
				bindingUtils={[PinBindingUtil]}
				tools={[PinTool]}
				overrides={overrides}
				components={components}
			/>
		</div>
	)
}