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,58 @@ 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.
+			//   TODO: Consider if we should do more fine-grained invalidation here. E.g. look at the specific deltas
+			//   returned by the `nodeChanged` event and only invalidate entries as needed.
+			//
+			// 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.
+
+			/**
+			 * 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 +1046,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 +1098,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 +1143,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 +1193,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.

packages/dds/tree/src/test/tableSchema.spec.ts

@@ -2503,6 +2503,140 @@ describe("TableFactory unit tests", () => {
 				});
 			});
 
+			// These tests exercise ID-based lookups across structural mutations
+			// (insert / remove / reorder) on the columns and rows lists. They are
+			// intended to catch stale-state regressions in the ID → node lookup
+			// path, including the lazy ID caches inside Table.
+			describe("ID lookup after structural changes", () => {
+				it("getColumn by ID resolves a column inserted after a prior lookup", () => {
+					const table = create2x2Table();
+
+					// Prime any internal lookup state with an existing-ID query.
+					assert(table.getColumn("column-0") !== undefined);
+
+					table.insertColumns({
+						index: table.columns.length,
+						columns: [{ id: "column-2", props: {} }],
+					});
+
+					const inserted = table.getColumn("column-2");
+					assert(inserted !== undefined);
+					assert.strictEqual(inserted.id, "column-2");
+					// Existing IDs still resolve correctly.
+					assert.strictEqual(table.getColumn("column-0")?.id, "column-0");
+				});
+
+				it("getColumn by ID returns undefined for a column removed after a prior lookup", () => {
+					const table = create2x2Table();
+
+					// Prime any internal lookup state for both columns.
+					assert(table.getColumn("column-0") !== undefined);
+					assert(table.getColumn("column-1") !== undefined);
+
+					table.removeColumns(["column-0"]);
+
+					assert.strictEqual(table.getColumn("column-0"), undefined);
+					assert.strictEqual(table.getColumn("column-1")?.id, "column-1");
+				});
+
+				it("getColumn by ID resolves the moved node after a column reorder", () => {
+					const table = create2x2Table();
+
+					// Prime any internal lookup state.
+					assert(table.getColumn("column-1") !== undefined);
+
+					// Move column-1 to the start.
+					table.columns.moveToStart(1);
+
+					const after = table.getColumn("column-1");
+					assert(after !== undefined);
+					assert.strictEqual(after.id, "column-1");
+					// The other column still resolves.
+					assert.strictEqual(table.getColumn("column-0")?.id, "column-0");
+					// And the table order reflects the move.
+					assert.deepEqual(
+						[...table.columns].map((c) => c.id),
+						["column-1", "column-0"],
+					);
+				});
+
+				it("getRow by ID resolves a row inserted after a prior lookup", () => {
+					const table = create2x2Table();
+
+					assert(table.getRow("row-0") !== undefined);
+
+					table.insertRows({
+						index: table.rows.length,
+						rows: [{ id: "row-2", cells: {}, props: {} }],
+					});
+
+					const inserted = table.getRow("row-2");
+					assert(inserted !== undefined);
+					assert.strictEqual(inserted.id, "row-2");
+					assert.strictEqual(table.getRow("row-0")?.id, "row-0");
+				});
+
+				it("getRow by ID returns undefined for a row removed after a prior lookup", () => {
+					const table = create2x2Table();
+
+					assert(table.getRow("row-0") !== undefined);
+					assert(table.getRow("row-1") !== undefined);
+
+					table.removeRows(["row-0"]);
+
+					assert.strictEqual(table.getRow("row-0"), undefined);
+					assert.strictEqual(table.getRow("row-1")?.id, "row-1");
+				});
+
+				it("getRow by ID resolves the moved node after a row reorder", () => {
+					const table = create2x2Table();
+
+					assert(table.getRow("row-1") !== undefined);
+
+					table.rows.moveToStart(1);
+
+					const after = table.getRow("row-1");
+					assert(after !== undefined);
+					assert.strictEqual(after.id, "row-1");
+					assert.strictEqual(table.getRow("row-0")?.id, "row-0");
+					assert.deepEqual(
+						[...table.rows].map((r) => r.id),
+						["row-1", "row-0"],
+					);
+				});
+
+				it("setCell by ID targets the correct cell after rows and columns are mutated", () => {
+					const table = create2x2Table();
+
+					// Prime any internal lookup state.
+					assert(table.getRow("row-0") !== undefined);
+					assert(table.getColumn("column-0") !== undefined);
+
+					// Remove an ID and reuse it on a freshly inserted node.
+					// If a stale lookup leaked, setCell would target the removed node
+					// instead of the new one.
+					table.removeRows(["row-0"]);
+					table.insertRows({
+						index: 0,
+						rows: [{ id: "row-0", cells: {}, props: {} }],
+					});
+					table.removeColumns(["column-0"]);
+					table.insertColumns({
+						index: 0,
+						columns: [{ id: "column-0", props: {} }],
+					});
+
+					table.setCell({
+						key: { row: "row-0", column: "column-0" },
+						cell: { value: "0-0" },
+					});
+
+					const cell = table.getCell({ row: "row-0", column: "column-0" });
+					assert(cell !== undefined);
+					assertEqualTrees(cell, { value: "0-0" });
+				});
+			});
+
 			describe("Recursive tables", () => {
 				it("Can create table schema with recursive types", () => {
 					const mySchemaFactory = new SchemaFactoryBeta("test-recursive");