feat: cache most of element selection (#6747)

src/scene/Scene.ts

@@ -11,6 +11,9 @@ import {
 } from "../element";
 import { LinearElementEditor } from "../element/linearElementEditor";
 import { isFrameElement } from "../element/typeChecks";
+import { getSelectedElements } from "./selection";
+import { AppState } from "../types";
+import { Assert, SameType } from "../utility-types";
 
 type ElementIdKey = InstanceType<typeof LinearElementEditor>["elementId"];
 type ElementKey = ExcalidrawElement | ElementIdKey;
@@ -18,6 +21,31 @@ type ElementKey = ExcalidrawElement | ElementIdKey;
 type SceneStateCallback = () => void;
 type SceneStateCallbackRemover = () => void;
 
+type SelectionHash = string & { __brand: "selectionHash" };
+
+const hashSelectionOpts = (
+  opts: Parameters<InstanceType<typeof Scene>["getSelectedElements"]>[0],
+) => {
+  const keys = ["includeBoundTextElement", "includeElementsInFrames"] as const;
+
+  type HashableKeys = Omit<typeof opts, "selectedElementIds" | "elements">;
+
+  // just to ensure we're hashing all expected keys
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  type _ = Assert<
+    SameType<
+      Required<HashableKeys>,
+      Pick<Required<HashableKeys>, typeof keys[number]>
+    >
+  >;
+
+  let hash = "";
+  for (const key of keys) {
+    hash += `${key}:${opts[key] ? "1" : "0"}`;
+  }
+  return hash as SelectionHash;
+};
+
 // ideally this would be a branded type but it'd be insanely hard to work with
 // in our codebase
 export type ExcalidrawElementsIncludingDeleted = readonly ExcalidrawElement[];
@@ -68,6 +96,15 @@ class Scene {
   private nonDeletedFrames: readonly NonDeleted<ExcalidrawFrameElement>[] = [];
   private frames: readonly ExcalidrawFrameElement[] = [];
   private elementsMap = new Map<ExcalidrawElement["id"], ExcalidrawElement>();
+  private selectedElementsCache: {
+    selectedElementIds: AppState["selectedElementIds"] | null;
+    elements: readonly NonDeletedExcalidrawElement[] | null;
+    cache: Map<SelectionHash, NonDeletedExcalidrawElement[]>;
+  } = {
+    selectedElementIds: null,
+    elements: null,
+    cache: new Map(),
+  };
 
   getElementsIncludingDeleted() {
     return this.elements;
@@ -81,6 +118,52 @@ class Scene {
     return this.frames;
   }
 
+  getSelectedElements(opts: {
+    // NOTE can be ommitted by making Scene constructor require App instance
+    selectedElementIds: AppState["selectedElementIds"];
+    /**
+     * for specific cases where you need to use elements not from current
+     * scene state. This in effect will likely result in cache-miss, and
+     * the cache won't be updated in this case.
+     */
+    elements?: readonly ExcalidrawElement[];
+    // selection-related options
+    includeBoundTextElement?: boolean;
+    includeElementsInFrames?: boolean;
+  }): NonDeleted<ExcalidrawElement>[] {
+    const hash = hashSelectionOpts(opts);
+
+    const elements = opts?.elements || this.nonDeletedElements;
+    if (
+      this.selectedElementsCache.elements === elements &&
+      this.selectedElementsCache.selectedElementIds === opts.selectedElementIds
+    ) {
+      const cached = this.selectedElementsCache.cache.get(hash);
+      if (cached) {
+        return cached;
+      }
+    } else if (opts?.elements == null) {
+      // if we're operating on latest scene elements and the cache is not
+      //  storing the latest elements, clear the cache
+      this.selectedElementsCache.cache.clear();
+    }
+
+    const selectedElements = getSelectedElements(
+      elements,
+      { selectedElementIds: opts.selectedElementIds },
+      opts,
+    );
+
+    // cache only if we're not using custom elements
+    if (opts?.elements == null) {
+      this.selectedElementsCache.selectedElementIds = opts.selectedElementIds;
+      this.selectedElementsCache.elements = this.nonDeletedElements;
+      this.selectedElementsCache.cache.set(hash, selectedElements);
+    }
+
+    return selectedElements;
+  }
+
   getNonDeletedFrames(): readonly NonDeleted<ExcalidrawFrameElement>[] {
     return this.nonDeletedFrames;
   }
@@ -168,11 +251,21 @@ class Scene {
   }
 
   destroy() {
+    this.nonDeletedElements = [];
+    this.elements = [];
+    this.nonDeletedFrames = [];
+    this.frames = [];
+    this.elementsMap.clear();
+    this.selectedElementsCache.selectedElementIds = null;
+    this.selectedElementsCache.elements = null;
+    this.selectedElementsCache.cache.clear();
+
     Scene.sceneMapById.forEach((scene, elementKey) => {
       if (scene === this) {
         Scene.sceneMapById.delete(elementKey);
       }
     });
+
     // done not for memory leaks, but to guard against possible late fires
     // (I guess?)
     this.callbacks.clear();