[fix] account for "draft" shapes when preserving selection state during replacePageContent (#427)

* account for "virtual" shapes when preserving appState

* rewrite merge logic

* More work on multiplayer

* Update TldrawApp.ts

* Improve logic around when to replace page content

Co-authored-by: Steve Ruiz <steveruizok@gmail.com>

examples/tldraw-example/src/multiplayer/useMultiplayerState.ts

@@ -4,6 +4,7 @@ import * as React from 'react'
 import type { TldrawApp, TDUser, TDShape, TDBinding, TDDocument } from '@tldraw/tldraw'
 import { useRedo, useUndo, useRoom, useUpdateMyPresence } from '@liveblocks/react'
 import { LiveMap, LiveObject } from '@liveblocks/client'
+import { Utils } from '@tldraw/core'
 
 declare const window: Window & { app: TldrawApp }
 
@@ -11,144 +12,15 @@ export function useMultiplayerState(roomId: string) {
   const [app, setApp] = React.useState<TldrawApp>()
   const [error, setError] = React.useState<Error>()
   const [loading, setLoading] = React.useState(true)
-  const rExpectingUpdate = React.useRef(false)
 
   const room = useRoom()
   const onUndo = useUndo()
   const onRedo = useRedo()
   const updateMyPresence = useUpdateMyPresence()
 
-  // Document Changes --------
-
   const rLiveShapes = React.useRef<LiveMap<string, TDShape>>()
   const rLiveBindings = React.useRef<LiveMap<string, TDBinding>>()
 
-  React.useEffect(() => {
-    const unsubs: (() => void)[] = []
-
-    if (!(app && room)) return
-    // Handle errors
-    unsubs.push(room.subscribe('error', (error) => setError(error)))
-
-    // Handle changes to other users' presence
-    unsubs.push(
-      room.subscribe('others', (others) => {
-        app.updateUsers(
-          others
-            .toArray()
-            .filter((other) => other.presence)
-            .map((other) => other.presence!.user)
-            .filter(Boolean)
-        )
-      })
-    )
-
-    // Handle events from the room
-    unsubs.push(
-      room.subscribe(
-        'event',
-        (e: { connectionId: number; event: { name: string; userId: string } }) => {
-          switch (e.event.name) {
-            case 'exit': {
-              app?.removeUser(e.event.userId)
-              break
-            }
-          }
-        }
-      )
-    )
-
-    // Send the exit event when the tab closes
-    function handleExit() {
-      if (!(room && app?.room)) return
-      room?.broadcastEvent({ name: 'exit', userId: app.room.userId })
-    }
-
-    window.addEventListener('beforeunload', handleExit)
-    unsubs.push(() => window.removeEventListener('beforeunload', handleExit))
-
-    // Setup the document's storage and subscriptions
-    async function setupDocument() {
-      const storage = await room.getStorage<any>()
-
-      // Initialize (get or create) shapes and bindings maps
-
-      let lShapes: LiveMap<string, TDShape> = storage.root.get('shapes')
-      if (!lShapes) {
-        storage.root.set('shapes', new LiveMap<string, TDShape>())
-        lShapes = storage.root.get('shapes')
-      }
-      rLiveShapes.current = lShapes
-
-      let lBindings: LiveMap<string, TDBinding> = storage.root.get('bindings')
-      if (!lBindings) {
-        storage.root.set('bindings', new LiveMap<string, TDBinding>())
-        lBindings = storage.root.get('bindings')
-      }
-      rLiveBindings.current = lBindings
-
-      // Subscribe to changes
-      function handleChanges() {
-        if (rExpectingUpdate.current) {
-          rExpectingUpdate.current = false
-          return
-        }
-
-        app?.replacePageContent(
-          Object.fromEntries(lShapes.entries()),
-          Object.fromEntries(lBindings.entries())
-        )
-      }
-
-      unsubs.push(room.subscribe(lShapes, handleChanges))
-      unsubs.push(room.subscribe(lBindings, handleChanges))
-
-      // Update the document with initial content
-      handleChanges()
-
-      // Migrate previous versions
-      const version = storage.root.get('version')
-
-      if (!version) {
-        // The doc object will only be present if the document was created
-        // prior to the current multiplayer implementation. At this time, the
-        // document was a single LiveObject named 'doc'. If we find a doc,
-        // then we need to move the shapes and bindings over to the new structures
-        // and then mark the doc as migrated.
-        const doc = storage.root.get('doc') as LiveObject<{
-          uuid: string
-          document: TDDocument
-          migrated?: boolean
-        }>
-
-        // No doc? No problem. This was likely a newer document
-        if (doc) {
-          const {
-            document: {
-              pages: {
-                page: { shapes, bindings },
-              },
-            },
-          } = doc.toObject()
-
-          Object.values(shapes).forEach((shape) => lShapes.set(shape.id, shape))
-          Object.values(bindings).forEach((binding) => lBindings.set(binding.id, binding))
-        }
-      }
-
-      // Save the version number for future migrations
-      storage.root.set('version', 2)
-
-      setLoading(false)
-    }
-
-    setupDocument()
-
-    return () => {
-      unsubs.forEach((unsub) => unsub())
-    }
-  }, [room, app])
-
   // Callbacks --------------
 
   // Put the state into the window, for debugging.
@@ -190,8 +62,6 @@ export function useMultiplayerState(roomId: string) {
             lBindings.set(binding.id, binding)
           }
         })
