fix: sort bound text elements to fix text duplication z-index error (#5130)

* fix: sort bound text elements to fix text duplication z-index error

* improve & sort groups & add tests

* fix backtracking and discontiguous groups

---------

Co-authored-by: dwelle <luzar.david@gmail.com>

src/actions/actionDuplicateSelection.tsx

@@ -16,8 +16,12 @@ import { AppState } from "../types";
 import { fixBindingsAfterDuplication } from "../element/binding";
 import { ActionResult } from "./types";
 import { GRID_SIZE } from "../constants";
-import { bindTextToShapeAfterDuplication } from "../element/textElement";
+import {
+  bindTextToShapeAfterDuplication,
+  getBoundTextElement,
+} from "../element/textElement";
 import { isBoundToContainer } from "../element/typeChecks";
+import { normalizeElementOrder } from "../element/sortElements";
 import { DuplicateIcon } from "../components/icons";
 
 export const actionDuplicateSelection = register({
@@ -64,6 +68,11 @@ const duplicateElements = (
   elements: readonly ExcalidrawElement[],
   appState: AppState,
 ): Partial<ActionResult> => {
+  // ---------------------------------------------------------------------------
+
+  // step (1)
+
+  const sortedElements = normalizeElementOrder(elements);
   const groupIdMap = new Map();
   const newElements: ExcalidrawElement[] = [];
   const oldElements: ExcalidrawElement[] = [];
@@ -85,42 +94,112 @@ const duplicateElements = (
     return newElement;
   };
 
-  const finalElements: ExcalidrawElement[] = [];
-
-  let index = 0;
   const selectedElementIds = arrayToMap(
-    getSelectedElements(elements, appState, true),
+    getSelectedElements(sortedElements, appState, true),
   );
-  while (index < elements.length) {
-    const element = elements[index];
+
+  // Ids of elements that have already been processed so we don't push them
+  // into the array twice if we end up backtracking when retrieving
+  // discontiguous group of elements (can happen due to a bug, or in edge
+  // cases such as a group containing deleted elements which were not selected).
+  //
+  // This is not enough to prevent duplicates, so we do a second loop afterwards
+  // to remove them.
+  //
+  // For convenience we mark even the newly created ones even though we don't
+  // loop over them.
+  const processedIds = new Map<ExcalidrawElement["id"], true>();
+
+  const markAsProcessed = (elements: ExcalidrawElement[]) => {
+    for (const element of elements) {
+      processedIds.set(element.id, true);
+    }
+    return elements;
+  };
+
+  const elementsWithClones: ExcalidrawElement[] = [];
+
+  let index = -1;
+
+  while (++index < sortedElements.length) {
+    const element = sortedElements[index];
+
+    if (processedIds.get(element.id)) {
+      continue;
+    }
+
+    const boundTextElement = getBoundTextElement(element);
     if (selectedElementIds.get(element.id)) {
-      if (element.groupIds.length) {
+      // if a group or a container/bound-text, duplicate atomically
+      if (element.groupIds.length || boundTextElement) {
         const groupId = getSelectedGroupForElement(appState, element);
-        // if group selected, duplicate it atomically
         if (groupId) {
-          const groupElements = getElementsInGroup(elements, groupId);
-          finalElements.push(
-            ...groupElements,
-            ...groupElements.map((element) =>
+          const groupElements = getElementsInGroup(sortedElements, groupId);
+          elementsWithClones.push(
+            ...markAsProcessed([
+              ...groupElements,
+              ...groupElements.map((element) =>
+                duplicateAndOffsetElement(element),
+              ),
+            ]),
+          );
+          continue;
+        }
+        if (boundTextElement) {
+          elementsWithClones.push(
+            ...markAsProcessed([
+              element,
+              boundTextElement,
               duplicateAndOffsetElement(element),
-            ),
+              duplicateAndOffsetElement(boundTextElement),
+            ]),
           );
-          index = index + groupElements.length;
           continue;
         }
       }
-      finalElements.push(element, duplicateAndOffsetElement(element));
+      elementsWithClones.push(
+        ...markAsProcessed([element, duplicateAndOffsetElement(element)]),
+      );
     } else {
-      finalElements.push(element);
+      elementsWithClones.push(...markAsProcessed([element]));
+    }
+  }
+
+  // step (2)
+
+  // second pass to remove duplicates. We loop from the end as it's likelier
+  // that the last elements are in the correct order (contiguous or otherwise).
+  // Thus we need to reverse as the last step (3).
+
+  const finalElementsReversed: ExcalidrawElement[] = [];
+
+  const finalElementIds = new Map<ExcalidrawElement["id"], true>();
+  index = elementsWithClones.length;
+
+  while (--index >= 0) {
+    const element = elementsWithClones[index];
+    if (!finalElementIds.get(element.id)) {
+      finalElementIds.set(element.id, true);
+      finalElementsReversed.push(element);
     }
-    index++;
   }
+
+  // step (3)
+
+  const finalElements = finalElementsReversed.reverse();
+
+  // ---------------------------------------------------------------------------
+
   bindTextToShapeAfterDuplication(
-    finalElements,
+    elementsWithClones,
+    oldElements,
+    oldIdToDuplicatedId,
+  );
+  fixBindingsAfterDuplication(
+    elementsWithClones,
     oldElements,
     oldIdToDuplicatedId,
   );
-  fixBindingsAfterDuplication(finalElements, oldElements, oldIdToDuplicatedId);
 
   return {
     elements: finalElements,

src/element/sortElements.test.ts

@@ -0,0 +1,402 @@
+import { API } from "../tests/helpers/api";
+import { mutateElement } from "./mutateElement";
+import { normalizeElementOrder } from "./sortElements";
+import { ExcalidrawElement } from "./types";
+
+const assertOrder = (
+  elements: readonly ExcalidrawElement[],
+  expectedOrder: string[],
+) => {
+  const actualOrder = elements.map((element) => element.id);
+  expect(actualOrder).toEqual(expectedOrder);
+};
+
+describe("normalizeElementsOrder", () => {
+  it("sort bound-text elements", () => {
+    const container = API.createElement({
+      id: "container",
+      type: "rectangle",
+    });
+    const boundText = API.createElement({
+      id: "boundText",
+      type: "text",
+      containerId: container.id,
+    });
+    const otherElement = API.createElement({
+      id: "otherElement",
+      type: "rectangle",
+      boundElements: [],
+    });
+    const otherElement2 = API.createElement({
+      id: "otherElement2",
+      type: "rectangle",
+      boundElements: [],
+    });
+
+    mutateElement(container, {
+      boundElements: [{ type: "text", id: boundText.id }],
+    });
+
+    assertOrder(normalizeElementOrder([container, boundText]), [
+      "container",
+      "boundText",
+    ]);
+    assertOrder(normalizeElementOrder([boundText, container]), [
+      "container",
+      "boundText",
+    ]);
+    assertOrder(
+      normalizeElementOrder([
+        boundText,
+        container,
+        otherElement,
+        otherElement2,
+      ]),
+      ["container", "boundText", "otherElement", "otherElement2"],
+    );
+    assertOrder(normalizeElementOrder([container, otherElement, boundText]), [
+      "container",
+      "boundText",
+      "otherElement",
+    ]);
+    assertOrder(
+      normalizeElementOrder([
+        container,
+        otherElement,
+        otherElement2,
+        boundText,
+      ]),
+      ["container", "boundText", "otherElement", "otherElement2"],
+    );
+
+    assertOrder(
+      normalizeElementOrder([
+        boundText,
+        otherElement,
+        container,
+        otherElement2,
+      ]),
+      ["otherElement", "container", "boundText", "otherElement2"],
+    );
+
+    // noop
+    assertOrder(
+      normalizeElementOrder([
+        otherElement,
+        container,
+        boundText,
+        otherElement2,
+      ]),
+      ["otherElement", "container", "boundText", "otherElement2"],
+    );
+
+    // text has existing containerId, but container doesn't list is
+    // as a boundElement
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "boundText",
+          type: "text",
+          containerId: "container",
+        }),
+        API.createElement({
+          id: "container",
+          type: "rectangle",
+        }),
+      ]),
+      ["boundText", "container"],
+    );
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "boundText",
+          type: "text",
+          containerId: "container",
+        }),
+      ]),
+      ["boundText"],
+    );
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "container",
+          type: "rectangle",
+          boundElements: [],
+        }),
+      ]),
+      ["container"],
+    );
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "container",
+          type: "rectangle",
+          boundElements: [{ id: "x", type: "text" }],
+        }),
+      ]),
+      ["container"],
+    );
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "arrow",
+          type: "arrow",
+        }),
+        API.createElement({
+          id: "container",
+          type: "rectangle",
+          boundElements: [{ id: "arrow", type: "arrow" }],
+        }),
+      ]),
+      ["arrow", "container"],
+    );
+  });
+
+  it("normalize group order", () => {
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "A_rect1",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+        API.createElement({
+          id: "rect2",
+          type: "rectangle",
+        }),
+        API.createElement({
+          id: "rect3",
+          type: "rectangle",
+        }),
+        API.createElement({
+          id: "A_rect4",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+        API.createElement({
+          id: "A_rect5",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+        API.createElement({
+          id: "rect6",
+          type: "rectangle",
+        }),
+        API.createElement({
+          id: "A_rect7",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+      ]),
+      ["A_rect1", "A_rect4", "A_rect5", "A_rect7", "rect2", "rect3", "rect6"],
+    );
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "A_rect1",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+        API.createElement({
+          id: "rect2",
+          type: "rectangle",
+        }),
+        API.createElement({
+          id: "B_rect3",
+          type: "rectangle",
+          groupIds: ["B"],
+        }),
+        API.createElement({
+          id: "A_rect4",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+        API.createElement({
+          id: "B_rect5",
+          type: "rectangle",
+          groupIds: ["B"],
+        }),
+        API.createElement({
+          id: "rect6",
+          type: "rectangle",
+        }),
+        API.createElement({
+          id: "A_rect7",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+      ]),
+      ["A_rect1", "A_rect4", "A_rect7", "rect2", "B_rect3", "B_rect5", "rect6"],
+    );
+    // nested groups
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "A_rect1",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+        API.createElement({
+          id: "BA_rect2",
+          type: "rectangle",
+          groupIds: ["B", "A"],
+        }),
+      ]),
+      ["A_rect1", "BA_rect2"],
+    );
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "BA_rect1",
+          type: "rectangle",
+          groupIds: ["B", "A"],
+        }),
+        API.createElement({
+          id: "A_rect2",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+      ]),
+      ["BA_rect1", "A_rect2"],
+    );
+    assertOrder(
+      normalizeElementOrder([
+        API.createElement({
+          id: "BA_rect1",
+          type: "rectangle",
+          groupIds: ["B", "A"],
+        }),
+        API.createElement({
+          id: "A_rect2",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+        API.createElement({
+          id: "CBA_rect3",
+          type: "rectangle",
+          groupIds: ["C", "B", "A"],
+        }),
+        API.createElement({
+          id: "rect4",
+          type: "rectangle",
+        }),
+        API.createElement({
+          id: "A_rect5",
+          type: "rectangle",
+          groupIds: ["A"],
+        }),
+        API.createElement({
+          id: "BA_rect5",
+          type: "rectangle",
+          groupIds: ["B", "A"],
+        }),
+        API.createElement({
+          id: "BA_rect6",
+          type: "rectangle",
+          groupIds: ["B", "A"],
+        }),
+        API.createElement({
+          id: "CBA_rect7",
+          type: "rectangle",
+          groupIds: ["C", "B", "A"],
+        }),
+        API.createElement({
+          id: "X_rect8",
+          type: "rectangle",
+          groupIds: ["X"],
+        }),
+        API.createElement({
+          id: "rect9",
+          type: "rectangle",
+        }),
+        API.createElement({
+          id: "YX_rect10",
+          type: "rectangle",
+          groupIds: ["Y", "X"],
+        }),
+        API.createElement({
+          id: "X_rect11",
+          type: "rectangle",
+          groupIds: ["X"],
+        }),
+      ]),
+      [
+        "BA_rect1",
+        "BA_rect5",
+        "BA_rect6",
+        "A_rect2",
+        "A_rect5",
+        "CBA_rect3",
+        "CBA_rect7",
+        "rect4",
+        "X_rect8",
+        "X_rect11",
+        "YX_rect10",
+        "rect9",
+      ],
+    );
+  });
+
+  // TODO
+  it.skip("normalize boundElements array", () => {
+    const container = API.createElement({
+      id: "container",
+      type: "rectangle",
+      boundElements: [],
+    });
+    const boundText = API.createElement({
+      id: "boundText",
+      type: "text",
+      containerId: container.id,
+    });
+
+    mutateElement(container, {
+      boundElements: [
+        { type: "text", id: boundText.id },
+        { type: "text", id: "xxx" },
+      ],
+    });
+
+    expect(normalizeElementOrder([container, boundText])).toEqual([
+      expect.objectContaining({
+        id: container.id,
+      }),
+      expect.objectContaining({ id: boundText.id }),
+    ]);
+  });
+
+  // should take around <100ms for 10K iterations (@dwelle's PC 22-05-25)
+  it.skip("normalizeElementsOrder() perf", () => {
+    const makeElements = (iterations: number) => {
+      const elements: ExcalidrawElement[] = [];
+      while (iterations--) {
+        const container = API.createElement({
+          type: "rectangle",
+          boundElements: [],
+          groupIds: ["B", "A"],
+        });
+        const boundText = API.createElement({
+          type: "text",
+          containerId: container.id,
+          groupIds: ["A"],
+        });
+        const otherElement = API.createElement({
+          type: "rectangle",
+          boundElements: [],
+          groupIds: ["C", "A"],
+        });
+        mutateElement(container, {
+          boundElements: [{ type: "text", id: boundText.id }],
+        });
+
+        elements.push(boundText, otherElement, container);
+      }
+      return elements;
+    };
+
+    const elements = makeElements(10000);
+    const t0 = Date.now();
+    normalizeElementOrder(elements);
+    console.info(`${Date.now() - t0}ms`);
+  });
+});

