fix(tree): prevent revert of commits that were rolled back

packages/dds/tree/src/shared-tree/treeCheckout.ts

@@ -486,9 +486,9 @@ export class TreeCheckout implements ITreeCheckout {
 	private readonly views = new Set<TreeView<ImplicitFieldSchema>>();
 
 	/**
-	 * Set of revertibles maintained for automatic disposal
+	 * Revertibles maintained for automatic disposal
 	 */
-	private readonly revertibles = new Set<RevertibleAlpha>();
+	private readonly revertibles = new Map<RevisionTag, RevertibleAlpha>();
 
 	/**
 	 * Each branch's head commit corresponds to a revertible commit.
@@ -754,7 +754,7 @@ export class TreeCheckout implements ITreeCheckout {
 								revision,
 								this.#transaction.activeBranch.fork(commit),
 							);
-							this.revertibles.add(revertible);
+							this.revertibles.set(revision, revertible);
 							return revertible;
 						};
 
@@ -791,6 +791,11 @@ export class TreeCheckout implements ITreeCheckout {
 				this.#events.emit("changed", metadata, getRevertible);
 				withinEventContext = false;
 			}
+		} else if (event.type === "remove") {
+			// Commits that are rolled back should no longer be revertible
+			for (const commit of event.removedCommits) {
+				this.revertibles.get(commit.revision)?.dispose();
+			}
 		} else if (this.isRemoteChangeEvent(event)) {
 			// TODO: figure out how to plumb through commit kind info for remote changes
 			this.#events.emit("changed", {
@@ -1097,7 +1102,7 @@ export class TreeCheckout implements ITreeCheckout {
 						"Unable to dispose a revertible that has already been disposed.",
 					);
 				}
-				checkout.disposeRevertible(revertible, revision);
+				checkout.disposeRevertible(revision);
 				onRevertibleDisposed?.(revertible);
 			},
 		};
@@ -1339,15 +1344,15 @@ export class TreeCheckout implements ITreeCheckout {
 	}
 
 	private purgeRevertibles(): void {
-		for (const revertible of this.revertibles) {
+		for (const revertible of this.revertibles.values()) {
 			revertible.dispose();
 		}
 	}
 
-	private disposeRevertible(revertible: RevertibleAlpha, revision: RevisionTag): void {
+	private disposeRevertible(revision: RevisionTag): void {
 		this.revertibleCommitBranches.get(revision)?.dispose();
 		this.revertibleCommitBranches.delete(revision);
-		this.revertibles.delete(revertible);
+		this.revertibles.delete(revision);
 	}
 
 	private revertRevertible(

packages/dds/tree/src/test/shared-tree/treeCheckout.spec.ts

@@ -1848,6 +1848,27 @@ describe("sharedTreeView", () => {
 			fork1.dispose();
 		});
 	});
+
+	it("Rolled back commits cannot be reverted", () => {
+		const sf = new SchemaFactory("Enrichment Schema");
+		class Node extends sf.object("Node", { id: sf.string }) {}
+		const NodeArray = sf.array(Node);
+		const provider = new TestTreeProviderLite(1);
+		const config = new TreeViewConfiguration({ schema: NodeArray, enableSchemaValidation });
+		const view = provider.trees[0].kernel.viewWith(config);
+		view.initialize([]);
+		const init = view.checkout.mainBranch.getHead();
+
+		const { undoStack, unsubscribe } = createTestUndoRedoStacks(view.events);
+		view.root.insertAtEnd({ id: "A" });
+		view.root[0].id = "B";
+		assert.equal(undoStack.length, 2);
+
+		view.checkout.mainBranch.removeAfter(init);
+		assert.equal(view.root.length, 0);
+		assert.equal(undoStack.length, 0);
+		unsubscribe();
+	});
 });
 
 const defaultSf = new SchemaFactory("Checkout and view test schema");

packages/dds/tree/src/test/utils.ts

@@ -119,6 +119,7 @@ import {
 	type FieldKindIdentifier,
 	type TreeNodeSchemaIdentifier,
 	type TreeFieldStoredSchema,
+	RevertibleStatus,
 } from "../core/index.js";
 import { FormatValidatorBasic } from "../external-utilities/index.js";
 import {
@@ -1347,10 +1348,14 @@ export function createTestUndoRedoStacks(
 	const unsubscribe = (): void => {
 		unsubscribeFromChangedEvent();
 		for (const revertible of undoStack) {
-			revertible.dispose();
+			if (revertible.status !== RevertibleStatus.Disposed) {
+				revertible.dispose();
+			}
 		}
 		for (const revertible of redoStack) {
-			revertible.dispose();
+			if (revertible.status !== RevertibleStatus.Disposed) {
+				revertible.dispose();
+			}
 		}
 	};
 	return { undoStack, redoStack, unsubscribe };