-
-        rExpectingUpdate.current = true
       })
     },
     [room]
@@ -205,6 +75,132 @@ export function useMultiplayerState(roomId: string) {
     [updateMyPresence]
   )
 
+  // Document Changes --------
+
+  React.useEffect(() => {
+    const unsubs: (() => void)[] = []
+    if (!(app && room)) return
+    // Handle errors
+    unsubs.push(room.subscribe('error', (error) => setError(error)))
+
+    // Handle changes to other users' presence
+    unsubs.push(
+      room.subscribe('others', (others) => {
+        app.updateUsers(
+          others
+            .toArray()
+            .filter((other) => other.presence)
+            .map((other) => other.presence!.user)
+            .filter(Boolean)
+        )
+      })
+    )
+
+    // Handle events from the room
+    unsubs.push(
+      room.subscribe(
+        'event',
+        (e: { connectionId: number; event: { name: string; userId: string } }) => {
+          switch (e.event.name) {
+            case 'exit': {
+              app?.removeUser(e.event.userId)
+              break
+            }
+          }
+        }
+      )
+    )
+
+    // Send the exit event when the tab closes
+    function handleExit() {
+      if (!(room && app?.room)) return
+      room?.broadcastEvent({ name: 'exit', userId: app.room.userId })
+    }
+
+    window.addEventListener('beforeunload', handleExit)
+    unsubs.push(() => window.removeEventListener('beforeunload', handleExit))
+
+    let stillAlive = true
+
+    // Setup the document's storage and subscriptions
+    async function setupDocument() {
+      const storage = await room.getStorage<any>()
+
+      // Initialize (get or create) shapes and bindings maps
+
+      let lShapes: LiveMap<string, TDShape> = storage.root.get('shapes')
+      if (!lShapes) {
+        storage.root.set('shapes', new LiveMap<string, TDShape>())
+        lShapes = storage.root.get('shapes')
+      }
+      rLiveShapes.current = lShapes
+
+      let lBindings: LiveMap<string, TDBinding> = storage.root.get('bindings')
+      if (!lBindings) {
+        storage.root.set('bindings', new LiveMap<string, TDBinding>())
+        lBindings = storage.root.get('bindings')
+      }
+      rLiveBindings.current = lBindings
+
+      // Migrate previous versions
+      const version = storage.root.get('version')
+
+      if (!version) {
+        // The doc object will only be present if the document was created
+        // prior to the current multiplayer implementation. At this time, the
+        // document was a single LiveObject named 'doc'. If we find a doc,
+        // then we need to move the shapes and bindings over to the new structures
+        // and then mark the doc as migrated.
+        const doc = storage.root.get('doc') as LiveObject<{
+          uuid: string
+          document: TDDocument
+          migrated?: boolean
+        }>
+
+        // No doc? No problem. This was likely a newer document
+        if (doc) {
+          const {
+            document: {
+              pages: {
+                page: { shapes, bindings },
+              },
+            },
+          } = doc.toObject()
+
+          Object.values(shapes).forEach((shape) => lShapes.set(shape.id, shape))
+          Object.values(bindings).forEach((binding) => lBindings.set(binding.id, binding))
+        }
+      }
+
+      // Save the version number for future migrations
+      storage.root.set('version', 2)
+
+      // Subscribe to changes
+      const handleChanges = () => {
+        app?.replacePageContent(
+          Object.fromEntries(lShapes.entries()),
+          Object.fromEntries(lBindings.entries())
+        )
+      }
+
+      if (stillAlive) {
+        unsubs.push(room.subscribe(lShapes, handleChanges))
+
+        // Update the document with initial content
+        handleChanges()
+
+        setLoading(false)
+      }
+    }
+
+    setupDocument()
+
+    return () => {
+      stillAlive = false
+      unsubs.forEach((unsub) => unsub())
+    }
+  }, [app])
+
   return {
     onUndo,
     onRedo,

packages/tldraw/src/state/TldrawApp.ts

@@ -469,7 +469,10 @@ export class TldrawApp extends StateManager<TDSnapshot> {
   }
 
   onPersist = () => {
-    this.broadcastPageChanges()
+    // If we are part of a room, send our changes to the server
+    if (this.callbacks.onChangePage) {
+      this.broadcastPageChanges()
+    }
   }
 
   private prevSelectedIds = this.selectedIds
@@ -493,6 +496,7 @@ export class TldrawApp extends StateManager<TDSnapshot> {
 
   /* ----------- Managing Multiplayer State ----------- */
 
+  private justSent = false
   private prevShapes = this.page.shapes
   private prevBindings = this.page.bindings
 
@@ -502,6 +506,26 @@ export class TldrawApp extends StateManager<TDSnapshot> {
     const changedShapes: Record<string, TDShape | undefined> = {}
     const changedBindings: Record<string, TDBinding | undefined> = {}
 
+    // const visitedIds = new Set<string>()
+    // const shapesToVisit = this.shapes
+
+    // while (shapesToVisit.length > 0) {
+    //   const shape = shapesToVisit.pop()
+    //   if (!shape) break
+    //   visitedIds.add(shape.id)
+    //   if (this.prevShapes[shape.id] !== shape) {
+    //     changedShapes[shape.id] = shape
+
+    //     if (shape.parentId !== this.currentPageId) {
+    //       shapesToVisit.push(this.page.shapes[shape.parentId])
+    //     }
+
+    //     if (shape.children) {
+
+    //     }
+    //   }
+    // }
+
     this.shapes.forEach((shape) => {
       visited.add(shape.id)
       if (this.prevShapes[shape.id] !== shape) {
@@ -512,6 +536,8 @@ export class TldrawApp extends StateManager<TDSnapshot> {
     Object.keys(this.prevShapes)
       .filter((id) => !visited.has(id))
       .forEach((id) => {
+        // After visiting all the current shapes, if we haven't visited a
+        // previously present shape, then it was deleted
         changedShapes[id] = undefined
       })
 
@@ -525,16 +551,73 @@ export class TldrawApp extends StateManager<TDSnapshot> {
     Object.keys(this.prevBindings)
       .filter((id) => !visited.has(id))
       .forEach((id) => {
+        // After visiting all the current bindings, if we haven't visited a
+        // previously present shape, then it was deleted
         changedBindings[id] = undefined
       })
 
+    this.justSent = true
     this.callbacks.onChangePage?.(this, changedShapes, changedBindings)
-
     this.callbacks.onPersist?.(this)
     this.prevShapes = this.page.shapes
     this.prevBindings = this.page.bindings
   }
 
+  getReservedContent = (ids: string[], pageId = this.currentPageId) => {
+    const { bindings } = this.document.pages[pageId]
+
+    // We want to know which shapes we need to
+    const reservedShapes: Record<string, TDShape> = {}
+    const reservedBindings: Record<string, TDBinding> = {}
+
+    // Quick lookup maps for bindings
+    const bindingsArr = Object.values(bindings)
+    const boundTos = new Map(bindingsArr.map((binding) => [binding.toId, binding]))
+    const boundFroms = new Map(bindingsArr.map((binding) => [binding.fromId, binding]))
+    const bindingMaps = [boundTos, boundFroms]
+
+    // Unique set of shape ids that are going to be reserved
+    const reservedShapeIds: string[] = []
+
+    if (this.session) ids.forEach((id) => reservedShapeIds.push(id))
+
+    const strongReservedShapeIds = new Set(reservedShapeIds)
+
+    // Which shape ids have we already visited?
+    const visited = new Set<string>()
+
+    // Time to visit every reserved shape and every related shape and binding.
+    while (reservedShapeIds.length > 0) {
+      const id = reservedShapeIds.pop()
+      if (!id) break
+      if (visited.has(id)) continue
+
+      // Add to set so that we don't process this id a second time
+      visited.add(id)
+
+      // Get the shape and reserve it
+      const shape = this.getShape(id)
+      reservedShapes[id] = shape
+
+      if (shape.parentId !== pageId) reservedShapeIds.push(shape.parentId)
+
+      // If the shape has children, add the shape's children to the list of ids to process
+      if (shape.children) reservedShapeIds.push(...shape.children)
+
+      // If there are binding for this shape, reserve the bindings and
+      // add its related shapes to the list of ids to process
+      bindingMaps
+        .map((map) => map.get(shape.id)!)
+        .filter(Boolean)
+        .forEach((binding) => {
+          reservedBindings[binding.id] = binding
+          reservedShapeIds.push(binding.toId, binding.fromId)
+        })
+    }
+
+    return { reservedShapes, reservedBindings, strongReservedShapeIds }
+  }
+
   /**
    * Manually patch a set of shapes.
    * @param shapes An array of shape partials, containing the changes to be made to each shape.
@@ -545,19 +628,75 @@ export class TldrawApp extends StateManager<TDSnapshot> {
     bindings: Record<string, TDBinding>,
     pageId = this.currentPageId
   ): this => {
+    // This will be called a few times: once by our own change,
+    // once by the change to shapes, and once by the change to bindings
+
+    if (this.justSent) {
+      this.justSent = false
+      return this
+    }
+
     this.useStore.setState((current) => {
       const { hoveredId, editingId, bindingId, selectedIds } = current.document.pageStates[pageId]
 
-      const keepShapes: Record<string, TDShape> = {}
+      const coreReservedIds = [...selectedIds]
-      const keepBindings: Record<string, TDBinding> = {}
 
-      if (this.session) {
+      if (editingId) coreReservedIds.push(editingId)
-        selectedIds.forEach((id) => (keepShapes[id] = this.getShape(id)))
+
-        Object.assign(keepBindings, this.bindings) // ROUGH
+      const { reservedShapes, reservedBindings, strongReservedShapeIds } = this.getReservedContent(
+        coreReservedIds,
+        this.currentPageId
+      )
+
+      // Merge in certain changes to reserved shapes
+      Object.values(reservedShapes)
+        // Don't merge updates to shapes with text (Text or Sticky)
+        .filter((reservedShape) => !('text' in reservedShape))
+        .forEach((reservedShape) => {
+          const incomingShape = shapes[reservedShape.id]
+          if (!incomingShape) return
+
+          // If the shape isn't "strongly reserved", then use the incoming shape;
+          // note that this is only if the incoming shape exists! If the shape was
+          // deleted in the incoming shapes, then we'll keep out reserved shape.
+          // This logic would need more work for arrows, because the incoming shape
+          // include a binding change that we'll need to resolve with our reserved bindings.
+          if (
+            !(
+              reservedShape.type === TDShapeType.Arrow ||
+              strongReservedShapeIds.has(reservedShape.id)
+            )
+          ) {
+            reservedShapes[reservedShape.id] = incomingShape
+            return
+          }
+
+          // Only allow certain merges.
+
+          // Allow decorations (of an arrow) to be changed
+          if ('decorations' in incomingShape && 'decorations' in reservedShape) {
+            reservedShape.decorations = incomingShape.decorations
+          }
+
+          // Allow the shape's style to be changed
+          reservedShape.style = incomingShape.style
+        })
+
+      // Use the incoming shapes / bindings as comparisons for what
+      // will have changed. This is important because we want to restore
+      // related shapes that may not have changed on our side, but which
+      // were deleted on the server.
+      this.prevShapes = shapes
+      this.prevBindings = bindings
+
+      const nextShapes = {
+        ...shapes,
+        ...reservedShapes,
       }
 
-      if (editingId) {
+      const nextBindings = {
-        keepShapes[editingId] = this.getShape(editingId)
+        ...bindings,
+        ...reservedBindings,
       }
 
       const next: TDSnapshot = {
@@ -567,29 +706,23 @@ export class TldrawApp extends StateManager<TDSnapshot> {
           pages: {
             [pageId]: {
               ...current.document.pages[pageId],
-              shapes: {
+              shapes: nextShapes,
-                ...shapes,
+              bindings: nextBindings,
-                ...keepShapes,
-              },
-              bindings: {
-                ...bindings,
-                ...keepBindings,
-              },
             },
           },
           pageStates: {
             ...current.document.pageStates,
             [pageId]: {
               ...current.document.pageStates[pageId],
-              selectedIds: selectedIds.filter((id) => shapes[id] !== undefined),
+              selectedIds: selectedIds.filter((id) => nextShapes[id] !== undefined),
               hoveredId: hoveredId
-                ? shapes[hoveredId] === undefined
+                ? nextShapes[hoveredId] === undefined
                   ? undefined
                   : hoveredId
                 : undefined,
               editingId: editingId,
               bindingId: bindingId
-                ? bindings[bindingId] === undefined
+                ? nextBindings[bindingId] === undefined
                   ? undefined
                   : bindingId
                 : undefined,
@@ -598,9 +731,66 @@ export class TldrawApp extends StateManager<TDSnapshot> {
         },
       }
 
+      // Get bindings related to the changed shapes
+      const bindingsToUpdate = TLDR.getRelatedBindings(next, Object.keys(nextShapes), pageId)
+
+      const page = next.document.pages[pageId]
+
+      // Update all of the bindings we've just collected
+      bindingsToUpdate.forEach((binding) => {
+        if (!page.bindings[binding.id]) {
+          return
+        }
+
+        const toShape = page.shapes[binding.toId]
+        const fromShape = page.shapes[binding.fromId]
+
+        const toUtils = TLDR.getShapeUtil(toShape)
+
+        const fromUtils = TLDR.getShapeUtil(fromShape)
+
+        // We only need to update the binding's "from" shape
+        const fromDelta = fromUtils.onBindingChange?.(
+          fromShape,
+          binding,
+          toShape,
+          toUtils.getBounds(toShape),
+          toUtils.getCenter(toShape)
+        )
+
+        if (fromDelta) {
+          const nextShape = {
+            ...fromShape,
+            ...fromDelta,
+          } as TDShape
+
+          page.shapes[fromShape.id] = nextShape
+        }
+      })
+
+      Object.values(nextShapes).forEach((shape) => {
+        if (shape.type !== TDShapeType.Group) return
+
+        const children = shape.children.filter((id) => page.shapes[id] !== undefined)
+
+        const commonBounds = Utils.getCommonBounds(
+          children
+            .map((id) => page.shapes[id])
+            .filter(Boolean)
+            .map((shape) => TLDR.getRotatedBounds(shape))
+        )
+
+        page.shapes[shape.id] = {
+          ...shape,
+          point: [commonBounds.minX, commonBounds.minY],
+          size: [commonBounds.width, commonBounds.height],
+          children,
+        }
+      })
+
       this.state.document = next.document
-      this.prevShapes = next.document.pages[this.currentPageId].shapes
+      // this.prevShapes = nextShapes
-      this.prevBindings = next.document.pages[this.currentPageId].bindings
+      // this.prevBindings = nextBindings
 
       return next
     }, true)