Fix unhydrated node parent tracking (#24708)

## Description

Fix unhydrated node parent tracking by centralizing the logic to update
parents, and applying it consistently after every edit.

This does make small edits in large sequences do O(N) parent updates,
but they were already O(N) from modifying the large array, and need to
do O(N) parent updates in the insertion is near the front anyway.

Author: CraigMacomber

.changeset/eight-ears-feel.md

@@ -0,0 +1,9 @@
+---
+"@fluidframework/tree": minor
+"fluid-framework": minor
+"__section": tree
+---
+Fix Tree.key and Tree.parent for Unhydrated nodes after edits
+
+In some cases, editing [Unhydrated](https://fluidframework.com/docs/api/fluid-framework/unhydrated-typealias) nodes could result in incorrect results being returned from [Tree.key](https://fluidframework.com/docs/data-structures/tree/nodes#treekey) and [Tree.parent](https://fluidframework.com/docs/data-structures/tree/nodes#treeparent).
+This has been fixed.

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

@@ -402,13 +402,24 @@ class UnhydratedFlexTreeField implements FlexTreeField {
 	 */
 	protected edit(edit: (mapTrees: ExclusiveMapTree[]) => void | ExclusiveMapTree[]): void {
 		const oldMapTrees = this.parent.mapTree.fields.get(this.key) ?? [];
+
+		// Clear parents for all old map trees.
+		for (const tree of oldMapTrees) {
+			tryUnhydratedFlexTreeNode(tree)?.adoptBy(undefined);
+		}
+
 		const newMapTrees = edit(oldMapTrees) ?? oldMapTrees;
 		if (newMapTrees.length > 0) {
 			this.parent.mapTree.fields.set(this.key, newMapTrees);
 		} else {
 			this.parent.mapTree.fields.delete(this.key);
 		}
 
+		// Set parents for all new map trees.
+		for (const [index, tree] of newMapTrees.entries()) {
+			tryUnhydratedFlexTreeNode(tree)?.adoptBy(this, index);
+		}
+
 		this.onEdit?.();
 	}
 
@@ -434,16 +445,6 @@ class EagerMapTreeOptionalField
 {
 	public readonly editor = {
 		set: (newContent: ExclusiveMapTree | undefined): void => {
-			// If the new content is a UnhydratedFlexTreeNode, it needs to have its parent pointer updated
-			if (newContent !== undefined) {
-				nodeCache.get(newContent)?.adoptBy(this, 0);
-			}
-			// If the old content is a UnhydratedFlexTreeNode, it needs to have its parent pointer unset
-			const oldContent = this.mapTrees[0];
-			if (oldContent !== undefined) {
-				nodeCache.get(oldContent)?.adoptBy(undefined);
-			}
-
 			this.edit((mapTrees) => {
 				if (newContent !== undefined) {
 					mapTrees[0] = newContent;
@@ -484,10 +485,8 @@ export class UnhydratedTreeSequenceField
 {
 	public readonly editor: UnhydratedTreeSequenceFieldEditBuilder = {
 		insert: (index, newContent): void => {
-			for (let i = 0; i < newContent.length; i++) {
-				const c = newContent[i];
+			for (const c of newContent) {
 				assert(c !== undefined, 0xa0a /* Unexpected sparse array content */);
-				nodeCache.get(c)?.adoptBy(this, index + i);
 			}
 			this.edit((mapTrees) => {
 				if (newContent.length < 1000) {
@@ -503,7 +502,6 @@ export class UnhydratedTreeSequenceField
 			for (let i = index; i < index + count; i++) {
 				const c = this.mapTrees[i];
 				assert(c !== undefined, 0xa0b /* Unexpected sparse array */);
-				nodeCache.get(c)?.adoptBy(undefined);
 			}
 			let removed: ExclusiveMapTree[] | undefined;
 			this.edit((mapTrees) => {

packages/dds/tree/src/test/simple-tree/api/treeNodeApi.spec.ts

@@ -35,7 +35,7 @@ import {
 	TestTreeProviderLite,
 	validateUsageError,
 } from "../../utils.js";
-import { getViewForForkedBranch, hydrate } from "../utils.js";
+import { describeHydration, getViewForForkedBranch, hydrate } from "../utils.js";
 import { brand, type areSafelyAssignable, type requireTrue } from "../../../util/index.js";
 
 import {
@@ -165,39 +165,70 @@ describe("treeNodeApi", () => {
 		});
 	});
 
-	it("key", () => {
-		class Child extends schema.object("Child", {
-			x: Point,
-			y: schema.optional(Point, { key: "stable-y" }),
-		}) {}
-		const Root = schema.array(Child);
-		const config = new TreeViewConfiguration({ schema: Root });
-		const view = getView(config);
-		view.initialize([
-			{ x: {}, y: undefined },
-			{ x: {}, y: {} },
-		]);
-		const { root } = view;
-		assert.equal(Tree.key(root), rootFieldKey);
-		assert.equal(Tree.key(root[0]), 0);
-		assert.equal(Tree.key(root[0].x), "x");
-		assert.equal(Tree.key(root[1]), 1);
-		assert.equal(Tree.key(root[1].x), "x");
-		assert(root[1].y !== undefined);
-		assert.equal(Tree.key(root[1].y), "y");
-	});
+	describeHydration("upward path", (init) => {
+		for (const [name, keyApi] of [
+			["key", (n: TreeNode): string | undefined | number => Tree.key(n)],
+			["key2", (n: TreeNode): string | undefined | number => TreeAlpha.key2(n)],
+		] as const) {
+			it(name, () => {
+				class Child extends schema.object("Child", {
+					x: Point,
+					y: schema.optional(Point, { key: "stable-y" }),
+				}) {}
+				class Root extends schema.array("Root", Child) {}
+				const root = init(Root, [
+					{ x: {}, y: undefined },
+					{ x: {}, y: {} },
+				]);
 
-	it("parent", () => {
-		class Child extends schema.object("Child", { x: Point }) {}
-		const Root = schema.array(Child);
-		const config = new TreeViewConfiguration({ schema: Root });
-		const view = getView(config);
-		view.initialize([{ x: {} }, { x: {} }]);
-		const { root } = view;
-		assert.equal(Tree.parent(root), undefined);
-		assert.equal(Tree.parent(root[0]), root);
-		assert.equal(Tree.parent(root[1]), root);
-		assert.equal(Tree.parent(root[1].x), root[1]);
+				// This is this how we handle root keys.
+				// Seems odd for detached fields other than root to have `rootFieldKey` key though.
+				// Exactly which key is given in this case is undocumented, it could change in the future.
+				// TreeAlpha.key2 just gives undefined, which is documented.
+				const rootKey = name === "key" ? rootFieldKey : undefined;
+
+				assert.equal(keyApi(root), rootKey);
+				assert.equal(keyApi(root[0]), 0);
+				assert.equal(keyApi(root[0].x), "x");
+				assert.equal(keyApi(root[1]), 1);
+				assert.equal(keyApi(root[1].x), "x");
+				assert(root[1].y !== undefined);
+				assert.equal(keyApi(root[1].y), "y");
+
+				const added = new Child({ x: {}, y: {} });
+
+				assert.equal(keyApi(added), rootKey);
+
+				// Check index is updated after insert.
+				root.insertAtStart(added);
+				assert.equal(keyApi(root[2]), 2);
+				assert.equal(keyApi(added), 0);
+
+				// Check index is updated after removal.
+				root.removeRange(0, 1);
+				assert.equal(keyApi(root[1]), 1);
+				assert.equal(keyApi(added), rootKey);
+			});
+		}
+
+		it("parent", () => {
+			class Child extends schema.object("Child", { x: Point }) {}
+			class Root extends schema.array("Root", Child) {}
+			const root = init(Root, [{ x: {} }, { x: {} }]);
+
+			assert.equal(Tree.parent(root), undefined);
+			assert.equal(Tree.parent(root[0]), root);
+			assert.equal(Tree.parent(root[1]), root);
+			assert.equal(Tree.parent(root[1].x), root[1]);
+
+			const added = new Child({ x: {} });
+
+			assert.equal(Tree.parent(added), undefined);
+			root.insertAtStart(added);
+			assert.equal(Tree.parent(added), root);
+			root.removeRange(0, 1);
+			assert.equal(Tree.parent(added), undefined);
+		});
 	});
 
 	it("treeStatus", () => {
@@ -215,29 +246,13 @@ describe("treeNodeApi", () => {
 		assert.equal(Tree.status(root), TreeStatus.InDocument);
 		assert.equal(Tree.status(child), TreeStatus.Removed);
 		assert.equal(Tree.status(newChild), TreeStatus.InDocument);
-		// TODO: test Deleted status.
-	});
 
-	it("key2", () => {
-		class Child extends schema.object("Child", {
-			x: Point,
-			y: schema.optional(Point, { key: "stable-y" }),
-		}) {}
-		const Root = schema.array(Child);
-		const config = new TreeViewConfiguration({ schema: Root });
-		const view = getView(config);
-		view.initialize([
-			{ x: {}, y: undefined },
-			{ x: {}, y: {} },
-		]);
-		const { root } = view;
-		assert.equal(TreeAlpha.key2(root), undefined);
-		assert.equal(TreeAlpha.key2(root[0]), 0);
-		assert.equal(TreeAlpha.key2(root[0].x), "x");
-		assert.equal(TreeAlpha.key2(root[1]), 1);
-		assert.equal(TreeAlpha.key2(root[1].x), "x");
-		assert(root[1].y !== undefined);
-		assert.equal(TreeAlpha.key2(root[1].y), "y");
+		view.dispose();
+		assert.equal(Tree.status(root), TreeStatus.Deleted);
+		assert.equal(Tree.status(child), TreeStatus.Deleted);
+		assert.equal(Tree.status(newChild), TreeStatus.Deleted);
+
+		// TODO: test Deleted status when caused by removal from the tree + expiring from removed status.
 	});
 
 	describe("shortID", () => {