Add re-entrancy errors for TreeCheckout (#23099)

## Description

This prevents the user from doing various actions to a SharedTree view
while in the middle of an ongoing tree edit. The current behavior for
these scenarios is undefined and currently tends to give unpredictable
internal error messages to users.

Author: noencke

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

@@ -383,7 +383,7 @@ export class SchematizingSimpleTreeView<
 
 	// #region Branching
 
-	public fork(): ReturnType<TreeBranch["fork"]> & TreeViewAlpha<TRootSchema> {
+	public fork(): ReturnType<TreeBranch["fork"]> & SchematizingSimpleTreeView<TRootSchema> {
 		return this.checkout.branch().viewWith(this.config);
 	}
 

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

@@ -73,6 +73,7 @@ import type {
 	UnsafeUnknownSchema,
 	ViewableTree,
 	TreeBranch,
+	TreeChangeEvents,
 } from "../simple-tree/index.js";
 import { getCheckout, SchematizingSimpleTreeView } from "./schematizingTreeView.js";
 
@@ -342,6 +343,8 @@ export interface RevertMetrics {
 export class TreeCheckout implements ITreeCheckoutFork {
 	public disposed = false;
 
+	private readonly editLock: EditLock;
+
 	private readonly views = new Set<TreeView<ImplicitFieldSchema>>();
 
 	/**
@@ -432,6 +435,8 @@ export class TreeCheckout implements ITreeCheckoutFork {
 			},
 		);
 
+		this.editLock = new EditLock(this.#transaction.activeBranchEditor);
+
 		branch.events.on("afterChange", (event) => {
 			// The following logic allows revertibles to be generated for the change.
 			// Currently only appends (including merges and transaction commits) are supported.
@@ -484,6 +489,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 	}
 
 	private readonly onAfterChange = (event: SharedTreeBranchChange<SharedTreeChange>): void => {
+		this.editLock.lock();
 		this.#events.emit("beforeBatch", event);
 		if (event.change !== undefined) {
 			const revision =
@@ -521,6 +527,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 			}
 		}
 		this.#events.emit("afterBatch");
+		this.editLock.unlock();
 		if (event.type === "append") {
 			event.newCommits.forEach((commit) => this.validateCommit(commit));
 		}
@@ -665,7 +672,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 
 	public get editor(): ISharedTreeEditor {
 		this.checkNotDisposed();
-		return this.#transaction.activeBranchEditor;
+		return this.editLock.editor;
 	}
 
 	public locate(anchor: Anchor): AnchorNode | undefined {
@@ -692,6 +699,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 		this.checkNotDisposed(
 			"The parent branch has already been disposed and can no longer create new branches.",
 		);
+		this.editLock.checkUnlocked("Branching");
 		const anchors = new AnchorSet();
 		const branch = this.#transaction.activeBranch.fork();
 		const storedSchema = this.storedSchema.clone();
@@ -720,6 +728,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 		checkout.checkNotDisposed(
 			"The source of the branch rebase has been disposed and cannot be rebased.",
 		);
+		this.editLock.checkUnlocked("Rebasing");
 		assert(
 			!checkout.transaction.isInProgress(),
 			0x9af /* A view cannot be rebased while it has a pending transaction */,
@@ -748,6 +757,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 		checkout.checkNotDisposed(
 			"The source of the branch merge has been disposed and cannot be merged.",
 		);
+		this.editLock.checkUnlocked("Merging");
 		assert(
 			!this.transaction.isInProgress(),
 			0x9b0 /* Views cannot be merged into a view while it has a pending transaction */,
@@ -768,6 +778,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 	}
 
 	public dispose(): void {
+		this.editLock.checkUnlocked("Disposing a view");
 		this[disposeSymbol]();
 	}
 
@@ -955,3 +966,103 @@ export class TreeCheckout implements ITreeCheckoutFork {
 
 	// #endregion Commit Validation
 }
+
+/**
+ * A helper class that assists {@link TreeCheckout} in preventing functionality from being used while the tree is in the middle of being edited.
+ */
+class EditLock {
+	/**
+	 * Edits the tree by calling the methods of the editor passed into the {@link EditLock} constructor.
+	 * @remarks Edits will throw an error if the lock is currently locked.
+	 */
+	public editor: ISharedTreeEditor;
+	private locked = false;
+
+	/**
+	 * @param editor - an editor which will be used to create a new editor that is monitored to determine if any changes are happening to the tree.
+	 * Use {@link EditLock.editor} in place of the original editor to ensure that changes are monitored.
+	 */
+	public constructor(editor: ISharedTreeEditor) {
+		const checkLock = (): void => this.checkUnlocked("Editing the tree");
+		this.editor = {
+			get schema() {
+				return editor.schema;
+			},
+			valueField(...fieldArgs) {
+				const valueField = editor.valueField(...fieldArgs);
+				return {
+					set(...editArgs) {
+						checkLock();
+						valueField.set(...editArgs);
+					},
+				};
+			},
+			optionalField(...fieldArgs) {
+				const optionalField = editor.optionalField(...fieldArgs);
+				return {
+					set(...editArgs) {
+						checkLock();
+						optionalField.set(...editArgs);
+					},
+				};
+			},
+			sequenceField(...fieldArgs) {
+				const sequenceField = editor.sequenceField(...fieldArgs);
+				return {
+					insert(...editArgs) {
+						checkLock();
+						sequenceField.insert(...editArgs);
+					},
+					remove(...editArgs) {
+						checkLock();
+						sequenceField.remove(...editArgs);
+					},
+				};
+			},
+			move(...moveArgs) {
+				checkLock();
+				editor.move(...moveArgs);
+			},
+			addNodeExistsConstraint(path) {
+				editor.addNodeExistsConstraint(path);
+			},
+		};
+	}
+
+	/**
+	 * Prevent further changes from being made to {@link EditLock.editor} until {@link EditLock.unlock} is called.
+	 * @remarks May only be called when the lock is not already locked.
+	 */
+	public lock(): void {
+		if (this.locked) {
+			debugger;
+		}
+		assert(!this.locked, "Checkout has already been locked");
+		this.locked = true;
+	}
+
+	/**
+	 * Throws an error if the lock is currently locked.
+	 * @param action - The current action being performed by the user.
+	 * This must start with a capital letter, as it shows up as the first part of the error message and we want it to look nice.
+	 */
+	public checkUnlocked<T extends string>(action: T extends Capitalize<T> ? T : never): void {
+		if (this.locked) {
+			// These type assertions ensure that the event name strings used here match the actual event names
+			const nodeChanged: keyof TreeChangeEvents = "nodeChanged";
+			const treeChanged: keyof TreeChangeEvents = "treeChanged";
+			throw new UsageError(
+				`${action} is forbidden during a ${nodeChanged} or ${treeChanged} event`,
+			);
+		}
+	}
+
+	/**
+	 * Allow changes to be made to {@link EditLock.editor} again.
+	 * @remarks May only be called when the lock is currently locked.
+	 */
+	public unlock(): void {
+		assert(this.locked, "Checkout has not been locked");
+		this.locked = false;
+	}
+}

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

@@ -51,6 +51,7 @@ import {
 	getOrCreateInnerNode,
 	toStoredSchema,
 	type InsertableField,
+	type TreeBranch,
 } from "../../simple-tree/index.js";
 // eslint-disable-next-line import/no-internal-modules
 import { stringSchema } from "../../simple-tree/leafNodeSchema.js";
@@ -1268,6 +1269,78 @@ describe("sharedTreeView", () => {
 			});
 		}
 	});
+
+	describe("throws an error if it is in the middle of an edit when a user attempts to", () => {
+		const sf = new SchemaFactory("Checkout and view test schema");
+		class NumberNode extends sf.object("Number", { number: sf.number }) {}
+
+		/** Tests that an error is thrown when a given action is taken during the execution of a nodeChanged/treeChanged listener */
+		function expectErrorDuringEdit(args: {
+			/**
+			 * Runs after the main view has been created but before the edit occurs
+			 * @returns (optionally) a view (e.g. a fork of the main view) that will be passed to `duringEdit`
+			 */
+			setup?: (
+				view: SchematizingSimpleTreeView<typeof NumberNode>,
+			) => void | SchematizingSimpleTreeView<typeof NumberNode>;
+			/** The code to run during the edit that should throw an error */
+			duringEdit: (view: SchematizingSimpleTreeView<typeof NumberNode>) => void;
+			/** The expected error message */
+			error: string;
+		}): void {
+			let view = getView(
+				new TreeViewConfiguration({ enableSchemaValidation, schema: NumberNode }),
+			);
+
+			view.initialize({ number: 3 });
+			view = args.setup?.(view) ?? view;
+
+			Tree.on(view.root, "nodeChanged", () => {
+				args.duringEdit(view);
+			});
+
+			assert.throws(() => (view.root.number = 0), new RegExp(args.error));
+		}
+
+		it("edit the tree", () => {
+			expectErrorDuringEdit({
+				duringEdit: (view) => {
+					view.root.number = 4;
+				},
+				error: "Editing the tree is forbidden during a nodeChanged or treeChanged event",
+			});
+		});
+
+		it("create a branch", () => {
+			expectErrorDuringEdit({
+				duringEdit: (view) => view.fork(),
+				error: ".*Branching is forbidden during a nodeChanged or treeChanged event.*",
+			});
+		});
+
+		it("rebase a branch", () => {
+			expectErrorDuringEdit({
+				duringEdit: (view) => view.rebaseOnto(view),
+				error: "Rebasing is forbidden during a nodeChanged or treeChanged event",
+			});
+		});
+
+		it("merge a branch", () => {
+			expectErrorDuringEdit({
+				duringEdit: (view) => view.merge(view),
+				error: "Merging is forbidden during a nodeChanged or treeChanged event",
+			});
+		});
+
+		it("dispose", () => {
+			let branch: TreeBranch | undefined;
+			expectErrorDuringEdit({
+				setup: (view) => (branch = view.fork()), // Create a fork of the view because the main view can't be disposed
+				duringEdit: (view) => view.dispose(),
+				error: "Disposing a view is forbidden during a nodeChanged or treeChanged event",
+			});
+		});
+	});
 });
 
 const defaultSf = new SchemaFactory("Checkout and view test schema");