@tldraw/editor 4.3.0-next.f4772c19540d → 4.4.0-canary.29afdff6bb04

package/src/lib/editor/managers/SpatialIndexManager/RBushIndex.ts

@@ -0,0 +1,144 @@
+import type { TLShapeId } from '@tldraw/tlschema'
+import RBush from 'rbush'
+import { Box } from '../../../primitives/Box'
+
+/**
+ * Element stored in the R-tree spatial index.
+ * Contains bounds (minX, minY, maxX, maxY) and shape ID.
+ */
+export interface SpatialElement {
+	minX: number
+	minY: number
+	maxX: number
+	maxY: number
+	id: TLShapeId
+}
+
+/**
+ * Custom RBush class for tldraw shapes.
+ */
+class TldrawRBush extends RBush<SpatialElement> {}
+
+/**
+ * Wrapper around RBush R-tree for efficient spatial queries.
+ * Maintains a map of elements currently in the tree for efficient updates.
+ */
+export class RBushIndex {
+	private rBush: TldrawRBush
+	private elementsInTree: Map<TLShapeId, SpatialElement>
+
+	constructor() {
+		this.rBush = new TldrawRBush()
+		this.elementsInTree = new Map()
+	}
+
+	/**
+	 * Search for shapes within the given bounds.
+	 * Returns set of shape IDs that intersect with the bounds.
+	 */
+	search(bounds: Box): Set<TLShapeId> {
+		const results = this.rBush.search({
+			minX: bounds.minX,
+			minY: bounds.minY,
+			maxX: bounds.maxX,
+			maxY: bounds.maxY,
+		})
+		return new Set(results.map((e: SpatialElement) => e.id))
+	}
+
+	/**
+	 * Insert or update a shape in the spatial index.
+	 * If the shape already exists, it will be removed first to prevent duplicates.
+	 */
+	upsert(id: TLShapeId, bounds: Box): void {
+		// Remove existing entry to prevent map-tree desync
+		const existing = this.elementsInTree.get(id)
+		if (existing) {
+			this.rBush.remove(existing)
+		}
+
+		const element: SpatialElement = {
+			minX: bounds.minX,
+			minY: bounds.minY,
+			maxX: bounds.maxX,
+			maxY: bounds.maxY,
+			id,
+		}
+		this.rBush.insert(element)
+		this.elementsInTree.set(id, element)
+	}
+
+	/**
+	 * Remove a shape from the spatial index.
+	 */
+	remove(id: TLShapeId): void {
+		const element = this.elementsInTree.get(id)
+		if (element) {
+			this.rBush.remove(element)
+			this.elementsInTree.delete(id)
+		}
+	}
+
+	/**
+	 * Bulk load elements into the spatial index.
+	 * More efficient than individual inserts for initial loading.
+	 */
+	bulkLoad(elements: SpatialElement[]): void {
+		this.rBush.load(elements)
+		for (const element of elements) {
+			this.elementsInTree.set(element.id, element)
+		}
+	}
+
+	/**
+	 * Clear all elements from the spatial index.
+	 */
+	clear(): void {
+		this.rBush.clear()
+		this.elementsInTree.clear()
+	}
+
+	/**
+	 * Check if a shape is in the spatial index.
+	 */
+	has(id: TLShapeId): boolean {
+		return this.elementsInTree.has(id)
+	}
+
+	/**
+	 * Get the number of elements in the spatial index.
+	 */
+	getSize(): number {
+		return this.elementsInTree.size
+	}
+
+	/**
+	 * Get all shape IDs currently in the spatial index.
+	 */
+	getAllShapeIds(): TLShapeId[] {
+		return Array.from(this.elementsInTree.keys())
+	}
+
+	/**
+	 * Get the bounds currently stored in the spatial index for a shape.
+	 * Returns undefined if the shape is not in the index.
+	 */
+	getBounds(id: TLShapeId): Box | undefined {
+		const element = this.elementsInTree.get(id)
+		if (!element) return undefined
+		return new Box(
+			element.minX,
+			element.minY,
+			element.maxX - element.minX,
+			element.maxY - element.minY
+		)
+	}
+
+	/**
+	 * Dispose of the spatial index.
+	 * Clears all data structures to prevent memory leaks.
+	 */
+	dispose(): void {
+		this.clear()
+	}
+}

package/src/lib/editor/managers/SpatialIndexManager/SpatialIndexManager.ts