src/element/sortElements.ts

@@ -0,0 +1,123 @@
+import { arrayToMapWithIndex } from "../utils";
+import { ExcalidrawElement } from "./types";
+
+const normalizeGroupElementOrder = (elements: readonly ExcalidrawElement[]) => {
+  const origElements: ExcalidrawElement[] = elements.slice();
+  const sortedElements = new Set<ExcalidrawElement>();
+
+  const orderInnerGroups = (
+    elements: readonly ExcalidrawElement[],
+  ): ExcalidrawElement[] => {
+    const firstGroupSig = elements[0]?.groupIds?.join("");
+    const aGroup: ExcalidrawElement[] = [elements[0]];
+    const bGroup: ExcalidrawElement[] = [];
+    for (const element of elements.slice(1)) {
+      if (element.groupIds?.join("") === firstGroupSig) {
+        aGroup.push(element);
+      } else {
+        bGroup.push(element);
+      }
+    }
+    return bGroup.length ? [...aGroup, ...orderInnerGroups(bGroup)] : aGroup;
+  };
+
+  const groupHandledElements = new Map<string, true>();
+
+  origElements.forEach((element, idx) => {
+    if (groupHandledElements.has(element.id)) {
+      return;
+    }
+    if (element.groupIds?.length) {
+      const topGroup = element.groupIds[element.groupIds.length - 1];
+      const groupElements = origElements.slice(idx).filter((element) => {
+        const ret = element?.groupIds?.some((id) => id === topGroup);
+        if (ret) {
+          groupHandledElements.set(element!.id, true);
+        }
+        return ret;
+      });
+
+      for (const elem of orderInnerGroups(groupElements)) {
+        sortedElements.add(elem);
+      }
+    } else {
+      sortedElements.add(element);
+    }
+  });
+
+  // if there's a bug which resulted in losing some of the elements, return
+  // original instead as that's better than losing data
+  if (sortedElements.size !== elements.length) {
+    console.error("normalizeGroupElementOrder: lost some elements... bailing!");
+    return elements;
+  }
+
+  return [...sortedElements];
+};
+
+/**
+ * In theory, when we have text elements bound to a container, they
+ * should be right after the container element in the elements array.
+ * However, this is not guaranteed due to old and potential future bugs.
+ *
+ * This function sorts containers and their bound texts together. It prefers
+ * original z-index of container (i.e. it moves bound text elements after
+ * containers).
+ */
+const normalizeBoundElementsOrder = (
+  elements: readonly ExcalidrawElement[],
+) => {
+  const elementsMap = arrayToMapWithIndex(elements);
+
+  const origElements: (ExcalidrawElement | null)[] = elements.slice();
+  const sortedElements = new Set<ExcalidrawElement>();
+
+  origElements.forEach((element, idx) => {
+    if (!element) {
+      return;
+    }
+    if (element.boundElements?.length) {
+      sortedElements.add(element);
+      origElements[idx] = null;
+      element.boundElements.forEach((boundElement) => {
+        const child = elementsMap.get(boundElement.id);
+        if (child && boundElement.type === "text") {
+          sortedElements.add(child[0]);
+          origElements[child[1]] = null;
+        }
+      });
+    } else if (element.type === "text" && element.containerId) {
+      const parent = elementsMap.get(element.containerId);
+      if (!parent?.[0].boundElements?.find((x) => x.id === element.id)) {
+        sortedElements.add(element);
+        origElements[idx] = null;
+
+        // if element has a container and container lists it, skip this element
+        // as it'll be taken care of by the container
+      }
+    } else {
+      sortedElements.add(element);
+      origElements[idx] = null;
+    }
+  });
+
+  // if there's a bug which resulted in losing some of the elements, return
+  // original instead as that's better than losing data
+  if (sortedElements.size !== elements.length) {
+    console.error(
+      "normalizeBoundElementsOrder: lost some elements... bailing!",
+    );
+    return elements;
+  }
+
+  return [...sortedElements];
+};
+
+export const normalizeElementOrder = (
+  elements: readonly ExcalidrawElement[],
+) => {
+  // console.time();
+  const ret = normalizeBoundElementsOrder(normalizeGroupElementOrder(elements));
+  // console.timeEnd();
+  return ret;
+};

