🤖 refactor: deduplicate Map/Record identity stabilization pattern

ammar-agent · ammar-agent · commit 0a9eca72066e · 2025-10-13T11:08:02.000-05:00
Extract common pattern from useUnreadTracking and useWorkspaceAggregators
into a reusable hook with comparator utilities.

Changes:
- Add useStableReference hook to manage identity stabilization
- Add compareMaps, compareRecords, compareArrays utility functions
- Comprehensive tests for all comparator functions
- Refactor useUnreadTracking to use useStableReference + compareMaps
- Refactor useWorkspaceAggregators to use useStableReference + compareRecords

Benefits:
- DRY - single implementation of the stabilization pattern
- Type-safe - generics preserve types throughout
- Testable - comparator functions are independently tested
- Reusable - easy to apply to new cases (e.g., arrays, custom objects)
- Maintainable - clear separation of comparison logic

The pattern maintains reference identity when deep values haven't changed,
preventing unnecessary re-renders in React components that depend on
reference equality.
diff --git a/src/hooks/useStableReference.test.ts b/src/hooks/useStableReference.test.ts
@@ -0,0 +1,138 @@
+/**
+ * Tests for stable reference utilities.
+ *
+ * Note: Hook tests (useStableReference) are omitted because they require jsdom setup.
+ * The comparator functions are the critical logic and are thoroughly tested here.
+ * The hook itself is a thin wrapper around useMemo and useRef with manual testing.
+ */
+import { compareMaps, compareRecords, compareArrays } from "./useStableReference";
+
+describe("compareMaps", () => {
+  it("returns true for empty maps", () => {
+    expect(compareMaps(new Map(), new Map())).toBe(true);
+  });
+
+  it("returns true for maps with same entries", () => {
+    const prev = new Map([
+      ["a", 1],
+      ["b", 2],
+    ]);
+    const next = new Map([
+      ["a", 1],
+      ["b", 2],
+    ]);
+    expect(compareMaps(prev, next)).toBe(true);
+  });
+
+  it("returns false for maps with different sizes", () => {
+    const prev = new Map([["a", 1]]);
+    const next = new Map([
+      ["a", 1],
+      ["b", 2],
+    ]);
+    expect(compareMaps(prev, next)).toBe(false);
+  });
+
+  it("returns false for maps with different keys", () => {
+    const prev = new Map([["a", 1]]);
+    const next = new Map([["b", 1]]);
+    expect(compareMaps(prev, next)).toBe(false);
+  });
+
+  it("returns false for maps with different values", () => {
+    const prev = new Map([["a", 1]]);
+    const next = new Map([["a", 2]]);
+    expect(compareMaps(prev, next)).toBe(false);
+  });
+
+  it("supports custom value equality function", () => {
+    const prev = new Map([["a", { id: 1 }]]);
+    const next = new Map([["a", { id: 1 }]]);
+
+    // Default comparison (reference equality) returns false
+    expect(compareMaps(prev, next)).toBe(false);
+
+    // Custom comparison (by id) returns true
+    expect(compareMaps(prev, next, (a, b) => a.id === b.id)).toBe(true);
+  });
+});
+
+describe("compareRecords", () => {
+  it("returns true for empty records", () => {
+    expect(compareRecords({}, {})).toBe(true);
+  });
+
+  it("returns true for records with same entries", () => {
+    const prev = { a: 1, b: 2 };
+    const next = { a: 1, b: 2 };
+    expect(compareRecords(prev, next)).toBe(true);
+  });
+
+  it("returns false for records with different sizes", () => {
+    const prev = { a: 1 };
+    const next = { a: 1, b: 2 };
+    expect(compareRecords(prev, next)).toBe(false);
+  });
+
+  it("returns false for records with different keys", () => {
+    const prev = { a: 1 };
+    const next = { b: 1 };
+    expect(compareRecords(prev, next)).toBe(false);
+  });
+
+  it("returns false for records with different values", () => {
+    const prev = { a: 1 };
+    const next = { a: 2 };
+    expect(compareRecords(prev, next)).toBe(false);
+  });
+
+  it("supports custom value equality function", () => {
+    const prev = { a: { id: 1 } };
+    const next = { a: { id: 1 } };
+
+    // Default comparison (reference equality) returns false
+    expect(compareRecords(prev, next)).toBe(false);
+
+    // Custom comparison (by id) returns true
+    expect(compareRecords(prev, next, (a, b) => a.id === b.id)).toBe(true);
+  });
+});
+
+describe("compareArrays", () => {
+  it("returns true for empty arrays", () => {
+    expect(compareArrays([], [])).toBe(true);
+  });
+
+  it("returns true for arrays with same elements", () => {
+    expect(compareArrays([1, 2, 3], [1, 2, 3])).toBe(true);
+  });
+
+  it("returns false for arrays with different lengths", () => {
+    expect(compareArrays([1, 2], [1, 2, 3])).toBe(false);
+  });
+
+  it("returns false for arrays with different values", () => {
+    expect(compareArrays([1, 2, 3], [1, 2, 4])).toBe(false);
+  });
+
+  it("returns false for arrays with same values in different order", () => {
+    expect(compareArrays([1, 2, 3], [3, 2, 1])).toBe(false);
+  });
+
+  it("supports custom value equality function", () => {
+    const prev = [{ id: 1 }, { id: 2 }];
+    const next = [{ id: 1 }, { id: 2 }];
+
+    // Default comparison (reference equality) returns false
+    expect(compareArrays(prev, next)).toBe(false);
+
+    // Custom comparison (by id) returns true
+    expect(compareArrays(prev, next, (a, b) => a.id === b.id)).toBe(true);
+  });
+});
+
+// Hook integration tests would require jsdom setup with bun.
+// The comparator functions above are the critical logic and are thoroughly tested.
+// The hook itself is tested manually through its usage in useUnreadTracking and
+// useWorkspaceAggregators.
+
diff --git a/src/hooks/useStableReference.ts b/src/hooks/useStableReference.ts
@@ -0,0 +1,134 @@
+import { useRef, useMemo, type DependencyList } from "react";
+
+/**
+ * Compare two Maps for deep equality (same keys and values).
+ * Uses === for value comparison by default.
+ *
+ * @param prev Previous Map
+ * @param next Next Map
+ * @param valueEquals Optional custom value equality function
+ * @returns true if Maps are equal, false otherwise
+ */
+export function compareMaps<K, V>(
+  prev: Map<K, V>,
+  next: Map<K, V>,
+  valueEquals: (a: V, b: V) => boolean = (a, b) => a === b
+): boolean {
+  if (prev.size !== next.size) return false;
+
+  for (const [key, value] of next) {
+    if (!prev.has(key)) return false;
+    if (!valueEquals(prev.get(key)!, value)) return false;
+  }
+
+  return true;
+}
+
+/**
+ * Compare two Records for deep equality (same keys and values).
+ * Uses === for value comparison by default.
+ *
+ * @param prev Previous Record
+ * @param next Next Record
+ * @param valueEquals Optional custom value equality function
+ * @returns true if Records are equal, false otherwise
+ */
+export function compareRecords<V>(
+  prev: Record<string, V>,
+  next: Record<string, V>,
+  valueEquals: (a: V, b: V) => boolean = (a, b) => a === b
+): boolean {
+  const prevKeys = Object.keys(prev);
+  const nextKeys = Object.keys(next);
+
+  if (prevKeys.length !== nextKeys.length) return false;
+
+  for (const key of nextKeys) {
+    if (!(key in prev)) return false;
+    if (!valueEquals(prev[key], next[key])) return false;
+  }
+
+  return true;
+}
+
+/**
+ * Compare two Arrays for deep equality (same length and values).
+ * Uses === for value comparison by default.
+ *
+ * @param prev Previous Array
+ * @param next Next Array
+ * @param valueEquals Optional custom value equality function
+ * @returns true if Arrays are equal, false otherwise
+ */
+export function compareArrays<V>(
+  prev: V[],
+  next: V[],
+  valueEquals: (a: V, b: V) => boolean = (a, b) => a === b
+): boolean {
+  if (prev.length !== next.length) return false;
+
+  for (let i = 0; i < next.length; i++) {
+    if (!valueEquals(prev[i], next[i])) return false;
+  }
+
+  return true;
+}
+
+/**
+ * Hook to stabilize reference identity for computed values.
+ *
+ * Returns the previous reference if the new value is deeply equal to the previous value,
+ * preventing unnecessary re-renders in components that depend on reference equality.
+ *
+ * Common use case: Stabilizing Map/Record/Array identities in useMemo when the
+ * underlying values haven't actually changed.
+ *
+ * @example
+ * ```typescript
+ * // Stabilize a Map<string, boolean>
+ * const unreadStatus = useStableReference(
+ *   () => {
+ *     const map = new Map<string, boolean>();
+ *     for (const [id, state] of workspaceStates) {
+ *       map.set(id, calculateUnread(state));
+ *     }
+ *     return map;
+ *   },
+ *   compareMaps,
+ *   [workspaceStates]
+ * );
+ * ```
+ *
+ * @param factory Function that creates the new value
+ * @param comparator Function to check equality between prev and next values
+ * @param deps Dependency list for useMemo
+ * @returns Stable reference to the value
+ */
+export function useStableReference<T>(
+  factory: () => T,
+  comparator: (prev: T, next: T) => boolean,
+  deps: DependencyList
+): T {
+  const ref = useRef<T | undefined>(undefined);
+
+  return useMemo(() => {
+    const next = factory();
+
+    // First render or no previous value
+    if (ref.current === undefined) {
+      ref.current = next;
+      return next;
+    }
+
+    // Compare with previous value
+    if (comparator(ref.current, next)) {
+      return ref.current; // Maintain identity
+    }
+
+    // Value changed, update ref and return new value
+    ref.current = next;
+    return next;
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, deps);
+}
+
diff --git a/src/hooks/useUnreadTracking.ts b/src/hooks/useUnreadTracking.ts
@@ -1,7 +1,8 @@
-import { useEffect, useCallback, useMemo, useRef } from "react";
+import { useEffect, useCallback, useRef } from "react";
 import type { WorkspaceSelection } from "@/components/ProjectSidebar";
 import type { WorkspaceState } from "./useWorkspaceAggregators";
 import { usePersistedState } from "./usePersistedState";
+import { useStableReference, compareMaps } from "./useStableReference";
 
 /**
  * Hook to track unread message status for all workspaces.
@@ -73,47 +74,35 @@ export function useUnreadTracking(
   }, [selectedWorkspace?.workspaceId, workspaceStates, markAsRead]);
 
   // Calculate unread status for all workspaces
-  const unreadStatusRef = useRef<Map<string, boolean>>(new Map());
-  const unreadStatus = useMemo(() => {
-    const next = new Map<string, boolean>();
-
-    for (const [workspaceId, state] of workspaceStates) {
-      // Streaming workspaces are never unread
-      if (state.canInterrupt) {
-        next.set(workspaceId, false);
-        continue;
-      }
-
-      // Check for any assistant-originated content newer than last-read timestamp:
-      // assistant text, tool calls/results (e.g., propose_plan), reasoning, and errors.
-      // Exclude user's own messages and UI markers.
-      const lastRead = lastReadMap[workspaceId] ?? 0;
-      const hasUnread = state.messages.some(
-        (msg) =>
-          msg.type !== "user" && msg.type !== "history-hidden" && (msg.timestamp ?? 0) > lastRead
-      );
+  // Use stable reference to prevent unnecessary re-renders when values haven't changed
+  const unreadStatus = useStableReference(
+    () => {
+      const map = new Map<string, boolean>();
+
+      for (const [workspaceId, state] of workspaceStates) {
+        // Streaming workspaces are never unread
+        if (state.canInterrupt) {
+          map.set(workspaceId, false);
+          continue;
+        }
 
-      next.set(workspaceId, hasUnread);
-    }
+        // Check for any assistant-originated content newer than last-read timestamp:
+        // assistant text, tool calls/results (e.g., propose_plan), reasoning, and errors.
+        // Exclude user's own messages and UI markers.
+        const lastRead = lastReadMap[workspaceId] ?? 0;
+        const hasUnread = state.messages.some(
+          (msg) =>
+            msg.type !== "user" && msg.type !== "history-hidden" && (msg.timestamp ?? 0) > lastRead
+        );
 
-    // Return previous Map reference if nothing actually changed to keep identity stable
-    const prev = unreadStatusRef.current;
-    if (prev.size === next.size) {
-      let same = true;
-      for (const [k, v] of next) {
-        if (prev.get(k) !== v) {
-          same = false;
-          break;
-        }
-      }
-      if (same) {
-        return prev;
+        map.set(workspaceId, hasUnread);
       }
-    }
 
-    unreadStatusRef.current = next;
-    return next;
-  }, [workspaceStates, lastReadMap]);
+      return map;
+    },
+    compareMaps,
+    [workspaceStates, lastReadMap]
+  );
 
   // Manual toggle function for clicking the indicator
   const toggleUnread = useCallback(
diff --git a/src/hooks/useWorkspaceAggregators.ts b/src/hooks/useWorkspaceAggregators.ts