@@ -0,0 +1,215 @@
+import { Computed, RESET_VALUE, computed, isUninitialized } from '@tldraw/state'
+import type { RecordsDiff } from '@tldraw/store'
+import type { TLRecord } from '@tldraw/tlschema'
+import { TLPageId, TLShape, TLShapeId, isShape } from '@tldraw/tlschema'
+import { objectMapValues } from '@tldraw/utils'
+import { Box } from '../../../primitives/Box'
+import type { Editor } from '../../Editor'
+import { RBushIndex, type SpatialElement } from './RBushIndex'
+
+/**
+ * Manages spatial indexing for efficient shape location queries.
+ *
+ * Uses an R-tree (via RBush) to enable O(log n) spatial queries instead of O(n) iteration.
+ * Handles shapes with computed bounds (arrows, groups, custom shapes) by checking all shapes'
+ * bounds on each update using the reactive bounds cache.
+ *
+ * Key features:
+ * - Incremental updates using filterHistory pattern
+ * - Leverages existing bounds cache reactivity for dependency tracking
+ * - Works with any custom shape type with computed bounds
+ * - Per-page index (rebuilds on page change)
+ * - Optimized for viewport culling queries
+ *
+ * @internal
+ */
+export class SpatialIndexManager {
+	private rbush: RBushIndex
+	private spatialIndexComputed: Computed<number>
+	private lastPageId: TLPageId | null = null
+
+	constructor(public readonly editor: Editor) {
+		this.rbush = new RBushIndex()
+		this.spatialIndexComputed = this.createSpatialIndexComputed()
+	}
+
+	private createSpatialIndexComputed() {
+		const shapeHistory = this.editor.store.query.filterHistory('shape')
+
+		return computed<number>('spatialIndex', (_prevValue, lastComputedEpoch) => {
+			if (isUninitialized(_prevValue)) {
+				return this.buildFromScratch(lastComputedEpoch)
+			}
+
+			const shapeDiff = shapeHistory.getDiffSince(lastComputedEpoch)
+
+			if (shapeDiff === RESET_VALUE) {
+				return this.buildFromScratch(lastComputedEpoch)
+			}
+
+			const currentPageId = this.editor.getCurrentPageId()
+			if (this.lastPageId !== currentPageId) {
+				return this.buildFromScratch(lastComputedEpoch)
+			}
+
+			// No shape changes - index is already up to date
+			if (shapeDiff.length === 0) {
+				return lastComputedEpoch
+			}
+
+			// Process incremental updates
+			this.processIncrementalUpdate(shapeDiff)
+
+			return lastComputedEpoch
+		})
+	}
+
+	private buildFromScratch(epoch: number): number {
+		this.rbush.clear()
+		this.lastPageId = this.editor.getCurrentPageId()
+
+		const shapes = this.editor.getCurrentPageShapes()
+		const elements: SpatialElement[] = []
+
+		// Collect all shape elements for bulk loading
+		for (const shape of shapes) {
+			const bounds = this.editor.getShapePageBounds(shape.id)
+			if (bounds) {
+				elements.push({
+					minX: bounds.minX,
+					minY: bounds.minY,
+					maxX: bounds.maxX,
+					maxY: bounds.maxY,
+					id: shape.id,
+				})
+			}
+		}
+
+		// Bulk load for efficiency
+		this.rbush.bulkLoad(elements)
+
+		return epoch
+	}
+
+	private processIncrementalUpdate(shapeDiff: RecordsDiff<TLRecord>[]): void {
+		// Track shapes we've already processed from the diff
+		const processedShapeIds = new Set<TLShapeId>()
+
+		// 1. Process shape additions, removals, and updates from diff
+		for (const changes of shapeDiff) {
+			// Handle additions (only for shapes on current page)
+			for (const shape of objectMapValues(changes.added) as TLShape[]) {
+				if (isShape(shape) && this.editor.getAncestorPageId(shape) === this.lastPageId) {
+					const bounds = this.editor.getShapePageBounds(shape.id)
+					if (bounds) {
+						this.rbush.upsert(shape.id, bounds)
+					}
+					processedShapeIds.add(shape.id)
+				}
+			}
+
+			// Handle removals
+			for (const shape of objectMapValues(changes.removed) as TLShape[]) {
+				if (isShape(shape)) {
+					this.rbush.remove(shape.id)
+					processedShapeIds.add(shape.id)
+				}
+			}
+
+			// Handle updated shapes: page changes and bounds updates
+			for (const [from, to] of objectMapValues(changes.updated) as [TLShape, TLShape][]) {
+				if (!isShape(to)) continue
+				processedShapeIds.add(to.id)
+
+				const wasOnPage = this.editor.getAncestorPageId(from) === this.lastPageId
+				const isOnPage = this.editor.getAncestorPageId(to) === this.lastPageId
+
+				if (isOnPage) {
+					const bounds = this.editor.getShapePageBounds(to.id)
+					if (bounds) {
+						this.rbush.upsert(to.id, bounds)
+					}
+				} else if (wasOnPage) {
+					this.rbush.remove(to.id)
+				}
+			}
+		}
+
+		// 2. Check remaining shapes in index for bounds changes
+		// This handles shapes with computed bounds (arrows bound to moved shapes, groups with moved children, etc.)
+		const allShapeIds = this.rbush.getAllShapeIds()
+
+		for (const shapeId of allShapeIds) {
+			if (processedShapeIds.has(shapeId)) continue
+
+			const currentBounds = this.editor.getShapePageBounds(shapeId)
+			const indexedBounds = this.rbush.getBounds(shapeId)
+
+			if (!this.areBoundsEqual(currentBounds, indexedBounds)) {
+				if (currentBounds) {
+					this.rbush.upsert(shapeId, currentBounds)
+				} else {
+					this.rbush.remove(shapeId)
+				}
+			}
+		}
+	}
+
+	private areBoundsEqual(a: Box | undefined, b: Box | undefined): boolean {
+		if (!a && !b) return true
+		if (!a || !b) return false
+		return a.minX === b.minX && a.minY === b.minY && a.maxX === b.maxX && a.maxY === b.maxY
+	}
+
+	/**
+	 * Get shape IDs within the given bounds.
+	 * Optimized for viewport culling queries.
+	 *
+	 * Note: Results are unordered. If you need z-order, combine with sorted shapes:
+	 * ```ts
+	 * const candidates = editor.spatialIndex.getShapeIdsInsideBounds(bounds)
+	 * const sorted = editor.getCurrentPageShapesSorted().filter(s => candidates.has(s.id))
+	 * ```
+	 *
+	 * @param bounds - The bounds to search within
+	 * @returns Unordered set of shape IDs within the bounds
+	 *
+	 * @public
+	 */
+	getShapeIdsInsideBounds(bounds: Box): Set<TLShapeId> {
+		this.spatialIndexComputed.get()
+		return this.rbush.search(bounds)
+	}
+
+	/**
+	 * Get shape IDs at a point (with optional margin).
+	 * Creates a small bounding box around the point and searches the spatial index.
+	 *
+	 * Note: Results are unordered. If you need z-order, combine with sorted shapes:
+	 * ```ts
+	 * const candidates = editor.spatialIndex.getShapeIdsAtPoint(point, margin)
+	 * const sorted = editor.getCurrentPageShapesSorted().filter(s => candidates.has(s.id))
+	 * ```
+	 *
+	 * @param point - The point to search at
+	 * @param margin - The margin around the point to search (default: 0)
+	 * @returns Unordered set of shape IDs that could potentially contain the point
+	 *
+	 * @public
+	 */
+	getShapeIdsAtPoint(point: { x: number; y: number }, margin = 0): Set<TLShapeId> {
+		this.spatialIndexComputed.get()
+		return this.rbush.search(new Box(point.x - margin, point.y - margin, margin * 2, margin * 2))
+	}
+
+	/**
+	 * Dispose of the spatial index manager.
+	 * Clears the R-tree to prevent memory leaks.
+	 *
+	 * @public
+	 */
+	dispose(): void {
+		this.rbush.dispose()
+		this.lastPageId = null
+	}
+}

