Replace O(n²) lookups with Map-based atoms in schema state layer

packages/graph-explorer/src/core/StateProvider/configuration.test.ts

@@ -307,6 +307,25 @@ describe("patchToRemoveDisplayLabel", () => {
       expect(attr).not.toHaveProperty("displayLabel");
     }
   });
+
+  it("should not mutate the original config", () => {
+    const config = createRandomVertexTypeConfig();
+    config.displayLabel = createRandomName("displayLabel");
+    config.attributes.forEach(
+      a => ((a as any).displayLabel = createRandomName("displayLabel")),
+    );
+    const originalDisplayLabel = config.displayLabel;
+    const originalAttrDisplayLabels = config.attributes.map(
+      a => (a as any).displayLabel,
+    );
+
+    patchToRemoveDisplayLabel(config);
+
+    expect(config.displayLabel).toBe(originalDisplayLabel);
+    config.attributes.forEach((a, i) => {
+      expect((a as any).displayLabel).toBe(originalAttrDisplayLabels[i]);
+    });
+  });
 });
 
 describe("normalizeConnection", () => {

packages/graph-explorer/src/core/StateProvider/configuration.ts

@@ -2,7 +2,7 @@ import type { ConnectionConfig } from "@shared/types";
 
 import { atom } from "jotai";
 import { selectAtom } from "jotai/utils";
-import { cloneDeep, isEqual, uniq } from "lodash";
+import { isEqual } from "lodash";
 
 import {
   activeConfigurationAtom,
@@ -63,30 +63,37 @@ export function mergeConfiguration(
   currentConfig: RawConfiguration,
   userStyling: UserStyling,
 ): RawConfiguration {
-  const configVLabels = currentConfig.schema?.vertices.map(v => v.type) || [];
-  const schemaVLabels = currentSchema?.vertices?.map(v => v.type) || [];
-  const allVertexLabels = uniq([...configVLabels, ...schemaVLabels]);
+  const configVertexMap = toMapByType(currentConfig.schema?.vertices);
+  const schemaVertexMap = toMapByType(currentSchema?.vertices);
+  const prefsVertexMap = toMapByType(userStyling.vertices);
+
+  const allVertexLabels = [
+    ...new Set([...configVertexMap.keys(), ...schemaVertexMap.keys()]),
+  ];
   const mergedVertices = allVertexLabels
-    .map(vLabel => {
-      const configVertex = currentConfig.schema?.vertices.find(
-        v => v.type === vLabel,
-      );
-      const schemaVertex = currentSchema?.vertices.find(v => v.type === vLabel);
-      const prefsVertex = userStyling.vertices?.find(v => v.type === vLabel);
-
-      return mergeVertex(configVertex, schemaVertex, prefsVertex);
-    })
+    .map(vLabel =>
+      mergeVertex(
+        configVertexMap.get(vLabel),
+        schemaVertexMap.get(vLabel),
+        prefsVertexMap.get(vLabel),
+      ),
+    )
     .toSorted((a, b) => a.type.localeCompare(b.type));
 
-  const configELabels = currentConfig.schema?.edges.map(v => v.type) || [];
-  const schemaELabels = currentSchema?.edges?.map(v => v.type) || [];
-  const allEdgeLabels = uniq([...configELabels, ...schemaELabels]);
-  const mergedEdges = allEdgeLabels.map(vLabel => {
-    const configEdge = currentConfig.schema?.edges.find(v => v.type === vLabel);
-    const schemaEdge = currentSchema?.edges.find(v => v.type === vLabel);
-    const prefsEdge = userStyling.edges?.find(v => v.type === vLabel);
-    return mergeEdge(configEdge, schemaEdge, prefsEdge);
-  });
+  const configEdgeMap = toMapByType(currentConfig.schema?.edges);
+  const schemaEdgeMap = toMapByType(currentSchema?.edges);
+  const prefsEdgeMap = toMapByType(userStyling.edges);
+
+  const allEdgeLabels = [
+    ...new Set([...configEdgeMap.keys(), ...schemaEdgeMap.keys()]),
+  ];
+  const mergedEdges = allEdgeLabels.map(eLabel =>
+    mergeEdge(
+      configEdgeMap.get(eLabel),
+      schemaEdgeMap.get(eLabel),
+      prefsEdgeMap.get(eLabel),
+    ),
+  );
 
   return {
     id: currentConfig.id,
@@ -126,20 +133,21 @@ const mergeAttributes = (
   config: VertexTypeConfig | EdgeTypeConfig | null,
   schema: VertexTypeConfig | EdgeTypeConfig | null,
 ): AttributeConfig[] => {
-  const configAttrLabels = config?.attributes.map(attr => attr.name) || [];
-  const schemaAttrLabels = schema?.attributes.map(attr => attr.name) || [];
-  const allAttrLabels = uniq([...configAttrLabels, ...schemaAttrLabels]);
-
-  return allAttrLabels.map(attrName => {
-    const configAttr = config?.attributes.find(attr => attr.name === attrName);
-    const schemaAttr = schema?.attributes.find(attr => attr.name === attrName);
-
-    return {
-      name: attrName,
-      ...(schemaAttr || {}),
-      ...(configAttr || {}),
-    };
-  });
+  const configAttrMap = new Map(
+    config?.attributes.map(attr => [attr.name, attr]),
+  );
+  const schemaAttrMap = new Map(
+    schema?.attributes.map(attr => [attr.name, attr]),
+  );
+  const allAttrNames = [
+    ...new Set([...configAttrMap.keys(), ...schemaAttrMap.keys()]),
+  ];
+
+  return allAttrNames.map(attrName => ({
+    name: attrName,
+    ...(schemaAttrMap.get(attrName) || {}),
+    ...(configAttrMap.get(attrName) || {}),
+  }));
 };
 
 const mergeVertex = (
@@ -276,15 +284,22 @@ export function getDefaultEdgeTypeConfig(edgeType: EdgeType): EdgeTypeConfig {
 export function patchToRemoveDisplayLabel<
   TypeConfig extends VertexTypeConfig | EdgeTypeConfig,
 >(config: TypeConfig): TypeConfig {
-  const cloned = cloneDeep(config);
+  const { displayLabel: _, ...rest } = config;
 
-  delete cloned.displayLabel;
-
-  // Remove any displayLabel values that were cached in old versions of Graph Explorer
-  for (const attr of cloned.attributes) {
-    // Cast to `any` since the type no longer has `displayLabel` defined
-    delete (attr as any).displayLabel;
-  }
+  return {
+    ...rest,
+    // Remove any displayLabel values that were cached in old versions of Graph Explorer
+    attributes: config.attributes.map(attr => {
+      const { displayLabel: _, ...attrRest } = attr as AttributeConfig & {
+        displayLabel?: string;
+      };
+      return attrRest;
+    }),
+  } as TypeConfig;
+}
 
-  return cloned;
+function toMapByType<T extends { type: string }>(
+  items: T[] | undefined | null,
+): Map<string, T> {
+  return new Map(items?.map(item => [item.type, item]));
 }

packages/graph-explorer/src/core/StateProvider/userPreferences.test.ts

@@ -1,4 +1,5 @@
 // @vitest-environment happy-dom
+import { useAtomValue } from "jotai";
 import { act } from "react";
 
 import { createEdgeType, createVertexType } from "@/core";
@@ -7,9 +8,11 @@ import { DbState, renderHookWithState } from "@/utils/testing";
 import {
   defaultEdgePreferences,
   defaultVertexPreferences,
+  edgePreferencesAtom,
   type EdgePreferencesStorageModel,
   useEdgeStyling,
   useVertexStyling,
+  vertexPreferencesAtom,
   type VertexPreferencesStorageModel,
 } from "./userPreferences";
 
@@ -395,3 +398,65 @@ describe("useDeferredAtom integration", () => {
     );
   });
 });
+
+describe("vertexPreferencesAtom", () => {
+  it("should return stored preferences for a known type", () => {
+    const dbState = new DbState();
+    const vertexType = createVertexType("Person");
+    dbState.addVertexStyle(vertexType, { color: "#ff0000" });
+
+    const { result } = renderHookWithState(
+      () => useAtomValue(vertexPreferencesAtom),
+      dbState,
+    );
+
+    expect(result.current.get(vertexType)).toStrictEqual(
+      createExpectedVertex({ type: vertexType, color: "#ff0000" }),
+    );
+  });
+
+  it("should return defaults for an unknown type", () => {
+    const dbState = new DbState();
+    const vertexType = createVertexType("Unknown");
+
+    const { result } = renderHookWithState(
+      () => useAtomValue(vertexPreferencesAtom),
+      dbState,
+    );
+
+    expect(result.current.get(vertexType)).toStrictEqual(
+      createExpectedVertex({ type: vertexType }),
+    );
+  });
+});
+
+describe("edgePreferencesAtom", () => {
+  it("should return stored preferences for a known type", () => {
+    const dbState = new DbState();
+    const edgeType = createEdgeType("KNOWS");
+    dbState.addEdgeStyle(edgeType, { lineColor: "#00ff00" });
+
+    const { result } = renderHookWithState(
+      () => useAtomValue(edgePreferencesAtom),
+      dbState,
+    );
+
+    expect(result.current.get(edgeType)).toStrictEqual(
+      createExpectedEdge({ type: edgeType, lineColor: "#00ff00" }),
+    );
+  });
+
+  it("should return defaults for an unknown type", () => {
+    const dbState = new DbState();
+    const edgeType = createEdgeType("Unknown");
+
+    const { result } = renderHookWithState(
+      () => useAtomValue(edgePreferencesAtom),
+      dbState,
+    );
+
+    expect(result.current.get(edgeType)).toStrictEqual(
+      createExpectedEdge({ type: edgeType }),
+    );
+  });
+});

packages/graph-explorer/src/core/StateProvider/userPreferences.ts

@@ -158,17 +158,35 @@ export type UserStyling = {
   edges?: Array<EdgePreferencesStorageModel>;
 };
 
-/** Get the stored user preferences for vertices and edges in a fast lookup Map. */
-function useStoredGraphPreferences() {
-  const graphPreferences = useAtomValue(userStylingAtom);
-  const vertices = new Map(
-    graphPreferences.vertices?.map(v => [v.type, v]) ?? [],
+/** Vertex preferences indexed by type for O(1) lookup with default fallback. */
+export const vertexPreferencesAtom = atom(get => {
+  const userStyling = get(userStylingAtom);
+  const lookup = new Map(
+    userStyling.vertices?.map(v => [
+      v.type,
+      createVertexPreference(v.type, v),
+    ]) ?? [],
   );
-  const edges = new Map(graphPreferences.edges?.map(e => [e.type, e]) ?? []);
-  const result = { vertices, edges };
-  const deferredResult = useDeferredValue(result);
-  return deferredResult;
-}
+  return {
+    get(type: VertexType) {
+      return lookup.get(type) ?? createVertexPreference(type);
+    },
+  };
+});
+
+/** Edge preferences indexed by type for O(1) lookup with default fallback. */
+export const edgePreferencesAtom = atom(get => {
+  const userStyling = get(userStylingAtom);
+  const lookup = new Map(
+    userStyling.edges?.map(e => [e.type, createEdgePreference(e.type, e)]) ??
+      [],
+  );
+  return {
+    get(type: EdgeType) {
+      return lookup.get(type) ?? createEdgePreference(type);
+    },
+  };
+});
 
 /** Combines the stored user preferences with the defined default values. */
 export function createVertexPreference(
@@ -196,22 +214,16 @@ export function createEdgePreference(
 
 /** Returns an array of vertex preferences based on the known vertex types in the schema. */
 export function useAllVertexPreferences(): VertexPreferences[] {
-  const { vertices: allPreferences } = useStoredGraphPreferences();
+  const prefs = useAtomValue(vertexPreferencesAtom);
   const { vertices: allSchemas } = useActiveSchema();
-
-  return allSchemas.map(({ type }) =>
-    createVertexPreference(type, allPreferences.get(type)),
-  );
+  return allSchemas.map(({ type }) => prefs.get(type));
 }
 
 /** Returns an array of edge preferences based on the known edge types in the schema. */
 export function useAllEdgePreferences(): EdgePreferences[] {
-  const { edges: allPreferences } = useStoredGraphPreferences();
+  const prefs = useAtomValue(edgePreferencesAtom);
   const { edges: allSchemas } = useActiveSchema();
-
-  return allSchemas.map(({ type }) =>
-    createEdgePreference(type, allPreferences.get(type)),
-  );
+  return allSchemas.map(({ type }) => prefs.get(type));
 }
 
 /** Returns the user preferences for the specified vertex type. */
@@ -228,22 +240,14 @@ export function useEdgePreferences(type: EdgeType): EdgePreferences {
  * Returns the user preferences for the specified vertex type.
  */
 export const vertexPreferenceByTypeAtom = atomFamily((type: VertexType) =>
-  atom(get => {
-    const userStyling = get(userStylingAtom);
-    const stored = userStyling.vertices?.find(v => v.type === type);
-    return createVertexPreference(type, stored);
-  }),
+  atom(get => get(vertexPreferencesAtom).get(type)),
 );
 
 /**
  * Returns the user preferences for the specified edge type.
  */
 export const edgePreferenceByTypeAtom = atomFamily((type: EdgeType) =>
-  atom(get => {
-    const userStyling = get(userStylingAtom);
-    const stored = userStyling.edges?.find(e => e.type === type);
-    return createEdgePreference(type, stored);
-  }),
+  atom(get => get(edgePreferencesAtom).get(type)),
 );
 
 type UpdatedVertexStyle = Partial<Omit<VertexPreferences, "type">>;