fix: rewrite collab element reconciliation to fix z-index issues (#4076)

src/excalidraw-app/collab/CollabWrapper.tsx

@@ -8,10 +8,7 @@ import {
   ExcalidrawElement,
   InitializedExcalidrawImageElement,
 } from "../../element/types";
-import {
-  getElementMap,
-  getSceneVersion,
-} from "../../packages/excalidraw/index";
+import { getSceneVersion } from "../../packages/excalidraw/index";
 import { Collaborator, Gesture } from "../../types";
 import {
   preventUnload,
@@ -64,6 +61,10 @@ import {
   isInitializedImageElement,
 } from "../../element/typeChecks";
 import { mutateElement } from "../../element/mutateElement";
+import {
+  ReconciledElements,
+  reconcileElements as _reconcileElements,
+} from "./reconciliation";
 
 interface CollabState {
   modalIsShown: boolean;
@@ -87,10 +88,6 @@ export interface CollabAPI {
   fetchImageFilesFromFirebase: CollabInstance["fetchImageFilesFromFirebase"];
 }
 
-type ReconciledElements = readonly ExcalidrawElement[] & {
-  _brand: "reconciledElements";
-};
-
 interface Props {
   excalidrawAPI: ExcalidrawImperativeAPI;
   onRoomClose?: () => void;
@@ -227,7 +224,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
   });
 
   saveCollabRoomToFirebase = async (
-    syncableElements: ExcalidrawElement[] = this.getSyncableElements(
+    syncableElements: readonly ExcalidrawElement[] = this.getSyncableElements(
       this.excalidrawAPI.getSceneElementsIncludingDeleted(),
     ),
   ) => {
@@ -484,65 +481,26 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
   };
 
   private reconcileElements = (
-    elements: readonly ExcalidrawElement[],
+    remoteElements: readonly ExcalidrawElement[],
   ): ReconciledElements => {
-    const currentElements = this.getSceneElementsIncludingDeleted();
-    // create a map of ids so we don't have to iterate
-    // over the array more than once.
-    const localElementMap = getElementMap(currentElements);
-
+    const localElements = this.getSceneElementsIncludingDeleted();
     const appState = this.excalidrawAPI.getAppState();
 
-    // Reconcile
-    const newElements: readonly ExcalidrawElement[] = elements
-      .reduce((elements, element) => {
-        // if the remote element references one that's currently
-        // edited on local, skip it (it'll be added in the next step)
-        if (
-          element.id === appState.editingElement?.id ||
-          element.id === appState.resizingElement?.id ||
-          element.id === appState.draggingElement?.id
-        ) {
-          return elements;
-        }
-
-        if (
-          localElementMap.hasOwnProperty(element.id) &&
-          localElementMap[element.id].version > element.version
-        ) {
-          elements.push(localElementMap[element.id]);
-          delete localElementMap[element.id];
-        } else if (
-          localElementMap.hasOwnProperty(element.id) &&
-          localElementMap[element.id].version === element.version &&
-          localElementMap[element.id].versionNonce !== element.versionNonce
-        ) {
-          // resolve conflicting edits deterministically by taking the one with the lowest versionNonce
-          if (localElementMap[element.id].versionNonce < element.versionNonce) {
-            elements.push(localElementMap[element.id]);
-          } else {
-            // it should be highly unlikely that the two versionNonces are the same. if we are
-            // really worried about this, we can replace the versionNonce with the socket id.
-            elements.push(element);
-          }
-          delete localElementMap[element.id];
-        } else {
-          elements.push(element);
-          delete localElementMap[element.id];
-        }
-
-        return elements;
-      }, [] as Mutable<typeof elements>)
-      // add local elements that weren't deleted or on remote
-      .concat(...Object.values(localElementMap));
+    const reconciledElements = _reconcileElements(
+      localElements,
+      remoteElements,
+      appState,
+    );
 
     // Avoid broadcasting to the rest of the collaborators the scene
     // we just received!
     // Note: this needs to be set before updating the scene as it
     // synchronously calls render.
-    this.setLastBroadcastedOrReceivedSceneVersion(getSceneVersion(newElements));
+    this.setLastBroadcastedOrReceivedSceneVersion(
+      getSceneVersion(reconciledElements),
+    );
 
-    return newElements as ReconciledElements;
+    return reconciledElements;
   };
 
   private loadImageFiles = throttle(async () => {
@@ -681,11 +639,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
       getSceneVersion(elements) >
       this.getLastBroadcastedOrReceivedSceneVersion()
     ) {
-      this.portal.broadcastScene(
-        SCENE.UPDATE,
-        this.getSyncableElements(elements),
-        false,
-      );
+      this.portal.broadcastScene(SCENE.UPDATE, elements, false);
       this.lastBroadcastedOrReceivedSceneVersion = getSceneVersion(elements);
       this.queueBroadcastAllElements();
     }
@@ -694,9 +648,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
   queueBroadcastAllElements = throttle(() => {
     this.portal.broadcastScene(
       SCENE.UPDATE,
-      this.getSyncableElements(
-        this.excalidrawAPI.getSceneElementsIncludingDeleted(),
-      ),
+      this.excalidrawAPI.getSceneElementsIncludingDeleted(),
       true,
     );
     const currentVersion = this.getLastBroadcastedOrReceivedSceneVersion();
@@ -722,8 +674,12 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
     });
   };
 
+  isSyncableElement = (element: ExcalidrawElement) => {
+    return element.isDeleted || !isInvisiblySmallElement(element);
+  };
+
   getSyncableElements = (elements: readonly ExcalidrawElement[]) =>
-    elements.filter((el) => el.isDeleted || !isInvisiblySmallElement(el));
+    elements.filter((element) => this.isSyncableElement(element));
 
   /** PRIVATE. Use `this.getContextValue()` instead. */
   private contextValue: CollabAPI | null = null;

src/excalidraw-app/collab/Portal.tsx

@@ -12,6 +12,7 @@ import { UserIdleState } from "../../types";
 import { trackEvent } from "../../analytics";
 import { throttle } from "lodash";
 import { mutateElement } from "../../element/mutateElement";
+import { BroadcastedExcalidrawElement } from "./reconciliation";
 
 class Portal {
   collab: CollabWrapper;
@@ -40,9 +41,7 @@ class Portal {
     this.socket.on("new-user", async (_socketId: string) => {
       this.broadcastScene(
         SCENE.INIT,
-        this.collab.getSyncableElements(
-          this.collab.getSceneElementsIncludingDeleted(),
-        ),
+        this.collab.getSceneElementsIncludingDeleted(),
         /* syncAll */ true,
       );
     });
@@ -124,24 +123,35 @@ class Portal {
 
   broadcastScene = async (
     sceneType: SCENE.INIT | SCENE.UPDATE,
-    syncableElements: ExcalidrawElement[],
+    allElements: readonly ExcalidrawElement[],
     syncAll: boolean,
   ) => {
     if (sceneType === SCENE.INIT && !syncAll) {
       throw new Error("syncAll must be true when sending SCENE.INIT");
     }
 
-    if (!syncAll) {
-      // sync out only the elements we think we need to to save bandwidth.
-      // periodically we'll resync the whole thing to make sure no one diverges
-      // due to a dropped message (server goes down etc).
-      syncableElements = syncableElements.filter(
-        (syncableElement) =>
-          !this.broadcastedElementVersions.has(syncableElement.id) ||
-          syncableElement.version >
-            this.broadcastedElementVersions.get(syncableElement.id)!,
-      );
-    }
+    // sync out only the elements we think we need to to save bandwidth.
+    // periodically we'll resync the whole thing to make sure no one diverges
+    // due to a dropped message (server goes down etc).
+    const syncableElements = allElements.reduce(
+      (acc, element: BroadcastedExcalidrawElement, idx, elements) => {
+        if (
+          (syncAll ||
+            !this.broadcastedElementVersions.has(element.id) ||
+            element.version >
+              this.broadcastedElementVersions.get(element.id)!) &&
+          this.collab.isSyncableElement(element)
+        ) {
+          acc.push({
+            ...element,
+            // z-index info for the reconciler
+            parent: idx === 0 ? "^" : elements[idx - 1]?.id,
+          });
+        }
+        return acc;
+      },
+      [] as BroadcastedExcalidrawElement[],
+    );
 
     const data: SocketUpdateDataSource[typeof sceneType] = {
       type: sceneType,

src/excalidraw-app/collab/reconciliation.ts

@@ -0,0 +1,162 @@
+import { ExcalidrawElement } from "../../element/types";
+import { AppState } from "../../types";
+
+export type ReconciledElements = readonly ExcalidrawElement[] & {
+  _brand: "reconciledElements";
+};
+
+export type BroadcastedExcalidrawElement = ExcalidrawElement & {
+  parent?: string;
+};
+
+const shouldDiscardRemoteElement = (
+  localAppState: AppState,
+  local: ExcalidrawElement | undefined,
+  remote: BroadcastedExcalidrawElement,
+): boolean => {
+  if (
+    local &&
+    // local element is being edited
+    (local.id === localAppState.editingElement?.id ||
+      local.id === localAppState.resizingElement?.id ||
+      local.id === localAppState.draggingElement?.id ||
+      // local element is newer
+      local.version > remote.version ||
+      // resolve conflicting edits deterministically by taking the one with
+      // the lowest versionNonce
+      (local.version === remote.version &&
+        local.versionNonce < remote.versionNonce))
+  ) {
+    return true;
+  }
+  return false;
+};
+
+const getElementsMapWithIndex = <T extends ExcalidrawElement>(
+  elements: readonly T[],
+) =>
+  elements.reduce(
+    (
+      acc: {
+        [key: string]: [element: T, index: number] | undefined;
+      },
+      element: T,
+      idx,
+    ) => {
+      acc[element.id] = [element, idx];
+      return acc;
+    },
+    {},
+  );
+
+export const reconcileElements = (
+  localElements: readonly ExcalidrawElement[],
+  remoteElements: readonly BroadcastedExcalidrawElement[],
+  localAppState: AppState,
+): ReconciledElements => {
+  const localElementsData = getElementsMapWithIndex<ExcalidrawElement>(
+    localElements,
+  );
+
+  const reconciledElements: ExcalidrawElement[] = localElements.slice();
+
+  const duplicates = new WeakMap<ExcalidrawElement, true>();
+
+  let cursor = 0;
+  let offset = 0;
+
+  let remoteElementIdx = -1;
+  for (const remoteElement of remoteElements) {
+    remoteElementIdx++;
+
+    const local = localElementsData[remoteElement.id];
+
+    if (shouldDiscardRemoteElement(localAppState, local?.[0], remoteElement)) {
+      if (remoteElement.parent) {
+        delete remoteElement.parent;
+      }
+
+      continue;
+    }
+
+    if (local) {
+      // mark for removal since it'll be replaced with the remote element
+      duplicates.set(local[0], true);
+    }
+
+    // parent may not be defined in case the remote client is running an older
+    // excalidraw version
+    const parent =
+      remoteElement.parent || remoteElements[remoteElementIdx - 1]?.id || null;
+
+    if (parent != null) {
+      delete remoteElement.parent;
+
+      // ^ indicates the element is the first in elements array
+      if (parent === "^") {
+        offset++;
+        if (cursor === 0) {
+          reconciledElements.unshift(remoteElement);
+          localElementsData[remoteElement.id] = [
+            remoteElement,
+            cursor - offset,
+          ];
+        } else {
+          reconciledElements.splice(cursor + 1, 0, remoteElement);
+          localElementsData[remoteElement.id] = [
+            remoteElement,
+            cursor + 1 - offset,
+          ];
+          cursor++;
+        }
+      } else {
+        let idx = localElementsData[parent]
+          ? localElementsData[parent]![1]
+          : null;
+        if (idx != null) {
+          idx += offset;
+        }
+        if (idx != null && idx >= cursor) {
+          reconciledElements.splice(idx + 1, 0, remoteElement);
+          offset++;
+          localElementsData[remoteElement.id] = [
+            remoteElement,
+            idx + 1 - offset,
+          ];
+          cursor = idx + 1;
+        } else if (idx != null) {
+          reconciledElements.splice(cursor + 1, 0, remoteElement);
+          offset++;
+          localElementsData[remoteElement.id] = [
+            remoteElement,
+            cursor + 1 - offset,
+          ];
+          cursor++;
+        } else {
+          reconciledElements.push(remoteElement);
+          localElementsData[remoteElement.id] = [
+            remoteElement,
+            reconciledElements.length - 1 - offset,
+          ];
+        }
+      }
+      // no parent z-index information, local element exists → replace in place
+    } else if (local) {
+      reconciledElements[local[1]] = remoteElement;
+      localElementsData[remoteElement.id] = [remoteElement, local[1]];
+      // otherwise push to the end
+    } else {
+      reconciledElements.push(remoteElement);
+      localElementsData[remoteElement.id] = [
+        remoteElement,
+        reconciledElements.length - 1 - offset,
+      ];
+    }
+  }
+
+  const ret: readonly ExcalidrawElement[] = reconciledElements.filter(
+    (element) => !duplicates.has(element),
+  );
+
+  return ret as ReconciledElements;
+};