improvement(tree): Maintain column/row ID -> node cache for faster lookup

packages/dds/tree/src/tableSchema.ts

@@ -7,7 +7,7 @@ import { fail } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import { EmptyKey } from "./core/index.js";
-import { TreeAlpha } from "./shared-tree/index.js";
+import { Tree, TreeAlpha } from "./shared-tree/index.js";
 import {
 	type FieldHasDefault,
 	type ImplicitAllowedTypes,
@@ -951,6 +951,62 @@ export namespace System_TableSchema {
 				return cell;
 			}
 
+			// #region ID lookup caches
+
+			// Looking up rows/columns by string ID is a hot path (every getCell / setCell / etc. call goes
+			// through #tryGetRow or #tryGetColumn).  Rather than scanning the arrays linearly on every call,
+			// we maintain lazily-built Maps from ID → node.
+			//
+			// Cache invalidation:
+			//   Each cache is marked stale (reset to `undefined`) by a `nodeChanged` listener registered on
+			//   the corresponding array node (this.table.rows / this.table.columns).  `nodeChanged` fires on
+			//   an array node whenever an element is inserted, removed, or moved — exactly the set of
+			//   operations that could change which ID maps to which node.  The event fires after the full
+			//   batch of edits has been applied (including remote edits received from collaborators), so the
+			//   cache is always rebuilt from a consistent, in-schema state.
+			//
+			// Listener lifetime:
+			//   The listener is registered exactly once per cache (guarded by the non-undefined check on the
+			//   stored unsubscribe callback).  Subsequent cache rebuilds after invalidation reuse the same
+			//   listener — no additional subscriptions accumulate.  The unsubscribe callback is stored so
+			//   that explicit cleanup is possible in the future if needed.
+			//
+			// Unhydrated trees:
+			//   `Tree.on` and `nodeChanged` work for unhydrated nodes as well — all mutations go through
+			//   `#applyEditsInBatch`, which wraps them in `withBufferedTreeEvents`, causing `nodeChanged`
+			//   to fire after each batch regardless of whether the node is hydrated.  The invalidation path
+			//   is therefore identical for hydrated and unhydrated tables.
+
+			/**
+			 * Cache from row ID → row node for O(1) lookups in {@link Table.#tryGetRow}.
+			 * `undefined` means the cache is stale and must be rebuilt before use.
+			 */
+			#rowCache: Map<string, RowValueType> | undefined = undefined;
+
+			/**
+			 * Unsubscribe function for the `nodeChanged` listener on `this.table.rows`.
+			 * `undefined` means the listener has not yet been registered (first cache build is pending).
+			 * After the first build, this is always defined and the listener remains active for the
+			 * lifetime of the Table node.
+			 */
+			#rowCacheUnsubscribe: (() => void) | undefined = undefined;
+
+			/**
+			 * Cache from column ID → column node for O(1) lookups in {@link Table.#tryGetColumn}.
+			 * `undefined` means the cache is stale and must be rebuilt before use.
+			 */
+			#columnCache: Map<string, ColumnValueType> | undefined = undefined;
+
+			/**
+			 * Unsubscribe function for the `nodeChanged` listener on `this.table.columns`.
+			 * `undefined` means the listener has not yet been registered (first cache build is pending).
+			 * After the first build, this is always defined and the listener remains active for the
+			 * lifetime of the Table node.
+			 */
+			#columnCacheUnsubscribe: (() => void) | undefined = undefined;
+
+			// #endregion
+
 			/**
 			 * Applies the provided edits in a "batch".
 			 *
@@ -994,6 +1050,40 @@ export namespace System_TableSchema {
 				});
 			}
 
+			/**
+			 * Returns the column ID → column node cache, building and caching it first if stale.
+			 *
+			 * @remarks
+			 * The first time this is called, it builds the map from the current contents of
+			 * `this.table.columns` and registers a `nodeChanged` listener on that array.
+			 * The listener invalidates the cache (sets `#columnCache` to `undefined`) whenever
+			 * the column array is structurally modified (insert / remove / move), so every
+			 * subsequent call that follows a structural change automatically triggers a rebuild.
+			 * The listener is registered only once — it is not re-registered on cache rebuilds.
+			 */
+			#getColumnCache(): Map<string, ColumnValueType> {
+				let cache = this.#columnCache;
+				if (cache === undefined) {
+					cache = new Map<string, ColumnValueType>();
+					for (const column of this.table.columns) {
+						// TypeScript is unable to narrow array element types correctly here.
+						// See: https://github.com/microsoft/TypeScript/issues/52144
+						cache.set((column as ColumnValueType).id, column as ColumnValueType);
+					}
+					this.#columnCache = cache;
+
+					// Register the invalidation listener once. The `nodeChanged` event fires on the
+					// array node itself after any structural change (insert / remove / move), which is
+					// exactly the set of changes that can alter the ID → node mapping.
+					if (this.#columnCacheUnsubscribe === undefined) {
+						this.#columnCacheUnsubscribe = Tree.on(this.columns, "nodeChanged", () => {
+							this.#columnCache = undefined;
+						});
+					}
+				}
+				return cache;
+			}
+
 			/**
 			 * Attempts to resolve the provided Column node or ID to a Column node in the table.
 			 * Returns `undefined` if there is no match.
@@ -1012,12 +1102,7 @@ export namespace System_TableSchema {
 				}
 
 				if (typeof columnOrIdOrIndex === "string") {
-					const columnId = columnOrIdOrIndex;
-					// TypeScript is unable to narrow the types correctly here, hence the casts.
-					// See: https://github.com/microsoft/TypeScript/issues/52144
-					return this.table.columns.find((col) => (col as ColumnValueType).id === columnId) as
-						| ColumnValueType
-						| undefined;
+					return this.#getColumnCache().get(columnOrIdOrIndex);
 				}
 
 				// If the user provided a node, ensure it actually exists in this table.
@@ -1062,6 +1147,40 @@ export namespace System_TableSchema {
 				);
 			}
 
+			/**
+			 * Returns the row ID → row node cache, building and caching it first if stale.
+			 *
+			 * @remarks
+			 * The first time this is called, it builds the map from the current contents of
+			 * `this.table.rows` and registers a `nodeChanged` listener on that array.
+			 * The listener invalidates the cache (sets `#rowCache` to `undefined`) whenever
+			 * the row array is structurally modified (insert / remove / move), so every
+			 * subsequent call that follows a structural change automatically triggers a rebuild.
+			 * The listener is registered only once — it is not re-registered on cache rebuilds.
+			 */
+			#getRowCache(): Map<string, RowValueType> {
+				let cache = this.#rowCache;
+				if (cache === undefined) {
+					cache = new Map<string, RowValueType>();
+					for (const row of this.table.rows) {
+						// TypeScript is unable to narrow array element types correctly here.
+						// See: https://github.com/microsoft/TypeScript/issues/52144
+						cache.set((row as RowValueType).id, row as RowValueType);
+					}
+					this.#rowCache = cache;
+
+					// Register the invalidation listener once. The `nodeChanged` event fires on the
+					// array node itself after any structural change (insert / remove / move), which is
+					// exactly the set of changes that can alter the ID → node mapping.
+					if (this.#rowCacheUnsubscribe === undefined) {
+						this.#rowCacheUnsubscribe = Tree.on(this.rows, "nodeChanged", () => {
+							this.#rowCache = undefined;
+						});
+					}
+				}
+				return cache;
+			}
+
 			/**
 			 * Attempts to resolve the provided Row node or ID to a Row node in the table.
 			 * Returns `undefined` if there is no match.
@@ -1078,12 +1197,7 @@ export namespace System_TableSchema {
 				}
 
 				if (typeof rowOrIdOrIndex === "string") {
-					const rowId = rowOrIdOrIndex;
-					// TypeScript is unable to narrow the types correctly here, hence the casts.
-					// See: https://github.com/microsoft/TypeScript/issues/52144
-					return this.table.rows.find((row) => (row as RowValueType).id === rowId) as
-						| RowValueType
-						| undefined;
+					return this.#getRowCache().get(rowOrIdOrIndex);
 				}
 
 				// If the user provided a node, ensure it actually exists in this table.