From 7e61e448ab485f48d0a33ddc5c5979a8b5594368 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mitja=20Bezen=C5=A1ek?= <mitja.bezensek@gmail.com>
Date: Sat, 13 Apr 2024 23:04:19 +0200
Subject: [PATCH] Perf: Improve perf of `getCurrentPageShapesSorted` (#3453)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This significantly improves performance. Here's a comparison with 2k
shapes. Top is the new logic, bottom the old one.


![image](https://github.com/tldraw/tldraw/assets/2523721/e17b3733-dfd1-4aec-a427-31537bb9d159)

One place where this does make a significant difference is when you have
a lot of shapes on the page and you start [creating a new
arrow](https://github.com/orgs/tldraw/projects/40?pane=issue&itemId=59296136):

Before:

![image](https://github.com/tldraw/tldraw/assets/2523721/e4550197-c2be-480e-8f9a-090cebe1c8e4)

![image](https://github.com/tldraw/tldraw/assets/2523721/7559fe14-ad08-4ee0-9c9e-de0b60d401b2)


After:

![image](https://github.com/tldraw/tldraw/assets/2523721/4c6a1df6-732f-48b4-a7ea-6ce0894cf46e)

![image](https://github.com/tldraw/tldraw/assets/2523721/1cd5f2aa-919c-4271-af9a-227e8babf458)


### Change Type

<!-- ❗ Please select a 'Scope' label ❗️ -->

- [ ] `sdk` — Changes the tldraw SDK
- [ ] `dotcom` — Changes the tldraw.com web app
- [ ] `docs` — Changes to the documentation, examples, or templates.
- [ ] `vs code` — Changes to the vscode plugin
- [x] `internal` — Does not affect user-facing stuff

<!-- ❗ Please select a 'Type' label ❗️ -->

- [ ] `bugfix` — Bug fix
- [ ] `feature` — New feature
- [x] `improvement` — Improving existing features
- [ ] `chore` — Updating dependencies, other boring stuff
- [ ] `galaxy brain` — Architectural changes
- [ ] `tests` — Changes to any test code
- [ ] `tools` — Changes to infrastructure, CI, internal scripts,
debugging tools, etc.
- [ ] `dunno` — I don't know

---------

Co-authored-by: Steve Ruiz <steveruizok@gmail.com>
---
 packages/editor/src/lib/editor/Editor.ts | 55 +++++++++++++++---------
 1 file changed, 34 insertions(+), 21 deletions(-)

diff --git a/packages/editor/src/lib/editor/Editor.ts b/packages/editor/src/lib/editor/Editor.ts
index 0f28bcbaa..24813f927 100644
--- a/packages/editor/src/lib/editor/Editor.ts
+++ b/packages/editor/src/lib/editor/Editor.ts
@@ -36,7 +36,6 @@ import {
 	createShapeId,
 	getShapePropKeysByStyle,
 	isPageId,
-	isShape,
 	isShapeId,
 } from '@tldraw/tlschema'
 import {
@@ -4575,31 +4574,31 @@ export class Editor extends EventEmitter<TLEventMap> {
 	 * @public
 	 */
 	@computed getCurrentPageShapesSorted(): TLShape[] {
-		// todo: consider making into a function call that includes options for selected-only, rendering, etc.
-		// todo: consider making a derivation or something, or merging with rendering shapes
-		const shapes = new Set(this.getCurrentPageShapes().sort(sortByIndex))
+		const shapes = this.getCurrentPageShapes().sort(sortByIndex)
+		const parentChildMap = new Map<TLShapeId, TLShape[]>()
+		const result: TLShape[] = []
+		const topLevelShapes: TLShape[] = []
+		let shape: TLShape, parent: TLShape | undefined
 
-		const results: TLShape[] = []
-
-		function pushShapeWithDescendants(shape: TLShape): void {
-			results.push(shape)
-			shapes.delete(shape)
-
-			shapes.forEach((otherShape) => {
-				if (otherShape.parentId === shape.id) {
-					pushShapeWithDescendants(otherShape)
+		for (let i = 0, n = shapes.length; i < n; i++) {
+			shape = shapes[i]
+			parent = this.getShape(shape.parentId)
+			if (parent) {
+				if (!parentChildMap.has(parent.id)) {
+					parentChildMap.set(parent.id, [])
 				}
-			})
+				parentChildMap.get(parent.id)!.push(shape)
+			} else {
+				// undefined if parent is a shape
+				topLevelShapes.push(shape)
+			}
 		}
 
-		shapes.forEach((shape) => {
-			const parent = this.getShape(shape.parentId)
-			if (!isShape(parent)) {
-				pushShapeWithDescendants(shape)
-			}
-		})
+		for (let i = 0, n = topLevelShapes.length; i < n; i++) {
+			pushShapeWithDescendants(topLevelShapes[i], parentChildMap, result)
+		}
 
-		return results
+		return result
 	}
 
 	/**
@@ -8884,3 +8883,17 @@ function applyPartialToShape<T extends TLShape>(prev: T, partial?: TLShapePartia
 	if (!next) return prev
 	return next
 }
+
+function pushShapeWithDescendants(
+	shape: TLShape,
+	parentChildMap: Map<TLShapeId, TLShape[]>,
+	result: TLShape[]
+): void {
+	result.push(shape)
+	const children = parentChildMap.get(shape.id)
+	if (children) {
+		for (let i = 0, n = children.length; i < n; i++) {
+			pushShapeWithDescendants(children[i], parentChildMap, result)
+		}
+	}
+}