From f7ae99dd1fc906089834c96055b83ad5871eba21 Mon Sep 17 00:00:00 2001
From: David Sheldrick <d.j.sheldrick@gmail.com>
Date: Fri, 8 Dec 2023 10:35:35 +0000
Subject: [PATCH] zoom to affected shapes after undo/redo (#2293)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR makes it so that any shapes affected by an undo/redo action,
along with any shapes that are selected after an undo/redo action, are
visible in the viewport.

### Change Type

- [x] `patch` — Bug fix


### Release Notes

- Make sure affected shapes are visible after undo/redo
---
 packages/editor/src/lib/editor/Editor.ts | 32 ++++++++-
 packages/tldraw/src/test/Editor.test.tsx | 87 ++++++++++++++++++++++++
 2 files changed, 117 insertions(+), 2 deletions(-)

diff --git a/packages/editor/src/lib/editor/Editor.ts b/packages/editor/src/lib/editor/Editor.ts
index c5aa2753b..597d2bcbd 100644
--- a/packages/editor/src/lib/editor/Editor.ts
+++ b/packages/editor/src/lib/editor/Editor.ts
@@ -99,6 +99,7 @@ import {
 	sortByIndex,
 } from '../utils/reordering/reordering'
 import { applyRotationToSnapshotShapes, getRotationSnapshot } from '../utils/rotation'
+import { uniq } from '../utils/uniq'
 import { uniqueId } from '../utils/uniqueId'
 import { arrowBindingsIndex } from './derivations/arrowBindingsIndex'
 import { parentsToChildren } from './derivations/parentsToChildren'
@@ -781,6 +782,29 @@ export class Editor extends EventEmitter<TLEventMap> {
 		}
 	)
 
+	private _ensureAffectedShapesAreMadeVisible(fn: () => void) {
+		// here we collect all the ids of shapes that are created or updated by the function
+		// along with the selectedIds of the current page
+		// then we attempt to make sure that they are all visible in the viewport after
+		// the function is run.
+		const changes = this.store.extractingChanges(fn)
+		const affectedRecordIds = uniq(
+			Object.keys(changes.added)
+				.concat(Object.keys(changes.updated))
+				.concat(this.getSelectedShapeIds())
+		)
+		const shapes = compact(
+			affectedRecordIds.map((id) => (isShapeId(id) ? this.getShape(id) : null))
+		)
+		if (!shapes.length) return this
+		const bounds = Box2d.Common(compact(shapes.map((shape) => this.getShapePageBounds(shape))))
+		const viewport = this.getViewportPageBounds()
+		if (!viewport.contains(bounds)) {
+			this.zoomToBounds(bounds, this.getCamera().z, { duration: 220 })
+		}
+		return this
+	}
+
 	/**
 	 * Undo to the last mark.
 	 *
@@ -792,7 +816,9 @@ export class Editor extends EventEmitter<TLEventMap> {
 	 * @public
 	 */
 	undo(): this {
-		this.history.undo()
+		this._ensureAffectedShapesAreMadeVisible(() => {
+			this.history.undo()
+		})
 		return this
 	}
 
@@ -825,7 +851,9 @@ export class Editor extends EventEmitter<TLEventMap> {
 	 * @public
 	 */
 	redo(): this {
-		this.history.redo()
+		this._ensureAffectedShapesAreMadeVisible(() => {
+			this.history.redo()
+		})
 		return this
 	}
 
diff --git a/packages/tldraw/src/test/Editor.test.tsx b/packages/tldraw/src/test/Editor.test.tsx
index c138e0b3e..46d92ab27 100644
--- a/packages/tldraw/src/test/Editor.test.tsx
+++ b/packages/tldraw/src/test/Editor.test.tsx
@@ -1,6 +1,7 @@
 import {
 	AssetRecordType,
 	BaseBoxShapeUtil,
+	Box2d,
 	PageRecordType,
 	TLShape,
 	createShapeId,
@@ -645,3 +646,89 @@ describe('when the user prefers light UI', () => {
 		expect(editor.user.getIsDarkMode()).toBe(false)
 	})
 })
+
+describe('undo and redo', () => {
+	test('cause the camera to move if the affected shapes are offscreen', () => {
+		editor = new TestEditor({})
+		editor.setScreenBounds(new Box2d(0, 0, 1000, 1000))
+		editor.user.updateUserPreferences({ animationSpeed: 0 })
+
+		const boxId = createShapeId('box')
+		editor.createShapes([{ id: boxId, type: 'geo', x: 100, y: 100, props: { w: 100, h: 100 } }])
+		editor.panZoomIntoView([boxId])
+		editor.mark()
+		const cameraBefore = editor.getCamera()
+
+		editor.updateShapes([
+			{
+				id: boxId,
+				type: 'geo',
+				x: 100,
+				y: 100,
+				props: {
+					geo: 'cloud',
+					w: 100,
+					h: 100,
+				},
+			},
+		])
+
+		expect(editor.getCamera()).toMatchInlineSnapshot(`
+		Object {
+		  "id": "camera:page:page",
+		  "meta": Object {},
+		  "typeName": "camera",
+		  "x": 0,
+		  "y": 0,
+		  "z": 1,
+		}
+	`)
+
+		editor.undo()
+		expect(editor.getCamera()).toEqual(cameraBefore)
+
+		editor.updateShapes([
+			{
+				id: boxId,
+				type: 'geo',
+				x: -500,
+				y: -500,
+			},
+		])
+		editor.mark()
+		editor.updateShapes([
+			{
+				id: boxId,
+				type: 'geo',
+				x: 500,
+				y: 500,
+			},
+		])
+		editor.undo()
+
+		expect(editor.getCamera()).not.toEqual(cameraBefore)
+		expect(editor.getCamera()).toMatchInlineSnapshot(`
+		Object {
+		  "id": "camera:page:page",
+		  "meta": Object {},
+		  "typeName": "camera",
+		  "x": 950,
+		  "y": 950,
+		  "z": 1,
+		}
+	`)
+
+		editor.redo()
+
+		expect(editor.getCamera()).toMatchInlineSnapshot(`
+		Object {
+		  "id": "camera:page:page",
+		  "meta": Object {},
+		  "typeName": "camera",
+		  "x": -50,
+		  "y": -50,
+		  "z": 1,
+		}
+	`)
+	})
+})