Use rbush to calculate shapes that are in view.

Fix.

Remove

packages/editor/api-report.md

@@ -23,6 +23,7 @@ import { JSX as JSX_2 } from 'react/jsx-runtime';
 import { Migrations } from '@tldraw/store';
 import { NamedExoticComponent } from 'react';
 import { PointerEventHandler } from 'react';
+import RBush from 'rbush';
 import { react } from '@tldraw/state';
 import { default as React_2 } from 'react';
 import * as React_3 from 'react';
@@ -900,6 +901,8 @@ export class Editor extends EventEmitter<TLEventMap> {
         speedThreshold?: number | undefined;
     }): this;
     readonly snaps: SnapManager;
+    // (undocumented)
+    _spatialIndex: SpatialIndex;
     stackShapes(shapes: TLShape[] | TLShapeId[], operation: 'horizontal' | 'vertical', gap: number): this;
     startFollowingUser(userId: string): this;
     stopCameraAnimation(): this;

packages/editor/api/api.json

@@ -7440,6 +7440,37 @@
           "name": "Editor",
           "preserveMemberOrder": false,
           "members": [
+            {
+              "kind": "Property",
+              "canonicalReference": "@tldraw/editor!Editor#_spatialIndex:member",
+              "docComment": "",
+              "excerptTokens": [
+                {
+                  "kind": "Content",
+                  "text": "_spatialIndex: "
+                },
+                {
+                  "kind": "Reference",
+                  "text": "SpatialIndex",
+                  "canonicalReference": "@tldraw/editor!~SpatialIndex:class"
+                },
+                {
+                  "kind": "Content",
+                  "text": ";"
+                }
+              ],
+              "isReadonly": false,
+              "isOptional": false,
+              "releaseTag": "Public",
+              "name": "_spatialIndex",
+              "propertyTypeTokenRange": {
+                "startIndex": 1,
+                "endIndex": 2
+              },
+              "isStatic": false,
+              "isProtected": false,
+              "isAbstract": false
+            },
             {
               "kind": "Constructor",
               "canonicalReference": "@tldraw/editor!Editor:constructor(1)",

packages/editor/package.json

@@ -59,7 +59,8 @@
 		"is-plain-object": "^5.0.0",
 		"lodash.throttle": "^4.1.1",
 		"lodash.uniq": "^4.5.0",
-		"nanoid": "4.0.2"
+		"nanoid": "4.0.2",
+		"rbush": "^3.0.1"
 	},
 	"peerDependencies": {
 		"react": "^18",
@@ -72,6 +73,7 @@
 		"@types/benchmark": "^2.1.2",
 		"@types/lodash.throttle": "^4.1.7",
 		"@types/lodash.uniq": "^4.5.7",
+		"@types/rbush": "3.0.3",
 		"@types/react-test-renderer": "^18.0.0",
 		"@types/wicg-file-system-access": "^2020.9.5",
 		"benchmark": "^2.1.4",

packages/editor/src/lib/editor/Editor.ts

@@ -111,6 +111,7 @@ import { HistoryManager } from './managers/HistoryManager'
 import { ScribbleManager } from './managers/ScribbleManager'
 import { SideEffectManager } from './managers/SideEffectManager'
 import { SnapManager } from './managers/SnapManager/SnapManager'
+import { SpatialIndex } from './managers/SpatialIndex'
 import { TextManager } from './managers/TextManager'
 import { TickManager } from './managers/TickManager'
 import { UserPreferencesManager } from './managers/UserPreferencesManager'
@@ -206,6 +207,8 @@ export class Editor extends EventEmitter<TLEventMap> {
 
 		this.getContainer = getContainer ?? (() => document.body)
 
+		this._spatialIndex = new SpatialIndex(this)
+
 		this.textMeasure = new TextManager(this)
 		this._tickManager = new TickManager(this)
 
@@ -3102,6 +3105,11 @@ export class Editor extends EventEmitter<TLEventMap> {
 		}
 	}
 
+	@computed
+	private getShapesInRenderingBoundsExpanded() {
+		return this._spatialIndex.getShapesInRenderingBoundsExpanded()
+	}
+
 	/** @internal */
 	getUnorderedRenderingShapes(
 		// The rendering state. We use this method both for rendering, which
@@ -3133,6 +3141,8 @@ export class Editor extends EventEmitter<TLEventMap> {
 
 		const erasingShapeIds = this.getErasingShapeIds()
 
+		const shapeInRenderingBoundsExpanded = new Set(this.getShapesInRenderingBoundsExpanded().get())
+
 		const addShapeById = (id: TLShapeId, opacity: number, isAncestorErasing: boolean) => {
 			const shape = this.getShape(id)
 			if (!shape) return
@@ -3197,7 +3207,10 @@ export class Editor extends EventEmitter<TLEventMap> {
 	 * @public
 	 */
 	@computed getRenderingShapes() {
+		let now = Date.now()
 		const renderingShapes = this.getUnorderedRenderingShapes(true)
+		console.log('unordered took', Date.now() - now, 'ms')
+		now = Date.now()
 
 		// Its IMPORTANT that the result be sorted by id AND include the index
 		// that the shape should be displayed at. Steve, this is the past you
@@ -3209,7 +3222,9 @@ export class Editor extends EventEmitter<TLEventMap> {
 		// drain. By always sorting by 'id' we keep the shapes always in the
 		// same order; but we later use index to set the element's 'z-index'
 		// to change the "rendered" position in z-space.
-		return renderingShapes.sort(sortById)
+		const sorted = renderingShapes.sort(sortById)
+		// console.log('sorting took', Date.now() - now, 'ms')
+		return sorted
 	}
 
 	/**
@@ -8859,6 +8874,8 @@ export class Editor extends EventEmitter<TLEventMap> {
 
 		return this
 	}
+
+	_spatialIndex: SpatialIndex
 }
 
 function alertMaxShapes(editor: Editor, pageId = editor.getCurrentPageId()) {

packages/editor/src/lib/editor/managers/SpatialIndex.ts

@@ -0,0 +1,124 @@
+import { RESET_VALUE, computed, isUninitialized } from '@tldraw/state'
+import { TLShape, TLShapeId, isShape, isShapeId } from '@tldraw/tlschema'
+import RBush from 'rbush'
+import { Box } from '../../primitives/Box'
+import { Editor } from '../Editor'
+
+type Element = {
+	minX: number
+	minY: number
+	maxX: number
+	maxY: number
+	id: TLShapeId
+}
+
+class TldrawRBush extends RBush<Element> {}
+
+export class SpatialIndex {
+	shapesInTree = new Map<TLShapeId, Element>()
+	rBush = new TldrawRBush()
+
+	constructor(private editor: Editor) {}
+
+	private getElement(shape: TLShape): Element | null {
+		const bounds = this.editor.getShapeMaskedPageBounds(shape)
+		if (!bounds) return null
+		return {
+			minX: bounds.x,
+			minY: bounds.y,
+			maxX: bounds.x + bounds.w,
+			maxY: bounds.y + bounds.h,
+			id: shape.id,
+		}
+	}
+
+	private addElementToArray(shape: TLShape, a: Element[]): Element | null {
+		const e = this.getElement(shape)
+		if (!e) return null
+		a.push(e)
+		return e
+	}
+
+	getShapesInRenderingBoundsExpanded() {
+		const { store } = this.editor
+		const shapeHistory = store.query.filterHistory('shape')
+
+		return computed<TLShapeId[]>('getShapesInView', (prevValue, lastComputedEpoch) => {
+			const renderingBoundsExpanded = this.editor.getRenderingBoundsExpanded()
+			const shapes = this.editor.getCurrentPageShapes()
+
+			if (isUninitialized(prevValue)) {
+				return this.fromScratch(renderingBoundsExpanded, shapes)
+			}
+			const diff = shapeHistory.getDiffSince(lastComputedEpoch)
+
+			if (diff === RESET_VALUE) {
+				return this.fromScratch(renderingBoundsExpanded, shapes)
+			}
+
+			const elementsToAdd: Element[] = []
+			for (const changes of diff) {
+				for (const record of Object.values(changes.added)) {
+					if (isShape(record)) {
+						const e = this.addElementToArray(record, elementsToAdd)
+						if (!e) continue
+						this.shapesInTree.set(record.id, e)
+					}
+				}
+
+				for (const [_from, to] of Object.values(changes.updated)) {
+					if (isShape(to)) {
+						const currentElement = this.shapesInTree.get(to.id)
+						if (currentElement) {
+							this.shapesInTree.delete(to.id)
+							this.rBush.remove(currentElement)
+						}
+						const newE = this.getElement(to)
+						if (!newE) continue
+						this.shapesInTree.set(to.id, newE)
+						elementsToAdd.push(newE)
+					}
+				}
+				this.rBush.load(elementsToAdd)
+
+				for (const id of Object.keys(changes.removed)) {
+					if (isShapeId(id)) {
+						const currentElement = this.shapesInTree.get(id)
+						if (currentElement) {
+							this.shapesInTree.delete(id)
+							this.rBush.remove(currentElement)
+						}
+					}
+				}
+			}
+			return this.searchTree(this.rBush, renderingBoundsExpanded)
+		})
+	}
+
+	private fromScratch(renderingBounds: Box, shapes: TLShape[]) {
+		this.rBush.clear()
+		this.shapesInTree = new Map<TLShapeId, Element>()
+		const elementsToAdd: Element[] = []
+
+		for (let i = 0; i < shapes.length; i++) {
+			const shape = shapes[i]
+			const e = this.addElementToArray(shape, elementsToAdd)
+			if (!e) continue
+			this.shapesInTree.set(shape.id, e)
+			elementsToAdd.push(e)
+		}
+		this.rBush.load(elementsToAdd)
+		return this.searchTree(this.rBush, renderingBounds)
+	}
+
+	private searchTree(tree: TldrawRBush, renderingBounds: Box): TLShapeId[] {
+		return tree
+			.search({
+				minX: renderingBounds.x,
+				minY: renderingBounds.y,
+				maxX: renderingBounds.x + renderingBounds.width,
+				maxY: renderingBounds.y + renderingBounds.height,
+			})
+			.map((b) => b.id)
+	}
+}

yarn.lock

@@ -7488,6 +7488,7 @@ __metadata:
     "@types/core-js": "npm:^2.5.5"
     "@types/lodash.throttle": "npm:^4.1.7"
     "@types/lodash.uniq": "npm:^4.5.7"
+    "@types/rbush": "npm:3.0.3"
     "@types/react-test-renderer": "npm:^18.0.0"
     "@types/wicg-file-system-access": "npm:^2020.9.5"
     "@use-gesture/react": "npm:^10.2.27"
@@ -7504,6 +7505,7 @@ __metadata:
     lodash.throttle: "npm:^4.1.1"
     lodash.uniq: "npm:^4.5.0"
     nanoid: "npm:4.0.2"
+    rbush: "npm:^3.0.1"
     react-test-renderer: "npm:^18.2.0"
     resize-observer-polyfill: "npm:^1.5.1"
   peerDependencies:
@@ -8330,6 +8332,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@types/rbush@npm:3.0.3":
+  version: 3.0.3
+  resolution: "@types/rbush@npm:3.0.3"
+  checksum: 59c75d20d3ebf95f8853a98f67d437adc047bf875df6e6bba90884fdfa8fa927402ccec762ecbc8724d98f9ed14c9e97d16eddb709a702021ce1874da5d0d8d7
+  languageName: node
+  linkType: hard
+
 "@types/react-dom@npm:^18.0.0, @types/react-dom@npm:^18.2.18":
   version: 18.2.18
   resolution: "@types/react-dom@npm:18.2.18"
@@ -21060,6 +21069,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"quickselect@npm:^2.0.0":
+  version: 2.0.0
+  resolution: "quickselect@npm:2.0.0"
+  checksum: ed2e78431050d223fb75da20ee98011aef1a03f7cb04e1a32ee893402e640be3cfb76d72e9dbe01edf3bb457ff6a62e5c2d85748424d1aa531f6ba50daef098c
+  languageName: node
+  linkType: hard
+
 "raf@npm:^3.4.1":
   version: 3.4.1
   resolution: "raf@npm:3.4.1"
@@ -21097,6 +21113,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"rbush@npm:^3.0.1":
+  version: 3.0.1
+  resolution: "rbush@npm:3.0.1"
+  dependencies:
+    quickselect: "npm:^2.0.0"
+  checksum: 489e2e7d9889888ad533518f194e3ab7cc19b1f1365a38ee99fbdda542a47f41cda7dc89870180050f4d04ea402e9ff294e1d767d03c0f1694e0028b7609eec9
+  languageName: node
+  linkType: hard
+
 "rc@npm:^1.2.7, rc@npm:^1.2.8, rc@npm:~1.2.7":
   version: 1.2.8
   resolution: "rc@npm:1.2.8"