package/src/lib/editor/managers/TickManager/TickManager.test.ts

@@ -21,6 +21,17 @@ describe('TickManager', () => {
 	let tickManager: TickManager
 	let mockEmit: Mock
 	let mockDisposablesAdd: Mock
+	let mockInputs: {
+		_currentScreenPoint: Vec
+		getCurrentScreenPoint(): Vec
+		currentScreenPoint: Vec
+		setCurrentScreenPoint(value: Vec): void
+		_pointerVelocity: Vec
+		getPointerVelocity(): Vec
+		pointerVelocity: Vec
+		setPointerVelocity(value: Vec): void
+		updatePointerVelocity(elapsed: number): void
+	}
 
 	beforeEach(() => {
 		vi.clearAllMocks()
@@ -41,16 +52,40 @@ describe('TickManager', () => {
 		mockEmit = vi.fn()
 		mockDisposablesAdd = vi.fn()
 
+		// Create a mock inputs object with getters and setters
+		mockInputs = {
+			_currentScreenPoint: new Vec(100, 100),
+			getCurrentScreenPoint() {
+				return this._currentScreenPoint
+			},
+			get currentScreenPoint() {
+				return this.getCurrentScreenPoint()
+			},
+			setCurrentScreenPoint(value: Vec) {
+				this._currentScreenPoint = value
+			},
+			_pointerVelocity: new Vec(0, 0),
+			getPointerVelocity() {
+				return this._pointerVelocity
+			},
+			get pointerVelocity() {
+				return this.getPointerVelocity()
+			},
+			setPointerVelocity(value: Vec) {
+				this._pointerVelocity = value
+			},
+			updatePointerVelocity(_elapsed: number) {
+				// Mock implementation - no-op for tests
+			},
+		}
+
 		editor = {
 			emit: mockEmit,
 			disposables: {
 				add: mockDisposablesAdd,
 			},
-			inputs: {
-				currentScreenPoint: new Vec(100, 100),
-				pointerVelocity: new Vec(0, 0),
-			},
-		} as any
+			inputs: mockInputs as unknown as Editor['inputs'],
+		} as unknown as Mocked<Editor>
 
 		tickManager = new TickManager(editor)
 	})
@@ -143,16 +178,6 @@ describe('TickManager', () => {
 			expect(mockEmit).toHaveBeenCalledTimes(2)
 		})
 
-		it('should update pointer velocity', () => {
-			const updatePointerVelocitySpy = vi.spyOn(tickManager as any, 'updatePointerVelocity')
-			tickManager.now = 1000
-			mockDateNow.mockReturnValue(1016)
-
-			tickManager.tick()
-
-			expect(updatePointerVelocitySpy).toHaveBeenCalledWith(16)
-		})
-
 		it('should schedule next RAF', () => {
 			tickManager.tick()
 			expect(mockRequestAnimationFrame).toHaveBeenCalled()
@@ -192,98 +217,6 @@ describe('TickManager', () => {
 		})
 	})
 
-	describe('updatePointerVelocity method', () => {
-		let updatePointerVelocity: any
-
-		beforeEach(() => {
-			// Access private method for testing
-			updatePointerVelocity = (tickManager as any).updatePointerVelocity.bind(tickManager)
-			// Reset the prevPoint to a known state
-			;(tickManager as any).prevPoint = new Vec(50, 50)
-		})
-
-		it('should return early if elapsed time is 0', () => {
-			const originalVelocity = editor.inputs.pointerVelocity.clone()
-
-			updatePointerVelocity(0)
-
-			expect(editor.inputs.pointerVelocity).toEqual(originalVelocity)
-		})
-
-		it('should calculate velocity based on pointer movement', () => {
-			editor.inputs.currentScreenPoint = new Vec(150, 150)
-			;(tickManager as any).prevPoint = new Vec(50, 50)
-			editor.inputs.pointerVelocity = new Vec(0, 0)
-
-			updatePointerVelocity(100) // 100ms elapsed
-
-			// Delta should be (100, 100), length = ~141.42, direction = (~0.707, ~0.707)
-			// Velocity should be length/elapsed = ~1.414 per ms in each direction
-			// But with linear interpolation (lrp factor 0.5), it will be halved
-			const expectedVelocity = new Vec(0.5, 0.5) // With lrp factor 0.5
-			expect(editor.inputs.pointerVelocity.x).toBeCloseTo(expectedVelocity.x, 1)
-			expect(editor.inputs.pointerVelocity.y).toBeCloseTo(expectedVelocity.y, 1)
-		})
-
-		it('should update prevPoint to current screen point', () => {
-			const newPoint = new Vec(200, 200)
-			editor.inputs.currentScreenPoint = newPoint
-
-			updatePointerVelocity(16)
-
-			expect((tickManager as any).prevPoint).toEqual(newPoint)
-		})
-
-		it('should use linear interpolation to smooth velocity', () => {
-			editor.inputs.currentScreenPoint = new Vec(150, 50)
-			;(tickManager as any).prevPoint = new Vec(50, 50)
-			editor.inputs.pointerVelocity = new Vec(2, 2)
-
-			updatePointerVelocity(100)
-
-			// Should interpolate between current velocity (2,2) and new velocity (1,0)
-			// lrp with factor 0.5 should give us something between them
-			expect(editor.inputs.pointerVelocity.x).toBeGreaterThan(0.5)
-			expect(editor.inputs.pointerVelocity.x).toBeLessThan(2)
-		})
-
-		it('should set very small velocity components to 0', () => {
-			editor.inputs.currentScreenPoint = new Vec(50.005, 50.005)
-			;(tickManager as any).prevPoint = new Vec(50, 50)
-			editor.inputs.pointerVelocity = new Vec(0, 0)
-
-			updatePointerVelocity(1000) // Long elapsed time = very small velocity
-
-			expect(editor.inputs.pointerVelocity.x).toBe(0)
-			expect(editor.inputs.pointerVelocity.y).toBe(0)
-		})
-
-		it('should handle zero movement (stationary pointer)', () => {
-			editor.inputs.currentScreenPoint = new Vec(100, 100)
-			;(tickManager as any).prevPoint = new Vec(100, 100)
-			editor.inputs.pointerVelocity = new Vec(1, 1)
-
-			updatePointerVelocity(16)
-
-			// Should interpolate towards zero velocity
-			expect(editor.inputs.pointerVelocity.x).toBeLessThan(1)
-			expect(editor.inputs.pointerVelocity.y).toBeLessThan(1)
-		})
-
-		it('should only update pointerVelocity if it actually changes', () => {
-			// Setup scenario where velocity won't change
-			editor.inputs.currentScreenPoint = new Vec(100, 100)
-			;(tickManager as any).prevPoint = new Vec(100, 100)
-			editor.inputs.pointerVelocity = new Vec(0, 0)
-
-			const originalVelocity = editor.inputs.pointerVelocity
-			updatePointerVelocity(16)
-
-			// Should still be the same object reference since no change occurred
-			expect(editor.inputs.pointerVelocity).toBe(originalVelocity)
-		})
-	})
-
 	describe('RAF throttling behavior', () => {
 		it('should use test-specific RAF behavior in test environment', () => {
 			// The TickManager should detect test environment and use requestAnimationFrame directly

package/src/lib/editor/managers/TickManager/TickManager.ts

@@ -1,5 +1,4 @@
 import { throttleToNextFrame as _throttleToNextFrame, bind } from '@tldraw/utils'
-import { Vec } from '../../../primitives/Vec'
 import { Editor } from '../../Editor'
 
 const throttleToNextFrame =
@@ -13,6 +12,7 @@ const throttleToNextFrame =
 			}
 		: _throttleToNextFrame
 
+/** @internal */
 export class TickManager {
 	constructor(public editor: Editor) {
 		this.editor.disposables.add(this.dispose)
@@ -40,7 +40,7 @@ export class TickManager {
 		const elapsed = now - this.now
 		this.now = now
 
-		this.updatePointerVelocity(elapsed)
+		this.editor.inputs.updatePointerVelocity(elapsed)
 		this.editor.emit('frame', elapsed)
 		this.editor.emit('tick', elapsed)
 		this.cancelRaf = throttleToNextFrame(this.tick)
@@ -53,34 +53,4 @@ export class TickManager {
 
 		this.cancelRaf?.()
 	}
-
-	private prevPoint = new Vec()
-
-	updatePointerVelocity(elapsed: number) {
-		const {
-			prevPoint,
-			editor: {
-				inputs: { currentScreenPoint, pointerVelocity },
-			},
-		} = this
-
-		if (elapsed === 0) return
-
-		const delta = Vec.Sub(currentScreenPoint, prevPoint)
-		this.prevPoint = currentScreenPoint.clone()
-
-		const length = delta.len()
-		const direction = length ? delta.div(length) : new Vec(0, 0)
-
-		// consider adjusting this with an easing rather than a linear interpolation
-		const next = pointerVelocity.clone().lrp(direction.mul(length / elapsed), 0.5)
-
-		// if the velocity is very small, just set it to 0
-		if (Math.abs(next.x) < 0.01) next.x = 0
-		if (Math.abs(next.y) < 0.01) next.y = 0
-
-		if (!pointerVelocity.equals(next)) {
-			this.editor.inputs.pointerVelocity = next
-		}
-	}
 }

package/src/lib/editor/shapes/BaseBoxShapeUtil.tsx

@@ -1,4 +1,4 @@
-import { TLBaseShape } from '@tldraw/tlschema'
+import { ExtractShapeByProps } from '@tldraw/tlschema'
 import { lerp } from '@tldraw/utils'
 import { Geometry2d } from '../../primitives/geometry/Geometry2d'
 import { Rectangle2d } from '../../primitives/geometry/Rectangle2d'
@@ -7,7 +7,7 @@ import { ShapeUtil, TLResizeInfo } from './ShapeUtil'
 import { resizeBox } from './shared/resizeBox'
 
 /** @public */
-export type TLBaseBoxShape = TLBaseShape<string, { w: number; h: number }>
+export type TLBaseBoxShape = ExtractShapeByProps<{ w: number; h: number }>
 
 /** @public */
 export abstract class BaseBoxShapeUtil<Shape extends TLBaseBoxShape> extends ShapeUtil<Shape> {