refactor(tree): Misc. column/row insert/remove improvements

packages/dds/tree/src/tableSchema.ts

@@ -3,7 +3,7 @@
  * Licensed under the MIT License.
  */
 
-import { assert, fail } from "@fluidframework/core-utils/internal";
+import { fail, Heap, oob } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import { Tree, TreeAlpha } from "./shared-tree/index.js";
@@ -550,6 +550,21 @@ export namespace System_TableSchema {
 		TCell extends ImplicitAllowedTypes = ImplicitAllowedTypes,
 	> = OptionsWithSchemaFactory<TSchemaFactory> & OptionsWithCellSchema<TCell>;
 
+	/**
+	 * An index range.
+	 */
+	interface IndexRange {
+		/**
+		 * Start of the range (inclusive).
+		 */
+		start: number;
+
+		/**
+		 * The number of items.
+		 */
+		count: number;
+	}
+
 	/**
 	 * Factory for creating table schema.
 	 * @system @alpha
@@ -759,42 +774,55 @@ export namespace System_TableSchema {
 						return [];
 					}
 
-					// Ensure the specified columns exists before starting transaction.
-					// Improves user-facing error experience.
-					for (const columnToRemove of columnsToRemove) {
-						this._assertContainsColumn(columnToRemove);
+					// To reduce the number of operations we need to perform, we will build up a series of
+					// contiguous ranges to be removed from the input list and remove each range as an operation
+					// (rather than each column individually).
+					const columnIndicesToRemove: number[] = [];
+					for (const column of columnsToRemove) {
+						this._assertContainsColumn(column);
+						const columnIndex =
+							this._getColumnIndex(column) ?? fail("Expected column to exist.");
+						columnIndicesToRemove.push(columnIndex);
 					}
 
+					const rangesToRemove = Table._groupIndicesIntoRanges(columnIndicesToRemove);
+
 					const removedColumns: ColumnValueType[] = [];
 					this._applyEditsInBatch(() => {
-						// Note, throwing an error within a transaction will abort the entire transaction.
-						// So if we throw an error here for any column, no columns will be removed.
-						for (const columnToRemove of columnsToRemove) {
-							const removedColumn = this._removeColumn(columnToRemove);
-							removedColumns.push(removedColumn);
+						// #region Remove columns in batches of contiguous ranges
+
+						// Note: Each range we remove from the list impacts the indices of the remaining items.
+						// As a result, we need to adjust each subsequent range we remove to account for the changes.
+						// `rangesToRemove` is sorted in ascending order, so it is sufficient to just track and modify based on a single offset.
+						let offset = 0;
+						for (const range of rangesToRemove) {
+							const removed = Table._removeRange(
+								{
+									index: range.start - offset,
+									count: range.count,
+								},
+								this.columns,
+							);
+							removedColumns.push(...removed);
+							offset = offset + range.count;
 						}
-					});
-					return removedColumns;
-				}
-			}
 
-			private _removeColumn(columnOrId: string | ColumnValueType): ColumnValueType {
-				const column = this._getColumn(columnOrId);
-				const index = this.columns.indexOf(column);
-				assert(index !== -1, "Column should exist");
+						// #endregion
 
-				this._applyEditsInBatch(() => {
-					// Remove the corresponding cell from all rows.
-					for (const row of this.rows) {
-						// TypeScript is unable to narrow the row type correctly here, hence the cast.
-						// See: https://github.com/microsoft/TypeScript/issues/52144
-						(row as RowValueType).removeCell(column);
-					}
+						// #region Remove cells corresponding with removed columns from each row
 
-					this.columns.removeAt(index);
-				});
+						for (const row of this.rows) {
+							for (const column of removedColumns) {
+								// TypeScript is unable to narrow the row type correctly here, hence the cast.
+								// See: https://github.com/microsoft/TypeScript/issues/52144
+								(row as RowValueType).removeCell(column);
+							}
+						}
 
-				return column;
+						// #endregion
+					});
+					return removedColumns;
+				}
 			}
 
 			public removeRows(
@@ -826,34 +854,41 @@ export namespace System_TableSchema {
 					return [];
 				}
 
-				// Ensure the specified rows exists before starting transaction.
-				// Improves user-facing error experience.
-				for (const rowToRemove of rowsToRemove) {
-					this._assertContainsRow(rowToRemove);
+				// To reduce the number of operations we need to perform, we will build up a series of
+				// contiguous ranges to be removed from the input list and remove each range as an operation
+				// (rather than each row individually).
+				const rowIndicesToRemove: number[] = [];
+				for (const row of rowsToRemove) {
+					this._assertContainsRow(row);
+					const rowIndex = this._getRowIndex(row) ?? fail("Expected row to exist.");
+					rowIndicesToRemove.push(rowIndex);
 				}
 
+				const rangesToRemove = Table._groupIndicesIntoRanges(rowIndicesToRemove);
+
 				const removedRows: RowValueType[] = [];
 				this._applyEditsInBatch(() => {
-					// Note, throwing an error within a transaction will abort the entire transaction.
-					// So if we throw an error here for any row, no rows will be removed.
-					for (const rowToRemove of rowsToRemove) {
-						const removedRow = this._removeRow(rowToRemove);
-						removedRows.push(removedRow);
+					// Note: Each range we remove from the list impacts the indices of the remaining items.
+					// As a result, we need to adjust each subsequent range we remove to account for the changes.
+					// `rangesToRemove` is sorted in ascending order, so it is sufficient to just track and modify based on a single offset.
+					let offset = 0;
+					for (const range of rangesToRemove) {
+						// Note, throwing an error within a transaction will abort the entire transaction.
+						// So if we throw an error here for any column, no columns will be removed.
+						const removed = Table._removeRange(
+							{
+								index: range.start - offset,
+								count: range.count,
+							},
+							this.rows,
+						);
+						removedRows.push(...removed);
+						offset = offset + range.count;
 					}
 				});
 				return removedRows;
 			}
 
-			private _removeRow(rowOrId: string | RowValueType): RowValueType {
-				const rowToRemove = this._getRow(rowOrId);
-
-				const index = this.rows.indexOf(rowToRemove);
-				assert(index !== -1, "Row should exist");
-
-				this.rows.removeAt(index);
-				return rowToRemove;
-			}
-
 			public removeCell(
 				key: TableSchema.CellKey<TColumnSchema, TRowSchema>,
 			): CellValueType | undefined {
@@ -1078,7 +1113,7 @@ export namespace System_TableSchema {
 				const { index, count } = range;
 				const end = index + count; // exclusive
 
-				// TypeScript is unable to narrow the array element type correctly here, hence the cast.
+				// TypeScript is unable to narrow the row type correctly here, hence the cast.
 				// See: https://github.com/microsoft/TypeScript/issues/52144
 				const removedRows = array.slice(
 					index,
@@ -1089,6 +1124,61 @@ export namespace System_TableSchema {
 				return removedRows;
 			}
 
+			/**
+			 * Groups the provided list of indices into a list of ordered, contiguous ranges.
+			 */
+			private static _groupIndicesIntoRanges(indices: readonly number[]): IndexRange[] {
+				// We will leverage a set for unique index lookup.
+				const indexSet = new Set(indices);
+
+				// Tracks which indices have already been visited by the walk to ensure
+				// indices are not double-counted, and to ensure the algorithm remains O(n).
+				const visited = new Set<number>();
+
+				// Tracks the set of disjoint index ranges, sorted in ascending order.
+				const sortedRanges = new Heap<IndexRange>({
+					compare: (a, b) => a.start - b.start,
+					min: {
+						start: Number.MIN_VALUE,
+						// Note: the count here is arbitrary - it isn't used by the comparison function,
+						// but the Heap type requires that "min" be of the same type as the elements being inserted.
+						count: Number.NaN,
+					},
+				});
+
+				for (const index of indices) {
+					if (visited.has(index)) {
+						continue;
+					}
+
+					let start = index; // inclusive
+					let end = index; // inclusive
+
+					// Expand backward
+					while (indexSet.has(start - 1)) {
+						start--;
+					}
+
+					// Expand forward
+					while (indexSet.has(end + 1)) {
+						end++;
+					}
+
+					// Mark all indices in this range as visited so we don't consider them again.
+					for (let i = start; i <= end; i++) {
+						visited.add(i);
+					}
+
+					sortedRanges.add({ start, count: end - start + 1 });
+				}
+
+				const result: IndexRange[] = [];
+				while (sortedRanges.peek() !== undefined) {
+					result.push(sortedRanges.get() ?? oob());
+				}
+				return result;
+			}
+
 			/**
 			 * Ensure that the specified index is a valid location for item insertion in the destination list.
 			 * @throws Throws a usage error if the destination is invalid.

packages/dds/tree/src/test/memory/tableTree.spec.ts

@@ -389,7 +389,7 @@ describe("SharedTree table APIs memory usage", () => {
 					},
 				});
 
-				// Test the memory usage of the SharedTree for undoing the insertion of a column and a row in the middle for a given number of times.
+				// Test the memory usage of undoing the insertion of a column and a row in the middle N times.
 				runBenchmark({
 					title: `Undo: ${scenarioName}`,
 					tableSize,
@@ -703,7 +703,7 @@ describe("SharedTree table APIs memory usage", () => {
 					},
 				});
 
-				// Test the memory usage of the SharedTree for undoing the removal of a column and a row in the middle for a given number of times.
+				// Test the memory usage of undoing the removal of a column and a row in the middle N times.
 				runBenchmark({
 					title: `Undo: ${scenarioName}`,
 					tableSize,
@@ -880,7 +880,7 @@ describe("SharedTree table APIs memory usage", () => {
 					},
 				});
 
-				// Test the memory usage of the SharedTree for undoing the setting of a cell value in the middle for a given number of times.
+				// Test the memory usage of undoing the setting of a cell value in the middle N times.
 				runBenchmark({
 					title: `Undo: ${scenarioName}`,
 					tableSize,
@@ -906,7 +906,7 @@ describe("SharedTree table APIs memory usage", () => {
 					},
 				});
 
-				// Test the memory usage of the SharedTree for redoing the setting of a cell value in the middle for a given number of times.
+				// Test the memory usage of redoing the setting of a cell value in the middle N times.
 				runBenchmark({
 					title: `Redo: ${scenarioName}`,
 					tableSize,

packages/dds/tree/src/test/shared-tree/tableTree.bench.ts

@@ -395,7 +395,7 @@ describe("SharedTree table APIs execution time", () => {
 					maxBenchmarkDurationSeconds,
 				});
 
-				// Test the execute time of undoing insert a row and a column in the middle for a given number of times.
+				// Test the execution time of undoing insert a row and a column in the middle of the table N times.
 				runBenchmark({
 					title: `Undo: ${scenarioName}`,
 					tableSize,
@@ -427,7 +427,7 @@ describe("SharedTree table APIs execution time", () => {
 				});
 
 				// TODO: AB#43364: Enable these tests back after allowing SharedTree to support undo/redo for removing cells when a column is removed.
-				// Test the execute time of the SharedTree for redoing a remove row and a column in the middle for a given number of times.
+				// Test the execution time of redoing a remove row and a column in the middle of the table N times.
 				// runBenchmark({
 				// 	title: `Redo: ${scenarioName}`,
 				// 	tableSize,
@@ -473,12 +473,14 @@ describe("SharedTree table APIs execution time", () => {
 					operation: (table) => {
 						for (let i = 0; i < count; i++) {
 							const column = new Column({});
-							const row = new Row({ cells: {} });
 							table.insertColumns({
 								index: Math.floor(table.columns.length / 2),
 								columns: [column],
 							});
+
+							const row = new Row({ cells: {} });
 							table.insertRows({ index: Math.floor(table.rows.length / 2), rows: [row] });
+
 							table.removeColumns([column]);
 							table.removeRows([row]);
 						}
@@ -487,7 +489,7 @@ describe("SharedTree table APIs execution time", () => {
 				});
 
 				// TODO: AB#43364: Enable these tests back after allowing SharedTree to support undo/redo for removing cells when a column is removed.
-				// Test the execute time of undoing insert a row and a column and removing them right away for a given number of times.
+				// Test the execution time of undoing insert a row and a column and removing them right away N times.
 				// runBenchmark({
 				// 	title: `Undo: ${scenarioName}`,
 				// 	tableSize,
@@ -522,7 +524,7 @@ describe("SharedTree table APIs execution time", () => {
 				// });
 
 				// TODO: AB#43364: Enable these tests back after allowing SharedTree to support undo/redo for removing cells when a column is removed.
-				// Test the execute time of the SharedTree for redoing an insert row and a column and removing them right away for a given number of times.
+				// Test the execution time of redoing an insert row and a column and removing them right away N times.
 				// runBenchmark({
 				// 	title: `Redo: ${scenarioName}`,
 				// 	tableSize,
@@ -833,18 +835,18 @@ describe("SharedTree table APIs execution time", () => {
 				});
 
 				// TODO: AB#43364: Enable these tests back after allowing SharedTree to support undo/redo for removing cells when a column is removed.
-				// Test the execute time of undoing insert a row and a column and removing them right away for a given number of times.
-				// Test the execute time of undoing remove a row and a column in the middle for a given number of times.
+				// Test the execution time of undoing insert a row and a column and removing them right away N times.
+				// Test the execution time of undoing remove a row and a column in the middle of the table N times.
 				// runBenchmark({
 				// 	title: `Undo: ${scenarioName}`,
 				// 	tableSize,
 				// 	initialCellValue,
 				// 	beforeOperation: (table, undoRedoManager) => {
 				// 		for (let i = 0; i < count; i++) {
 				// 			const column = table.columns[Math.floor(table.columns.length / 2)];
-				// 			table.removeColumn(column);
+				// 			table.removeColumns([column]);
 				// 			const row = table.rows[Math.floor(table.rows.length / 2)];
-				// 			table.removeRow(row);
+				// 			table.removeRows([row]);
 				// 		}
 				// 		assert(undoRedoManager.canUndo);
 				// 	},
@@ -863,17 +865,17 @@ describe("SharedTree table APIs execution time", () => {
 				// });
 
 				// TODO: AB#43364: Enable these tests back after allowing SharedTree to support undo/redo for removing cells when a column is removed.
-				// Test the execute time of the SharedTree for redoing a remove row and a column in the middle for a given number of times.
+				// Test the execution time of redoing a remove row and a column in the middle of the table N times.
 				// runBenchmark({
 				// 	title: `Redo: ${scenarioName}`,
 				// 	tableSize,
 				// 	initialCellValue,
 				// 	beforeOperation: (table, undoRedoManager) => {
 				// 		for (let i = 0; i < count; i++) {
 				// 			const column = table.columns[Math.floor(table.columns.length / 2)];
-				// 			table.removeColumn(column);
+				// 			table.removeColumns([column]);
 				// 			const row = table.rows[Math.floor(table.rows.length / 2)];
-				// 			table.removeRow(row);
+				// 			table.removeRows([row]);
 				// 		}
 				// 		for (let i = 0; i < count; i++) {
 				// 			undoRedoManager.undo();
@@ -920,7 +922,7 @@ describe("SharedTree table APIs execution time", () => {
 					maxBenchmarkDurationSeconds,
 				});
 
-				// Test the execute time of undoing set a cell value for a given number of times.
+				// Test the execution time of undoing set a cell value N times.
 				runBenchmark({
 					title: `Undo: ${scenarioName}`,
 					tableSize,
@@ -950,7 +952,7 @@ describe("SharedTree table APIs execution time", () => {
 					maxBenchmarkDurationSeconds,
 				});
 
-				// Test the execute time of the SharedTree for redoing a set cell value for a given number of times.
+				// Test the execution time of redoing a set cell value N times.
 				runBenchmark({
 					title: `Redo: ${scenarioName}`,
 					tableSize,