Fix handling of deleted TreeNodes

packages/dds/tree/src/feature-libraries/flex-tree/context.ts

@@ -3,7 +3,7 @@
  * Licensed under the MIT License.
  */
 
-import { assert } from "@fluidframework/core-utils/internal";
+import { assert, debugAssert } from "@fluidframework/core-utils/internal";
 
 import {
 	type ForestEvents,
@@ -43,6 +43,11 @@ export interface FlexTreeContext {
 	 * If false, this context was created for use in a unhydrated tree, and the full document schema is unknown.
 	 */
 	isHydrated(): this is FlexTreeHydratedContext;
+
+	/**
+	 * If true, none of the nodes in this context can be used.
+	 */
+	isDisposed(): boolean;
 }
 
 /**
@@ -106,9 +111,14 @@ export class Context implements FlexTreeHydratedContext, IDisposable {
 	}
 
 	public isHydrated(): this is FlexTreeHydratedContext {
+		debugAssert(() => !this.disposed || "Disposed");
 		return true;
 	}
 
+	public isDisposed(): boolean {
+		return this.disposed;
+	}
+
 	public get schema(): TreeStoredSchema {
 		return this.checkout.storedSchema;
 	}

packages/dds/tree/src/feature-libraries/flex-tree/flexTreeTypes.ts

@@ -94,6 +94,14 @@ export enum TreeStatus {
 
 	/**
 	 * Is removed and cannot be added back to the original document tree.
+	 * @remarks
+	 * Nodes can enter this state for multiple reasons:
+	 * - The node was removed and nothing (ex: undo/redo history) kept it from being cleaned up.
+	 * - The {@link TreeView} was disposed or had a schema change which made the tree incompatible.
+	 * @privateRemarks
+	 * There was planned work (AB#17948) to make the first reason a node could become "Deleted" impossible,
+	 * at least as an opt in feature,
+	 * by lifetime extending all nodes which are still possible to reach automatically.
 	 */
 	Deleted = 2,
 

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

@@ -34,6 +34,8 @@ export class CheckoutFlexTreeView<out TCheckout extends ITreeCheckout = ITreeChe
 	 */
 	public readonly flexTree: FlexTreeField;
 
+	private disposed = false;
+
 	public constructor(
 		/**
 		 * Access non-view schema specific aspects of this branch.
@@ -52,6 +54,9 @@ export class CheckoutFlexTreeView<out TCheckout extends ITreeCheckout = ITreeChe
 	}
 
 	public [disposeSymbol](): void {
+		assert(!this.disposed, "Double disposed");
+		this.disposed = true;
+
 		for (const anchorNode of this.checkout.forest.anchors) {
 			tryDisposeTreeNode(anchorNode);
 		}
@@ -65,6 +70,7 @@ export class CheckoutFlexTreeView<out TCheckout extends ITreeCheckout = ITreeChe
 	 * Any mutations of the new view will not apply to this view until the new view is merged back into this view via `merge()`.
 	 */
 	public fork(): CheckoutFlexTreeView<ITreeCheckout & ITreeCheckoutFork> {
+		assert(!this.disposed, "disposed");
 		const branch = this.checkout.branch();
 		return new CheckoutFlexTreeView(branch, this.schema, this.nodeKeyManager);
 	}

packages/dds/tree/src/simple-tree/core/treeNodeKernel.ts

@@ -3,7 +3,8 @@
  * Licensed under the MIT License.
  */
 
-import { assert, Lazy, fail } from "@fluidframework/core-utils/internal";
+import { assert, Lazy, fail, debugAssert } from "@fluidframework/core-utils/internal";
+import { UsageError } from "@fluidframework/telemetry-utils/internal";
 import { createEmitter } from "@fluid-internal/client-utils";
 import type { Listenable, Off } from "@fluidframework/core-interfaces";
 import type { InternalTreeNode, TreeNode, Unhydrated } from "./types.js";
@@ -72,17 +73,17 @@ export function tryGetTreeNodeSchema(value: unknown): undefined | TreeNodeSchema
 /** The {@link HydrationState} of a {@link TreeNodeKernel} before the kernel is hydrated */
 interface UnhydratedState {
 	off: Off;
-	innerNode: UnhydratedFlexTreeNode;
+	readonly innerNode: UnhydratedFlexTreeNode;
 }
 
 /** The {@link HydrationState} of a {@link TreeNodeKernel} after the kernel is hydrated */
 interface HydratedState {
 	/** The flex node for this kernel (lazy - undefined if it has not yet been demanded) */
 	innerNode?: FlexTreeNode;
 	/** The {@link AnchorNode} that this node is associated with. */
-	anchorNode: AnchorNode;
+	readonly anchorNode: AnchorNode;
 	/** All {@link Off | event deregistration functions} that should be run when the kernel is disposed. */
-	offAnchorNode: Set<Off>;
+	readonly offAnchorNode: Set<Off>;
 }
 
 /** State within a {@link TreeNodeKernel} that is related to the hydration process */
@@ -287,35 +288,46 @@ export class TreeNodeKernel {
 	 *
 	 * For hydrated nodes it returns a FlexTreeNode backed by the forest.
 	 * Note that for "marinated" nodes, this FlexTreeNode exists and returns it: it does not return the MapTreeNode which is the current InnerNode.
+	 *
+	 * If `allowDeleted` is false, this will throw a UsageError if the node is deleted.
 	 */
-	public getOrCreateInnerNode(allowFreed = false): InnerNode {
+	public getOrCreateInnerNode(allowDeleted = false): InnerNode {
 		if (!isHydrated(this.#hydrationState)) {
+			debugAssert(
+				() =>
+					this.#hydrationState.innerNode?.context.isDisposed() === false ||
+					"Unhydrated node should never be disposed",
+			);
 			return this.#hydrationState.innerNode; // Unhydrated case
 		}
 
-		if (this.#hydrationState.innerNode !== undefined) {
-			return this.#hydrationState.innerNode; // Cooked case
+		if (this.#hydrationState.innerNode === undefined) {
+			// Marinated case -> cooked
+			const anchorNode = this.#hydrationState.anchorNode;
+			// The proxy is bound to an anchor node, but it may or may not have an actual flex node yet
+			const flexNode = anchorNode.slots.get(flexTreeSlot);
+			if (flexNode !== undefined) {
+				// If the flex node already exists, use it...
+				this.#hydrationState.innerNode = flexNode;
+			} else {
+				// ...otherwise, the flex node must be created
+				const context =
+					anchorNode.anchorSet.slots.get(ContextSlot) ?? fail(0xb41 /* missing context */);
+				const cursor = context.checkout.forest.allocateCursor("getFlexNode");
+				context.checkout.forest.moveCursorToPath(anchorNode, cursor);
+				this.#hydrationState.innerNode = makeTree(context, cursor);
+				cursor.free();
+				// Calling this is a performance improvement, however, do this only after demand to avoid momentarily having no anchors to anchorNode
+				anchorForgetters?.get(this.node)?.();
+				if (!allowDeleted) {
+					assertFlexTreeEntityNotFreed(this.#hydrationState.innerNode);
+				}
+			}
 		}
 
-		// Marinated case -> cooked
-		const anchorNode = this.#hydrationState.anchorNode;
-		// The proxy is bound to an anchor node, but it may or may not have an actual flex node yet
-		const flexNode = anchorNode.slots.get(flexTreeSlot);
-		if (flexNode !== undefined) {
-			// If the flex node already exists, use it...
-			this.#hydrationState.innerNode = flexNode;
-		} else {
-			// ...otherwise, the flex node must be created
-			const context =
-				anchorNode.anchorSet.slots.get(ContextSlot) ?? fail(0xb41 /* missing context */);
-			const cursor = context.checkout.forest.allocateCursor("getFlexNode");
-			context.checkout.forest.moveCursorToPath(anchorNode, cursor);
-			this.#hydrationState.innerNode = makeTree(context, cursor);
-			cursor.free();
-			// Calling this is a performance improvement, however, do this only after demand to avoid momentarily having no anchors to anchorNode
-			anchorForgetters?.get(this.node)?.();
-			if (!allowFreed) {
-				assertFlexTreeEntityNotFreed(this.#hydrationState.innerNode);
+		if (!allowDeleted) {
+			if (this.#hydrationState.innerNode.context.isDisposed()) {
+				throw new UsageError("Cannot access a Deleted node.");
 			}
 		}
 
@@ -417,11 +429,15 @@ export const unhydratedFlexTreeNodeToTreeNode =
  */
 export const proxySlot = anchorSlot<TreeNode>();
 
+/**
+ * Dispose a TreeNode (if any) for an existing anchor without disposing the anchor.
+ */
 export function tryDisposeTreeNode(anchorNode: AnchorNode): void {
 	const treeNode = anchorNode.slots.get(proxySlot);
 	if (treeNode !== undefined) {
 		const kernel = getKernel(treeNode);
 		kernel.dispose();
+		anchorNode.slots.delete(proxySlot);
 	}
 }
 
@@ -454,10 +470,12 @@ export function getSimpleContextFromInnerNode(innerNode: InnerNode): Context {
  *
  * For hydrated nodes it returns a FlexTreeNode backed by the forest.
  * Note that for "marinated" nodes, this FlexTreeNode exists and returns it: it does not return the MapTreeNode which is the current InnerNode.
+ *
+ * If `allowDeleted` is false, this will throw a UsageError if the node is deleted.
  */
-export function getOrCreateInnerNode(treeNode: TreeNode, allowFreed = false): InnerNode {
+export function getOrCreateInnerNode(treeNode: TreeNode, allowDeleted = false): InnerNode {
 	const kernel = getKernel(treeNode);
-	return kernel.getOrCreateInnerNode(allowFreed);
+	return kernel.getOrCreateInnerNode(allowDeleted);
 }
 
 /**

packages/dds/tree/src/simple-tree/core/unhydratedFlexTree.ts

@@ -73,7 +73,7 @@ interface LocationInField {
  *
  * Create a `UnhydratedFlexTreeNode` by calling {@link getOrCreate}.
  */
-export class UnhydratedFlexTreeNode implements UnhydratedFlexTreeNode {
+export class UnhydratedFlexTreeNode implements FlexTreeNode {
 	public get schema(): TreeNodeSchemaIdentifier {
 		return this.mapTree.type;
 	}
@@ -277,6 +277,10 @@ export class UnhydratedContext implements FlexTreeContext {
 		public readonly schema: TreeStoredSchema,
 	) {}
 
+	public isDisposed(): boolean {
+		return false;
+	}
+
 	public isHydrated(): this is FlexTreeHydratedContext {
 		return false;
 	}

packages/dds/tree/src/simple-tree/objectNode.ts

@@ -3,7 +3,7 @@
  * Licensed under the MIT License.
  */
 
-import { assert, Lazy, fail } from "@fluidframework/core-utils/internal";
+import { assert, Lazy, fail, debugAssert } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import type { FieldKey, SchemaPolicy } from "../core/index.js";
@@ -174,6 +174,8 @@ function createFlexKeyMapping(fields: Record<string, ImplicitFieldSchema>): Simp
 }
 
 /**
+ * Creates a proxy handler for the given schema.
+ *
  * @param allowAdditionalProperties - If true, setting of unexpected properties will be forwarded to the target object.
  * Otherwise setting of unexpected properties will error.
  * TODO: consider implementing this using `Object.preventExtension` instead.
@@ -196,9 +198,9 @@ function createProxyHandler(
 	const handler: ProxyHandler<TreeNode> = {
 		get(target, propertyKey, proxy): unknown {
 			const fieldInfo = schema.flexKeyMap.get(propertyKey);
-
 			if (fieldInfo !== undefined) {
 				const flexNode = getOrCreateInnerNode(proxy);
+				debugAssert(() => !flexNode.context.isDisposed() || "FlexTreeNode is disposed");
 				const field = flexNode.tryGetField(fieldInfo.storedKey);
 				if (field !== undefined) {
 					return getTreeNodeForField(field);

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

@@ -35,6 +35,7 @@ import {
 	// eslint-disable-next-line import/no-internal-modules
 } from "../../feature-libraries/chunked-forest/chunkedForest.js";
 import {
+	flexTreeSlot,
 	MockNodeIdentifierManager,
 	TreeCompressionStrategy,
 	TreeStatus,
@@ -111,6 +112,9 @@ import {
 	// eslint-disable-next-line import/no-internal-modules
 } from "../../simple-tree/api/index.js";
 import type { IChannel } from "@fluidframework/datastore-definitions/internal";
+import { configureDebugAsserts } from "@fluidframework/core-utils/internal";
+// eslint-disable-next-line import/no-internal-modules
+import { proxySlot } from "../../simple-tree/core/treeNodeKernel.js";
 
 const enableSchemaValidation = true;
 
@@ -129,6 +133,15 @@ class MockSharedTreeRuntime extends MockFluidDataStoreRuntime {
 }
 
 describe("SharedTree", () => {
+	let debugAssertsDefault: boolean;
+	beforeEach(() => {
+		debugAssertsDefault = configureDebugAsserts(true);
+	});
+
+	afterEach(() => {
+		configureDebugAsserts(debugAssertsDefault);
+	});
+
 	describe("viewWith", () => {
 		it("initialize tree", () => {
 			const tree = treeTestFactory();
@@ -160,9 +173,28 @@ describe("SharedTree", () => {
 			assert.deepEqual(view2.root, 10);
 		});
 
-		// TODO (AB#31456): Enable this test once the bug is fixed.
-		it.skip("initialize-dispose-view with object schema", () => {
+		it("re-view after view disposal with TreeNodes", () => {
 			const tree = treeTestFactory();
+
+			// Scan AnchorSet and check its slots for cached invalid data.
+			function checkAnchors(allowNodes: boolean) {
+				const anchors = tree.kernel.checkout.forest.anchors;
+				for (const anchor of anchors) {
+					const node = anchor.slots.get(flexTreeSlot);
+					if (node !== undefined) {
+						assert(node.context.isDisposed() === false);
+						assert(allowNodes);
+					}
+					const proxy = anchor.slots.get(proxySlot);
+					if (proxy !== undefined) {
+						assert.equal(Tree.status(proxy), TreeStatus.InDocument);
+						assert(allowNodes);
+					}
+				}
+			}
+
+			checkAnchors(false);
+
 			assert.deepEqual(tree.contentSnapshot().schema.rootFieldSchema, storedEmptyFieldSchema);
 
 			const factory = new SchemaFactory("my-factory");
@@ -182,10 +214,27 @@ describe("SharedTree", () => {
 			view1.initialize(new MySchema({ number: 10 }));
 			assert.deepEqual(view1.root, expectedContents);
 
+			const root1 = view1.root;
+
+			assert(Tree.status(root1) === TreeStatus.InDocument);
+
+			checkAnchors(true);
+
 			view1.dispose();
 
+			assert(Tree.status(root1) === TreeStatus.Deleted);
+			assert.throws(() => root1.number, validateUsageError(/Deleted/));
+
+			checkAnchors(false);
+
 			const view2 = tree.viewWith(config);
-			assert.deepEqual(view2.root, expectedContents); // <-- This throws with assert 0x778
+
+			checkAnchors(false);
+
+			const root2 = view2.root;
+			assert.notEqual(root1, root2);
+			assert.equal(view2.root.number, 10);
+			assert.deepEqual(view2.root, expectedContents);
 		});
 
 		it("concurrent initialize", () => {