[fix] Protect document from missing parents / children. (#622)

* Scan document for missing parents / children.

* fix tests

* Apply fixes to vscode extension

apps/vscode/editor/package.json

@@ -16,7 +16,7 @@
     "build:apps": "yarn build"
   },
   "devDependencies": {
-    "@tldraw/tldraw": "^1.6.1",
+    "@tldraw/tldraw": "*",
     "@types/node": "^17.0.14",
     "@types/react": "^17.0.38",
     "@types/react-dom": "^17.0.11",

apps/vscode/editor/src/app.tsx

@@ -10,6 +10,7 @@ import { exportToImage } from 'utils/export'
 declare let currentFile: TDFile
 
 export default function App(): JSX.Element {
+  const rLoaded = React.useRef(false)
   const rTldrawApp = React.useRef<TldrawApp>()
   const rInitialDocument = React.useRef<TDDocument>(
     currentFile ? currentFile.document : defaultDocument
@@ -39,7 +40,12 @@ export default function App(): JSX.Element {
         try {
           const { document } = JSON.parse(data.text) as TDFile
           const app = rTldrawApp.current!
-          app.updateDocument(document)
+          if (rLoaded.current) {
+            app.updateDocument(document)
+          } else {
+            app.loadDocument(document)
+            rLoaded.current = true
+          }
           app.zoomToFit()
         } catch (e) {
           console.warn('Failed to parse file:', data.text)

apps/vscode/extension/package.json

@@ -131,4 +131,4 @@
 		"vsce": "^2.2.0"
 	},
 	"gitHead": "4b1137849ad07da36fc8f0f19cb64e7535a79296"
-}
+}

package.json

@@ -12,6 +12,7 @@
   "workspaces": [
     "packages/*",
     "apps/*",
+    "apps/vscode/*",
     "examples/*"
   ],
   "scripts": {

packages/tldraw/package.json

@@ -94,8 +94,7 @@
     ],
     "testEnvironment": "jsdom",
     "modulePathIgnorePatterns": [
-      "<rootDir>/dist/",
-      "<rootDir>/src/test/"
+      "<rootDir>/dist/"
     ],
     "moduleNameMapper": {
       "@tldraw/tldraw": "<rootDir>/src",

packages/tldraw/src/state/TldrawApp.spec.ts

@@ -1,6 +1,7 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
 import { mockDocument, TldrawTestApp } from '~test'
 import { ArrowShape, ColorStyle, SessionType, TDShapeType } from '~types'
+import { deepCopy } from './StateManager/copy'
 import type { SelectTool } from './tools/SelectTool'
 
 describe('TldrawTestApp', () => {
@@ -659,38 +660,48 @@ describe('TldrawTestApp', () => {
       const app2 = new TldrawTestApp('migrate_1')
       await app2.ready
       expect(app2.getShape('rect1')).toBeTruthy()
+
       return
     })
   })
 
   describe('When replacing the page content', () => {
-    it('Should update the page with the correct shapes and bindings.', () => {
-      const shapes = mockDocument.pages.page1.shapes
-      const bindings = mockDocument.pages.page1.bindings
+    it('Should update the page with the correct shapes and bindings.', async () => {
+      const shapes = deepCopy(mockDocument.pages.page1.shapes)
+      const bindings = deepCopy(mockDocument.pages.page1.bindings)
       const app = new TldrawTestApp('multiplayer', {
         onChangePage: () => {
           //
         },
-      }).createPage()
+      })
+      await app.ready
+      app.createPage('page2')
+      expect(app.currentPageId).toBe('page2')
       app.replacePageContent(shapes, bindings, {})
-
-      expect(app.shapes).toEqual(Object.values(shapes))
-      expect(app.bindings).toEqual(Object.values(bindings))
+      expect(app.shapes).toMatchObject(
+        Object.values(shapes).map((s) => ({ ...s, parentId: 'page2' }))
+      )
+      expect(app.bindings).toMatchObject(Object.values(bindings))
     })
 
-    it('It should update the page shapes after the settings have been updated', () => {
-      const shapes = mockDocument.pages.page1.shapes
-      const bindings = mockDocument.pages.page1.bindings
+    it('It should update the page shapes after the settings have been updated', async () => {
+      const shapes = deepCopy(mockDocument.pages.page1.shapes)
+      const bindings = deepCopy(mockDocument.pages.page1.bindings)
       const app = new TldrawTestApp('multiplayer', {
         onChangePage: () => {
           //
         },
-      }).createPage()
+      })
+      await app.ready
+      app.createPage('page2')
+      expect(app.currentPageId).toBe('page2')
+      expect(shapes.rect1.parentId).toBe('page1')
       app.setSetting('isDebugMode', true)
       app.replacePageContent(shapes, bindings, {})
-
-      expect(app.shapes).toEqual(Object.values(shapes))
-      expect(app.bindings).toEqual(Object.values(bindings))
+      expect(app.shapes).toMatchObject(
+        Object.values(shapes).map((s) => ({ ...s, parentId: 'page2' }))
+      )
+      expect(app.bindings).toMatchObject(Object.values(bindings))
     })
   })
 
@@ -717,6 +728,7 @@ describe('When adding a video', () => {
 
 describe('When space panning', () => {
   it('pans camera when spacebar is down', () => {
+    // global.console.warn = jest.fn()
     const app = new TldrawTestApp()
     expect(app.pageState.camera.point).toMatchObject([0, 0])
     app.movePointer([0, 0])
@@ -731,6 +743,7 @@ describe('When space panning', () => {
     app.releaseKey(' ')
     app.stopPointing()
     expect(app.pageState.camera.point).toMatchObject([100, 100])
+    // expect(global.console.warn).not.toHaveBeenCalled()
   })
 
   it('pans camera in any state', () => {

packages/tldraw/src/state/TldrawApp.ts

@@ -684,6 +684,15 @@ export class TldrawApp extends StateManager<TDSnapshot> {
       return this
     }
 
+    const page = this.document.pages[this.currentPageId]
+
+    Object.values(shapes).forEach((shape) => {
+      if (shape.parentId !== pageId && !(page.shapes[shape.parentId] || shapes[shape.parentId])) {
+        console.warn('Added a shape without a parent on the page')
+        shape.parentId = pageId
+      }
+    })
+
     this.useStore.setState((current) => {
       const { hoveredId, editingId, bindingId, selectedIds } = current.document.pageStates[pageId]
 

packages/tldraw/src/state/data/migrate.ts

@@ -4,22 +4,43 @@ import { Decoration, FontStyle, TDDocument, TDShapeType, TextShape } from '~type
 export function migrate(document: TDDocument, newVersion: number): TDDocument {
   const { version = 0 } = document
 
+  if (!('assets' in document)) {
+    document.assets = {}
+  }
+
   // Remove unused assets when loading a document
-  if ('assets' in document) {
-    const assetIdsInUse = new Set<string>()
+  const assetIdsInUse = new Set<string>()
 
-    Object.values(document.pages).forEach((page) =>
-      Object.values(page.shapes).forEach((shape) => {
-        if (shape.assetId) assetIdsInUse.add(shape.assetId)
-      })
-    )
+  Object.values(document.pages).forEach((page) =>
+    Object.values(page.shapes).forEach((shape) => {
+      const { parentId, children, assetId } = shape
 
-    Object.keys(document.assets).forEach((assetId) => {
-      if (!assetIdsInUse.has(assetId)) {
-        delete document.assets[assetId]
+      if (assetId) {
+        assetIdsInUse.add(assetId)
+      }
+
+      // Fix missing parent bug
+      if (parentId !== page.id && !page.shapes[parentId]) {
+        console.warn('Encountered a shape with a missing parent!')
+        shape.parentId = page.id
+      }
+
+      if (children) {
+        children.forEach((childId) => {
+          if (!page.shapes[childId]) {
+            console.warn('Encountered a parent with a missing child!', shape.id, childId)
+            children?.splice(children.indexOf(childId), 1)
+          }
+        })
       }
     })
-  }
+  )
+
+  Object.keys(document.assets).forEach((assetId) => {
+    if (!assetIdsInUse.has(assetId)) {
+      delete document.assets[assetId]
+    }
+  })
 
   if (version === newVersion) return document
 

packages/tldraw/src/state/sessions/TransformSession/TransformSession.spec.ts

@@ -216,7 +216,7 @@ describe('Transform session', () => {
 })
 
 describe('When creating with a transform session', () => {
-  it.only('Deletes the shape on undo', () => {
+  it('Deletes the shape on undo', () => {
     const app = new TldrawTestApp()
       .selectTool(TDShapeType.Rectangle)
       .pointCanvas([0, 0])

packages/tldraw/src/test/badDocument.spec.ts

@@ -0,0 +1,27 @@
+/* eslint-disable @typescript-eslint/no-non-null-assertion */
+import { badDocument } from './documents/badDocument'
+import { TldrawTestApp } from './TldrawTestApp'
+
+describe('When loading a bad document', () => {
+  it('Fixes the document', () => {
+    const app = new TldrawTestApp()
+
+    global.console.warn = jest.fn()
+
+    // This doc has parents that are missing and children set to missing shapes
+    app.loadDocument(badDocument as any)
+
+    for (const pageId in app.document.pages) {
+      const page = app.document.pages[pageId]
+      for (const shape of Object.values(page.shapes)) {
+        if (shape.parentId === pageId) continue
+        expect(page.shapes[shape.parentId]).toBeDefined()
+      }
+    }
+
+    expect(app.getShape('rect2').parentId).toBe('page1')
+    expect(app.getShape('group1').children!.length).toBe(0)
+
+    expect(global.console.warn).toHaveBeenCalled()
+  })
+})

packages/tldraw/src/test/documents/badDocument.ts

@@ -0,0 +1,72 @@
+import { TDDocument, ColorStyle, DashStyle, SizeStyle, TDShapeType } from '~types'
+
+export const badDocument: TDDocument = {
+  version: 0,
+  id: 'doc',
+  name: 'New Document',
+  pages: {
+    page1: {
+      id: 'page1',
+      shapes: {
+        rect1: {
+          id: 'rect1',
+          parentId: 'page1',
+          name: 'Rectangle',
+          childIndex: 1,
+          type: TDShapeType.Rectangle,
+          point: [0, 0],
+          size: [100, 100],
+          style: {
+            dash: DashStyle.Draw,
+            size: SizeStyle.Medium,
+            color: ColorStyle.Blue,
+          },
+          label: '',
+        },
+        rect2: {
+          id: 'rect2',
+          parentId: 'MISSING_PARENT',
+          name: 'Rectangle',
+          childIndex: 2,
+          type: TDShapeType.Rectangle,
+          point: [100, 100],
+          size: [100, 100],
+          style: {
+            dash: DashStyle.Draw,
+            size: SizeStyle.Medium,
+            color: ColorStyle.Blue,
+          },
+          label: '',
+          labelPoint: [0.5, 0.5],
+        },
+        group1: {
+          id: 'group1',
+          parentId: 'page1',
+          name: 'Group',
+          childIndex: 3,
+          type: TDShapeType.Group,
+          point: [20, 20],
+          size: [100, 100],
+          style: {
+            dash: DashStyle.Draw,
+            size: SizeStyle.Medium,
+            color: ColorStyle.Blue,
+          },
+          children: ['MISSING_CHILD'],
+        },
+      },
+      bindings: {},
+    },
+  },
+  pageStates: {
+    page1: {
+      id: 'page1',
+      selectedIds: [],
+      camera: {
+        point: [0, 0],
+        zoom: 1,
+      },
+    },
+  },
+  assets: {},
+}