tree(feat): Revertible lifetime management (microsoft#19153)

packages/dds/tree/api-report/tree.api.md

@@ -205,7 +205,8 @@ export type ChangesetLocalId = Brand<number, "ChangesetLocalId">;
 // @internal
 export interface CheckoutEvents {
     afterBatch(): void;
-    revertible(revertible: Revertible): void;
+    newRevertible(revertible: Revertible): void;
+    revertibleDisposed(revertible: Revertible): void;
 }
 
 // @internal
@@ -388,12 +389,6 @@ export type DetachedPlaceUpPath = Brand<Omit<PlaceUpPath, "parent">, "DetachedRa
 // @internal
 export type DetachedRangeUpPath = Brand<Omit<RangeUpPath, "parent">, "DetachedRangeUpPath">;
 
-// @internal
-export enum DiscardResult {
-    Failure = 1,
-    Success = 0
-}
-
 // @public
 export const disposeSymbol: unique symbol;
 
@@ -1357,12 +1352,14 @@ export type RestrictiveReadonlyRecord<K extends symbol | string, T> = {
 
 // @internal
 export interface Revertible {
-    discard(): DiscardResult;
+    discard(): RevertibleResult;
     readonly kind: RevertibleKind;
     readonly origin: {
         readonly isLocal: boolean;
     };
-    revert(): RevertResult;
+    retain(): RevertibleResult;
+    revert(): RevertibleResult;
+    readonly status: RevertibleStatus;
 }
 
 // @internal
@@ -1374,11 +1371,17 @@ export enum RevertibleKind {
 }
 
 // @internal
-export enum RevertResult {
+export enum RevertibleResult {
     Failure = 1,
     Success = 0
 }
 
+// @internal
+export enum RevertibleStatus {
+    Disposed = 1,
+    Valid = 0
+}
+
 // @internal
 export type RevisionTag = SessionSpaceCompressedId | "root";
 

packages/dds/tree/src/core/index.ts

@@ -192,4 +192,9 @@ export {
 	AllowedUpdateType,
 } from "./schema-view/index.js";
 
-export { Revertible, RevertibleKind, RevertResult, DiscardResult } from "./revertible/index.js";
+export {
+	Revertible,
+	RevertibleKind,
+	RevertibleStatus,
+	RevertibleResult,
+} from "./revertible/index.js";

packages/dds/tree/src/core/revertible/index.ts

@@ -3,4 +3,4 @@
  * Licensed under the MIT License.
  */
 
-export { Revertible, RevertibleKind, RevertResult, DiscardResult } from "./revertible.js";
+export { Revertible, RevertibleKind, RevertibleStatus, RevertibleResult } from "./revertible.js";

packages/dds/tree/src/core/revertible/revertible.ts

@@ -22,13 +22,22 @@ export interface Revertible {
 		readonly isLocal: boolean;
 	};
 	/**
-	 * Can be called in order to revert a change. A successful revert will automatically discard resources.
+	 * The current status of the revertible.
 	 */
-	revert(): RevertResult;
+	readonly status: RevertibleStatus;
 	/**
-	 * Should be called to garbage collect any resources associated with the revertible.
+	 * Reverts the associated change and decrements the reference count of the revertible.
 	 */
-	discard(): DiscardResult;
+	revert(): RevertibleResult;
+	/**
+	 * Increments the reference count of the revertible.
+	 * Should be called to prevent/delay the garbage collection of the resources associated with this revertible.
+	 */
+	retain(): RevertibleResult;
+	/**
+	 * Decrements the reference count of the revertible.
+	 */
+	discard(): RevertibleResult;
 }
 
 /**
@@ -51,25 +60,25 @@ export enum RevertibleKind {
 }
 
 /**
- * The result of a revert operation.
+ * The status of a {@link Revertible}.
  *
  * @internal
  */
-export enum RevertResult {
-	/** The revert was successful. */
-	Success,
-	/** The revert failed. */
-	Failure,
+export enum RevertibleStatus {
+	/** The revertible can be reverted. */
+	Valid,
+	/** The revertible has been disposed. Reverting it will have no effect. */
+	Disposed,
 }
 
 /**
- * The result of a discard operation.
+ * The result of a revert operation.
  *
  * @internal
  */
-export enum DiscardResult {
-	/** The discard was successful. */
+export enum RevertibleResult {
+	/** The operation was successful. */
 	Success,
-	/** The discard failed. */
+	/** The operation failed. This occurs when attempting an operation on a disposed revertible */
 	Failure,
 }

packages/dds/tree/src/index.ts

@@ -84,8 +84,8 @@ export {
 	MapTree,
 	Revertible,
 	RevertibleKind,
-	RevertResult,
-	DiscardResult,
+	RevertibleStatus,
+	RevertibleResult,
 	forbiddenFieldKindIdentifier,
 	StoredSchemaCollection,
 	ErasedTreeNodeSchemaDataFormat,

packages/dds/tree/src/shared-tree-core/branch.ts

@@ -18,14 +18,13 @@ import {
 	makeAnonChange,
 	Revertible,
 	RevertibleKind,
-	RevertResult,
-	DiscardResult,
+	RevertibleResult,
+	RevertibleStatus,
 	BranchRebaseResult,
 	rebaseChangeOverChanges,
 	tagRollbackInverse,
 } from "../core/index.js";
 import { EventEmitter, ISubscribable } from "../events/index.js";
-import { fail } from "../util/index.js";
 import { TransactionStack } from "./transactionStack.js";
 
 /**
@@ -110,12 +109,16 @@ export interface SharedTreeBranchEvents<TEditor extends ChangeFamilyEditor, TCha
 	/**
 	 * Fired when a revertible change is made to this branch.
 	 */
-	revertible(type: Revertible): void;
+	newRevertible(revertible: Revertible): void;
 
 	/**
 	 * Fired when a revertible made on this branch is disposed.
+	 *
+	 * @param revertible - The revertible that was disposed.
+	 * This revertible was previously passed to the `newRevertible` event.
+	 * @param revision - The revision associated with the revertible that was disposed.
 	 */
-	revertibleDispose(revision: RevisionTag): void;
+	revertibleDisposed(revertible: Revertible, revision: RevisionTag): void;
 
 	/**
 	 * Fired when this branch forks
@@ -225,7 +228,7 @@ export class SharedTreeBranch<TEditor extends ChangeFamilyEditor, TChange> exten
 
 		// If this is not part of a transaction, emit a revertible event
 		if (!this.isTransacting()) {
-			this.emit("revertible", this.makeSharedTreeRevertible(newHead, revertibleKind));
+			this.emitNewRevertible(newHead, revertibleKind);
 		}
 
 		this.emit("afterChange", changeEvent);
@@ -302,7 +305,7 @@ export class SharedTreeBranch<TEditor extends ChangeFamilyEditor, TChange> exten
 
 		// If this transaction is not nested, emit a revertible event
 		if (!this.isTransacting()) {
-			this.emit("revertible", this.makeSharedTreeRevertible(newHead, RevertibleKind.Default));
+			this.emitNewRevertible(newHead, RevertibleKind.Default);
 		}
 
 		this.emit("afterChange", changeEvent);
@@ -389,59 +392,46 @@ export class SharedTreeBranch<TEditor extends ChangeFamilyEditor, TChange> exten
 		}
 	}
 
-	private makeSharedTreeRevertible(
-		commit: GraphCommit<TChange>,
-		kind: RevertibleKind,
-	): Revertible {
-		this._revertibleCommits.set(commit.revision, commit);
-		let discarded = false;
-		const revertible = {
+	private emitNewRevertible(commit: GraphCommit<TChange>, kind: RevertibleKind): void {
+		if (!this.hasListeners("newRevertible")) {
+			// No point generating revertibles if no one cares about them
+			return;
+		}
+		const revertible = new RevertibleRevision(
 			kind,
-			origin: {
-				// This is currently always the case, but we may want to support reverting remote ops
-				isLocal: true,
-			},
-			revert: () => {
-				if (discarded) {
-					fail("revertible has already been discarded");
-				}
-				const revertCommit = this.revert(commit.revision, kind);
-				if (revertCommit !== undefined) {
-					revertible.discard();
-					return RevertResult.Success;
-				}
-				return RevertResult.Failure;
-			},
-			discard: () => {
-				if (discarded) {
-					fail("revertible has already been discarded");
-				}
-				// TODO: delete the repair data from the forest
-				this._revertibleCommits.delete(commit.revision);
-				this.revertibles.delete(revertible);
-				discarded = true;
-				this.emit("revertibleDispose", commit.revision);
-				return DiscardResult.Success;
-			},
-		};
+			commit.revision,
+			this.revertRevertible.bind(this),
+			this.disposeRevertible.bind(this),
+		);
+		this._revertibleCommits.set(commit.revision, commit);
 		this.revertibles.add(revertible);
-		return revertible;
+		this.emit("newRevertible", revertible);
+		// Decrements the ref count for the revertible.
+		// This ensures that the revertible is disposed if no listener has retained it.
+		revertible.discard();
 	}
 
-	private revert(
-		revision: RevisionTag,
-		revertibleKind: RevertibleKind,
-	): [change: TChange, newCommit: GraphCommit<TChange>] | undefined {
+	private disposeRevertible(revertible: RevertibleRevision): void {
+		// TODO: delete the repair data from the forest
+		this._revertibleCommits.delete(revertible.revision);
+		this.revertibles.delete(revertible);
+		this.emit("revertibleDisposed", revertible, revertible.revision);
+	}
+
+	private revertRevertible(revertible: RevertibleRevision): void {
 		assert(!this.isTransacting(), 0x7cb /* Undo is not yet supported during transactions */);
 
-		const commit = this._revertibleCommits.get(revision);
+		const commit = this._revertibleCommits.get(revertible.revision);
 		assert(commit !== undefined, 0x7cc /* expected to find a revertible commit */);
 
-		let change = this.changeFamily.rebaser.invert(tagChange(commit.change, revision), false);
+		let change = this.changeFamily.rebaser.invert(
+			tagChange(commit.change, revertible.revision),
+			false,
+		);
 
 		const headCommit = this.getHead();
 		// Rebase the inverted change onto any commits that occurred after the undoable commits.
-		if (revision !== headCommit.revision) {
+		if (revertible.revision !== headCommit.revision) {
 			const pathAfterUndoable: GraphCommit<TChange>[] = [];
 			const ancestor = findCommonAncestor([commit], [headCommit, pathAfterUndoable]);
 			assert(
@@ -451,10 +441,10 @@ export class SharedTreeBranch<TEditor extends ChangeFamilyEditor, TChange> exten
 			change = rebaseChangeOverChanges(this.changeFamily.rebaser, change, pathAfterUndoable);
 		}
 
-		return this.applyChange(
+		this.applyChange(
 			change,
 			this.mintRevisionTag(),
-			revertibleKind === RevertibleKind.Default || revertibleKind === RevertibleKind.Redo
+			revertible.kind === RevertibleKind.Default || revertible.kind === RevertibleKind.Redo
 				? RevertibleKind.Undo
 				: RevertibleKind.Redo,
 		);
@@ -656,3 +646,63 @@ export function onForkTransitive<T extends ISubscribable<{ fork: (t: T) => void
 	);
 	return () => offs.forEach((off) => off());
 }
+
+class RevertibleRevision implements Revertible {
+	public readonly origin: Revertible["origin"];
+
+	private referenceCount = 1;
+
+	public constructor(
+		public readonly kind: RevertibleKind,
+		public readonly revision: RevisionTag,
+		private readonly onRevert: (revertible: RevertibleRevision) => void,
+		private readonly onDispose: (revertible: RevertibleRevision) => void,
+	) {
+		this.kind = kind;
+		this.revision = revision;
+		// This is currently always the case, but we may want to support reverting remote ops
+		this.origin = { isLocal: true };
+	}
+
+	public get status(): RevertibleStatus {
+		return this.referenceCount === 0 ? RevertibleStatus.Disposed : RevertibleStatus.Valid;
+	}
+
+	public revert(): RevertibleResult {
+		if (this.status === RevertibleStatus.Valid) {
+			this.onRevert(this);
+			this.dispose();
+			return RevertibleResult.Success;
+		}
+		return RevertibleResult.Failure;
+	}
+
+	public retain(): RevertibleResult {
+		if (this.status === RevertibleStatus.Valid) {
+			this.referenceCount += 1;
+			return RevertibleResult.Success;
+		}
+		return RevertibleResult.Failure;
+	}
+
+	public discard(): RevertibleResult {
+		if (this.status === RevertibleStatus.Valid) {
+			if (this.referenceCount === 1) {
+				this.dispose();
+			} else {
+				this.referenceCount -= 1;
+			}
+			return RevertibleResult.Success;
+		}
+		return RevertibleResult.Failure;
+	}
+
+	private dispose(): void {
+		assert(
+			this.status === RevertibleStatus.Valid,
+			"Cannot dispose already disposed revertible",
+		);
+		this.referenceCount = 0;
+		this.onDispose(this);
+	}
+}