From 31a1a8b5ae89ade5d6de6ba3cc137870f73ae378 Mon Sep 17 00:00:00 2001
From: Steve Ruiz <steveruizok@gmail.com>
Date: Thu, 6 Jan 2022 13:36:40 +0000
Subject: [PATCH] Fix bug with missing draw shape after undo

---
 .../src/state/tools/DrawTool/DrawTool.spec.ts | 60 +++++++++++++++++++
 .../src/state/tools/DrawTool/DrawTool.ts      |  5 +-
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/packages/tldraw/src/state/tools/DrawTool/DrawTool.spec.ts b/packages/tldraw/src/state/tools/DrawTool/DrawTool.spec.ts
index 1ef3c9846..b9ed5039d 100644
--- a/packages/tldraw/src/state/tools/DrawTool/DrawTool.spec.ts
+++ b/packages/tldraw/src/state/tools/DrawTool/DrawTool.spec.ts
@@ -1,4 +1,6 @@
 import { TldrawApp } from '~state'
+import { TldrawTestApp } from '~test'
+import { TDShapeType } from '~types'
 import { DrawTool } from '.'
 
 describe('DrawTool', () => {
@@ -7,3 +9,61 @@ describe('DrawTool', () => {
     new DrawTool(app)
   })
 })
+
+describe('When shift+clicking to extend a shape', () => {
+  it('extends the same shape', () => {
+    const app = new TldrawTestApp()
+    app.reset()
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas([0, 0])
+    app.movePointer([100, 100])
+    app.movePointer([200, 200])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(1)
+    app.pointCanvas({ x: 300, y: 300, shiftKey: true })
+    app.movePointer([400, 400])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(1)
+  })
+
+  it('does not extend after switching tools the same shape', () => {
+    const app = new TldrawTestApp()
+    app.reset()
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas([0, 0])
+    app.movePointer([100, 100])
+    app.movePointer([200, 200])
+    app.stopPointing()
+    app.selectTool('select')
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas({ x: 300, y: 300, shiftKey: true })
+    app.movePointer([400, 400])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(2)
+  })
+
+  it('does not extend after undo', () => {
+    const app = new TldrawTestApp()
+    app.reset()
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas([0, 0])
+    app.movePointer([100, 100])
+    app.movePointer([200, 200])
+    app.stopPointing()
+    app.undo()
+    app.pointCanvas({ x: 300, y: 300, shiftKey: true })
+    app.movePointer([400, 400])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(1)
+  })
+
+  it('does not extend if no shape is present', () => {
+    const app = new TldrawTestApp()
+    app.reset()
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas({ x: 300, y: 300, shiftKey: true })
+    app.movePointer([400, 400])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(1)
+  })
+})
diff --git a/packages/tldraw/src/state/tools/DrawTool/DrawTool.ts b/packages/tldraw/src/state/tools/DrawTool/DrawTool.ts
index 5f6c33dd7..30af4d22f 100644
--- a/packages/tldraw/src/state/tools/DrawTool/DrawTool.ts
+++ b/packages/tldraw/src/state/tools/DrawTool/DrawTool.ts
@@ -42,9 +42,10 @@ export class DrawTool extends BaseTool {
       currentPoint,
       appState: { currentPageId, currentStyle },
     } = this.app
-    if (info.shiftKey && this.lastShapeId) {
+    const previous = this.lastShapeId && this.app.getShape(this.lastShapeId)
+    if (info.shiftKey && previous) {
       // Extend the previous shape
-      this.app.startSession(SessionType.Draw, this.lastShapeId)
+      this.app.startSession(SessionType.Draw, previous.id)
       this.setStatus(Status.Extending)
     } else {
       // Create a new shape