src/element/textElement.ts

@@ -127,10 +127,16 @@ export const bindTextToShapeAfterDuplication = (
         const newContainer = sceneElementMap.get(newElementId);
         if (newContainer) {
           mutateElement(newContainer, {
-            boundElements: (newContainer.boundElements || []).concat({
-              type: "text",
-              id: newTextElementId,
-            }),
+            boundElements: (element.boundElements || [])
+              .filter(
+                (boundElement) =>
+                  boundElement.id !== newTextElementId &&
+                  boundElement.id !== boundTextElementId,
+              )
+              .concat({
+                type: "text",
+                id: newTextElementId,
+              }),
           });
         }
         const newTextElement = sceneElementMap.get(newTextElementId);

src/excalidraw-app/collab/reconciliation.ts

@@ -1,6 +1,7 @@
 import { PRECEDING_ELEMENT_KEY } from "../../constants";
 import { ExcalidrawElement } from "../../element/types";
 import { AppState } from "../../types";
+import { arrayToMapWithIndex } from "../../utils";
 
 export type ReconciledElements = readonly ExcalidrawElement[] & {
   _brand: "reconciledElements";
@@ -33,30 +34,13 @@ const shouldDiscardRemoteElement = (
   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);
+    arrayToMapWithIndex<ExcalidrawElement>(localElements);
 
   const reconciledElements: ExcalidrawElement[] = localElements.slice();
 
@@ -69,7 +53,7 @@ export const reconcileElements = (
   for (const remoteElement of remoteElements) {
     remoteElementIdx++;
 
-    const local = localElementsData[remoteElement.id];
+    const local = localElementsData.get(remoteElement.id);
 
     if (shouldDiscardRemoteElement(localAppState, local?.[0], remoteElement)) {
       if (remoteElement[PRECEDING_ELEMENT_KEY]) {
@@ -105,21 +89,21 @@ export const reconcileElements = (
         offset++;
         if (cursor === 0) {
           reconciledElements.unshift(remoteElement);
-          localElementsData[remoteElement.id] = [
+          localElementsData.set(remoteElement.id, [
             remoteElement,
             cursor - offset,
-          ];
+          ]);
         } else {
           reconciledElements.splice(cursor + 1, 0, remoteElement);
-          localElementsData[remoteElement.id] = [
+          localElementsData.set(remoteElement.id, [
             remoteElement,
             cursor + 1 - offset,
-          ];
+          ]);
           cursor++;
         }
       } else {
-        let idx = localElementsData[parent]
-          ? localElementsData[parent]![1]
+        let idx = localElementsData.has(parent)
+          ? localElementsData.get(parent)![1]
           : null;
         if (idx != null) {
           idx += offset;
@@ -127,38 +111,38 @@ export const reconcileElements = (
         if (idx != null && idx >= cursor) {
           reconciledElements.splice(idx + 1, 0, remoteElement);
           offset++;
-          localElementsData[remoteElement.id] = [
+          localElementsData.set(remoteElement.id, [
             remoteElement,
             idx + 1 - offset,
-          ];
+          ]);
           cursor = idx + 1;
         } else if (idx != null) {
           reconciledElements.splice(cursor + 1, 0, remoteElement);
           offset++;
-          localElementsData[remoteElement.id] = [
+          localElementsData.set(remoteElement.id, [
             remoteElement,
             cursor + 1 - offset,
-          ];
+          ]);
           cursor++;
         } else {
           reconciledElements.push(remoteElement);
-          localElementsData[remoteElement.id] = [
+          localElementsData.set(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]];
+      localElementsData.set(remoteElement.id, [remoteElement, local[1]]);
       // otherwise push to the end
     } else {
       reconciledElements.push(remoteElement);
-      localElementsData[remoteElement.id] = [
+      localElementsData.set(remoteElement.id, [
         remoteElement,
         reconciledElements.length - 1 - offset,
-      ];
+      ]);
     }
   }
 

src/tests/zindex.test.tsx

@@ -11,6 +11,7 @@ import {
 } from "../actions";
 import { AppState } from "../types";
 import { API } from "./helpers/api";
+import { selectGroupsForSelectedElements } from "../groups";
 
 // Unmount ReactDOM from root
 ReactDOM.unmountComponentAtNode(document.getElementById("root")!);
@@ -34,6 +35,7 @@ const populateElements = (
     height?: number;
     containerId?: string;
   }[],
+  appState?: Partial<AppState>,
 ) => {
   const selectedElementIds: any = {};
 
@@ -84,6 +86,11 @@ const populateElements = (
   });
 
   h.setState({
+    ...selectGroupsForSelectedElements(
+      { ...h.state, ...appState, selectedElementIds },
+      h.elements,
+    ),
+    ...appState,
     selectedElementIds,
   });
 
@@ -111,11 +118,7 @@ const assertZindex = ({
   appState?: Partial<AppState>;
   operations: [Actions, string[]][];
 }) => {
-  const selectedElementIds = populateElements(elements);
-
-  h.setState({
-    editingGroupId: appState?.editingGroupId || null,
-  });
+  const selectedElementIds = populateElements(elements, appState);
 
   operations.forEach(([action, expected]) => {
     h.app.actionManager.executeAction(action);
@@ -884,9 +887,6 @@ describe("z-index manipulation", () => {
       { id: "A", groupIds: ["g1"], isSelected: true },
       { id: "B", groupIds: ["g1"], isSelected: true },
     ]);
-    h.setState({
-      selectedGroupIds: { g1: true },
-    });
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements).toMatchObject([
       { id: "A" },
@@ -908,9 +908,6 @@ describe("z-index manipulation", () => {
       { id: "B", groupIds: ["g1"], isSelected: true },
       { id: "C" },
     ]);
-    h.setState({
-      selectedGroupIds: { g1: true },
-    });
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements).toMatchObject([
       { id: "A" },
@@ -933,9 +930,6 @@ describe("z-index manipulation", () => {
       { id: "B", groupIds: ["g1"], isSelected: true },
       { id: "C", isSelected: true },
     ]);
-    h.setState({
-      selectedGroupIds: { g1: true },
-    });
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements.map((element) => element.id)).toEqual([
       "A",
@@ -952,9 +946,6 @@ describe("z-index manipulation", () => {
       { id: "C", groupIds: ["g2"], isSelected: true },
       { id: "D", groupIds: ["g2"], isSelected: true },
     ]);
-    h.setState({
-      selectedGroupIds: { g1: true, g2: true },
-    });
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements.map((element) => element.id)).toEqual([
       "A",
@@ -967,14 +958,16 @@ describe("z-index manipulation", () => {
       "D_copy",
     ]);
 
-    populateElements([
-      { id: "A", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "B", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "C", groupIds: ["g2"], isSelected: true },
-    ]);
-    h.setState({
-      selectedGroupIds: { g1: true },
-    });
+    populateElements(
+      [
+        { id: "A", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "B", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "C", groupIds: ["g2"], isSelected: true },
+      ],
+      {
+        selectedGroupIds: { g1: true },
+      },
+    );
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements.map((element) => element.id)).toEqual([
       "A",
@@ -985,14 +978,16 @@ describe("z-index manipulation", () => {
       "C_copy",
     ]);
 
-    populateElements([
-      { id: "A", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "B", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "C", groupIds: ["g2"], isSelected: true },
-    ]);
-    h.setState({
-      selectedGroupIds: { g2: true },
-    });
+    populateElements(
+      [
+        { id: "A", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "B", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "C", groupIds: ["g2"], isSelected: true },
+      ],
+      {
+        selectedGroupIds: { g2: true },
+      },
+    );
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements.map((element) => element.id)).toEqual([
       "A",
@@ -1003,17 +998,19 @@ describe("z-index manipulation", () => {
       "C_copy",
     ]);
 
-    populateElements([
-      { id: "A", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "B", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "C", groupIds: ["g2"], isSelected: true },
-      { id: "D", groupIds: ["g3", "g4"], isSelected: true },
-      { id: "E", groupIds: ["g3", "g4"], isSelected: true },
-      { id: "F", groupIds: ["g4"], isSelected: true },
-    ]);
-    h.setState({
-      selectedGroupIds: { g2: true, g4: true },
-    });
+    populateElements(
+      [
+        { id: "A", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "B", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "C", groupIds: ["g2"], isSelected: true },
+        { id: "D", groupIds: ["g3", "g4"], isSelected: true },
+        { id: "E", groupIds: ["g3", "g4"], isSelected: true },
+        { id: "F", groupIds: ["g4"], isSelected: true },
+      ],
+      {
+        selectedGroupIds: { g2: true, g4: true },
+      },
+    );
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements.map((element) => element.id)).toEqual([
       "A",
@@ -1030,11 +1027,14 @@ describe("z-index manipulation", () => {
       "F_copy",
     ]);
 
-    populateElements([
-      { id: "A", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "B", groupIds: ["g1", "g2"] },
-      { id: "C", groupIds: ["g2"] },
-    ]);
+    populateElements(
+      [
+        { id: "A", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "B", groupIds: ["g1", "g2"] },
+        { id: "C", groupIds: ["g2"] },
+      ],
+      { editingGroupId: "g1" },
+    );
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements.map((element) => element.id)).toEqual([
       "A",
@@ -1043,32 +1043,73 @@ describe("z-index manipulation", () => {
       "C",
     ]);
 
-    populateElements([
-      { id: "A", groupIds: ["g1", "g2"] },
-      { id: "B", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "C", groupIds: ["g2"] },
+    populateElements(
+      [
+        { id: "A", groupIds: ["g1", "g2"] },
+        { id: "B", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "C", groupIds: ["g2"] },
+      ],
+      { editingGroupId: "g1" },
+    );
+    h.app.actionManager.executeAction(actionDuplicateSelection);
+    expect(h.elements.map((element) => element.id)).toEqual([
+      "A",
+      "B",
+      "B_copy",
+      "C",
     ]);
+
+    populateElements(
+      [
+        { id: "A", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "B", groupIds: ["g1", "g2"], isSelected: true },
+        { id: "C", groupIds: ["g2"] },
+      ],
+      { editingGroupId: "g1" },
+    );
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements.map((element) => element.id)).toEqual([
       "A",
+      "A_copy",
       "B",
       "B_copy",
       "C",
     ]);
+  });
 
+  it("duplicating incorrectly interleaved elements (group elements should be together) should still produce reasonable result", () => {
     populateElements([
-      { id: "A", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "B", groupIds: ["g1", "g2"], isSelected: true },
-      { id: "C", groupIds: ["g2"], isSelected: true },
+      { id: "A", groupIds: ["g1"], isSelected: true },
+      { id: "B" },
+      { id: "C", groupIds: ["g1"], isSelected: true },
     ]);
     h.app.actionManager.executeAction(actionDuplicateSelection);
     expect(h.elements.map((element) => element.id)).toEqual([
       "A",
+      "C",
       "A_copy",
+      "C_copy",
+      "B",
+    ]);
+  });
+
+  it("group-selected duplication should includes deleted elements that weren't selected on account of being deleted", () => {
+    populateElements([
+      { id: "A", groupIds: ["g1"], isDeleted: true },
+      { id: "B", groupIds: ["g1"], isSelected: true },
+      { id: "C", groupIds: ["g1"], isSelected: true },
+      { id: "D" },
+    ]);
+    expect(h.state.selectedGroupIds).toEqual({ g1: true });
+    h.app.actionManager.executeAction(actionDuplicateSelection);
+    expect(h.elements.map((element) => element.id)).toEqual([
+      "A",
       "B",
-      "B_copy",
       "C",
+      "A_copy",
+      "B_copy",
       "C_copy",
+      "D",
     ]);
   });
 

src/utils.ts

@@ -607,6 +607,14 @@ export const arrayToMap = <T extends { id: string } | string>(
   }, new Map());
 };
 
+export const arrayToMapWithIndex = <T extends { id: string }>(
+  elements: readonly T[],
+) =>
+  elements.reduce((acc, element: T, idx) => {
+    acc.set(element.id, [element, idx]);
+    return acc;
+  }, new Map<string, [element: T, index: number]>());
+
 export const isTestEnv = () =>
   typeof process !== "undefined" && process.env?.NODE_ENV === "test";