tree: Fix bug with updating ancestry table when composing moved nodes (…

…microsoft#22463)

## Description

This fixes a bug where ModularChangeset's `nodeToParent` table was not
updated when a node was moved to a different field during `compose`.

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

@@ -158,6 +158,7 @@ export {
 
 export {
 	areEqualChangeAtomIds,
+	areEqualChangeAtomIdOpts,
 	makeChangeAtomId,
 	asChangeAtomId,
 	type ChangeRebaser,

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

@@ -5,6 +5,7 @@
 
 export {
 	areEqualChangeAtomIds,
+	areEqualChangeAtomIdOpts,
 	makeChangeAtomId,
 	asChangeAtomId,
 	mintCommit,

packages/dds/tree/src/core/rebase/types.ts

@@ -81,6 +81,17 @@ export function areEqualChangeAtomIds(a: ChangeAtomId, b: ChangeAtomId): boolean
 	return a.localId === b.localId && a.revision === b.revision;
 }
 
+export function areEqualChangeAtomIdOpts(
+	a: ChangeAtomId | undefined,
+	b: ChangeAtomId | undefined,
+): boolean {
+	if (a === undefined || b === undefined) {
+		return a === b;
+	}
+
+	return areEqualChangeAtomIds(a, b);
+}
+
 /**
  * @returns a ChangeAtomId with the given revision and local ID.
  */

packages/dds/tree/src/feature-libraries/modular-schema/modularChangeFamily.ts

@@ -38,6 +38,7 @@ import {
 	revisionMetadataSourceFromInfo,
 	areEqualChangeAtomIds,
 	type ChangeAtomId,
+	areEqualChangeAtomIdOpts,
 } from "../../core/index.js";
 import {
 	type IdAllocationState,
@@ -263,8 +264,6 @@ export class ModularChangeFamily
 		const genId: IdAllocator = idAllocatorFromState(idState);
 		const revisionMetadata: RevisionMetadataSource = revisionMetadataSourceFromInfo(revInfos);
 
-		const crossFieldTable = newComposeTable(change1, change2);
-
 		// We merge nodeChanges, nodeToParent, and nodeAliases from the two changesets.
 		// The merged tables will have correct entries for all nodes which are only referenced in one of the input changesets.
 		// During composeFieldMaps and composeInvalidatedElements we will find all nodes referenced in both input changesets
@@ -284,9 +283,12 @@ export class ModularChangeFamily
 			mergeBTrees(change1.nodeAliases, change2.nodeAliases),
 		);
 
+		const crossFieldTable = newComposeTable(change1, change2, composedNodeToParent);
+
 		const composedFields = this.composeFieldMaps(
 			change1.fieldChanges,
 			change2.fieldChanges,
+			undefined,
 			genId,
 			crossFieldTable,
 			revisionMetadata,
@@ -321,7 +323,7 @@ export class ModularChangeFamily
 	): void {
 		const context = crossFieldTable.fieldToContext.get(fieldChange);
 		assert(context !== undefined, 0x8cc /* Should have context for every invalidated field */);
-		const { change1: fieldChange1, change2: fieldChange2, composedChange } = context;
+		const { fieldId, change1: fieldChange1, change2: fieldChange2, composedChange } = context;
 
 		const rebaser = getChangeHandler(this.fieldKinds, composedChange.fieldKind).rebaser;
 		const composeNodes = (child1: NodeId | undefined, child2: NodeId | undefined): NodeId => {
@@ -342,7 +344,7 @@ export class ModularChangeFamily
 			fieldChange2,
 			composeNodes,
 			genId,
-			new ComposeManager(crossFieldTable, fieldChange, false),
+			new ComposeManager(crossFieldTable, fieldChange, fieldId, false),
 			revisionMetadata,
 		);
 		composedChange.change = brand(amendedChange);
@@ -470,7 +472,14 @@ export class ModularChangeFamily
 				? [fieldChange, emptyChange]
 				: [emptyChange, fieldChange];
 
-			const composedField = this.composeFieldChanges(change1, change2, genId, table, metadata);
+			const composedField = this.composeFieldChanges(
+				fieldId,
+				change1,
+				change2,
+				genId,
+				table,
+				metadata,
+			);
 
 			if (fieldId.nodeId === undefined) {
 				composedFields.set(fieldId.field, composedField);
@@ -499,6 +508,7 @@ export class ModularChangeFamily
 	private composeFieldMaps(
 		change1: FieldChangeMap | undefined,
 		change2: FieldChangeMap | undefined,
+		parentId: NodeId | undefined,
 		genId: IdAllocator,
 		crossFieldTable: ComposeTable,
 		revisionMetadata: RevisionMetadataSource,
@@ -509,10 +519,12 @@ export class ModularChangeFamily
 		}
 
 		for (const [field, fieldChange1] of change1) {
+			const fieldId: FieldId = { nodeId: parentId, field };
 			const fieldChange2 = change2.get(field);
 			const composedField =
 				fieldChange2 !== undefined
 					? this.composeFieldChanges(
+							fieldId,
 							fieldChange1,
 							fieldChange2,
 							genId,
@@ -545,6 +557,7 @@ export class ModularChangeFamily
 	 * Any composed `FieldChange` which is invalidated by new cross-field information will be added to `crossFieldTable.invalidatedFields`.
 	 */
 	private composeFieldChanges(
+		fieldId: FieldId,
 		change1: FieldChange,
 		change2: FieldChange,
 		idAllocator: IdAllocator,
@@ -558,7 +571,7 @@ export class ModularChangeFamily
 			change2: change2Normalized,
 		} = this.normalizeFieldChanges(change1, change2, idAllocator, revisionMetadata);
 
-		const manager = new ComposeManager(crossFieldTable, change1);
+		const manager = new ComposeManager(crossFieldTable, change1, fieldId);
 
 		const composedChange = changeHandler.rebaser.compose(
 			change1Normalized,
@@ -581,6 +594,7 @@ export class ModularChangeFamily
 		};
 
 		crossFieldTable.fieldToContext.set(change1, {
+			fieldId,
 			change1: change1Normalized,
 			change2: change2Normalized,
 			composedChange: composedField,
@@ -605,6 +619,7 @@ export class ModularChangeFamily
 		const nodeChangeset1 = nodeChangeFromId(nodeChanges1, id1);
 		const nodeChangeset2 = nodeChangeFromId(nodeChanges2, id2);
 		const composedNodeChangeset = this.composeNodeChanges(
+			id1,
 			nodeChangeset1,
 			nodeChangeset2,
 			idAllocator,
@@ -627,6 +642,7 @@ export class ModularChangeFamily
 	}
 
 	private composeNodeChanges(
+		nodeId: NodeId,
 		change1: NodeChangeset,
 		change2: NodeChangeset,
 		genId: IdAllocator,
@@ -638,6 +654,7 @@ export class ModularChangeFamily
 		const composedFieldChanges = this.composeFieldMaps(
 			change1.fieldChanges,
 			change2.fieldChanges,
+			nodeId,
 			genId,
 			crossFieldTable,
 			revisionMetadata,
@@ -887,7 +904,7 @@ export class ModularChangeFamily
 			rebasedNodes,
 		);
 
-		return makeModularChangeset(
+		const rebased = makeModularChangeset(
 			this.pruneFieldMap(rebasedFields, rebasedNodes),
 			rebasedNodes,
 			crossFieldTable.rebasedNodeToParent,
@@ -900,6 +917,8 @@ export class ModularChangeFamily
 			change.destroys,
 			change.refreshers,
 		);
+
+		return rebased;
 	}
 
 	// This performs a first pass on all fields which have both new and base changes.
@@ -1584,6 +1603,57 @@ export class ModularChangeFamily
 		const emptyChange = getChangeHandler(this.fieldKinds, fieldKind).createEmpty();
 		return { fieldKind, change: brand(emptyChange) };
 	}
+
+	public validateChangeset(change: ModularChangeset): void {
+		let numNodes = this.validateFieldChanges(change, change.fieldChanges, undefined);
+
+		for (const [[revision, localId], node] of change.nodeChanges.entries()) {
+			if (node.fieldChanges === undefined) {
+				continue;
+			}
+
+			const nodeId: NodeId = { revision, localId };
+			const numChildren = this.validateFieldChanges(change, node.fieldChanges, nodeId);
+
+			numNodes += numChildren;
+		}
+
+		assert(numNodes === change.nodeChanges.size, "Node table contains unparented nodes");
+	}
+
+	/**
+	 * Asserts that each child and cross field key in each field has a correct entry in
+	 * `nodeToParent` or `crossFieldKeyTable`.
+	 * @returns the number of children found.
+	 */
+	private validateFieldChanges(
+		change: ModularChangeset,
+		fieldChanges: FieldChangeMap,
+		nodeParent: NodeId | undefined,
+	): number {
+		let numChildren = 0;
+		for (const [field, fieldChange] of fieldChanges.entries()) {
+			const fieldId = { nodeId: nodeParent, field };
+			const handler = getChangeHandler(this.fieldKinds, fieldChange.fieldKind);
+			for (const [child, _index] of handler.getNestedChanges(fieldChange.change)) {
+				const parentFieldId = getParentFieldId(change, child);
+				assert(areEqualFieldIds(parentFieldId, fieldId), "Inconsistent node parentage");
+				numChildren += 1;
+			}
+
+			for (const keyRange of handler.getCrossFieldKeys(fieldChange.change)) {
+				const fields = getFieldsForCrossFieldKey(change, keyRange);
+				assert(
+					fields.length === 1 &&
+						fields[0] !== undefined &&
+						areEqualFieldIds(fields[0], fieldId),
+					"Inconsistent cross field keys",
+				);
+			}
+		}
+
+		return numChildren;
+	}
 }
 
 function replaceCrossFieldKeyTableRevisions(
@@ -2059,6 +2129,7 @@ interface RebaseFieldContext {
 function newComposeTable(
 	baseChange: ModularChangeset,
 	newChange: ModularChangeset,
+	composedNodeToParent: ChangeAtomIdBTree<FieldId>,
 ): ComposeTable {
 	return {
 		...newCrossFieldTable<FieldChange>(),
@@ -2068,6 +2139,7 @@ function newComposeTable(
 		newFieldToBaseField: new Map(),
 		newToBaseNodeId: newTupleBTree(),
 		composedNodes: new Set(),
+		composedNodeToParent,
 		pendingCompositions: {
 			nodeIdsToCompose: [],
 			affectedBaseFields: newTupleBTree(),
@@ -2087,6 +2159,7 @@ interface ComposeTable extends CrossFieldTable<FieldChange> {
 	readonly newFieldToBaseField: Map<FieldChange, FieldChange>;
 	readonly newToBaseNodeId: ChangeAtomIdBTree<NodeId>;
 	readonly composedNodes: Set<NodeChangeset>;
+	readonly composedNodeToParent: ChangeAtomIdBTree<FieldId>;
 	readonly pendingCompositions: PendingCompositions;
 }
 
@@ -2109,6 +2182,10 @@ interface PendingCompositions {
 }
 
 interface ComposeFieldContext {
+	/**
+	 * The field ID for this field in the composed changeset.
+	 */
+	fieldId: FieldId;
 	change1: FieldChangeset;
 	change2: FieldChangeset;
 	composedChange: FieldChange;
@@ -2324,7 +2401,12 @@ class RebaseManager extends CrossFieldManagerI<FieldChange> {
 
 // TODO: Deduplicate this with RebaseTable
 class ComposeManager extends CrossFieldManagerI<FieldChange> {
-	public constructor(table: ComposeTable, currentField: FieldChange, allowInval = true) {
+	public constructor(
+		table: ComposeTable,
+		currentField: FieldChange,
+		private readonly fieldId: FieldId,
+		allowInval = true,
+	) {
 		super(table, currentField, allowInval);
 	}
 
@@ -2377,15 +2459,16 @@ class ComposeManager extends CrossFieldManagerI<FieldChange> {
 	}
 
 	public override onMoveIn(id: ChangeAtomId): void {
-		throw new Error("Method not implemented.");
+		setInChangeAtomIdMap(this.table.composedNodeToParent, id, this.fieldId);
 	}
+
 	public override moveKey(
 		target: CrossFieldTarget,
 		revision: RevisionTag | undefined,
 		id: ChangesetLocalId,
 		count: number,
 	): void {
-		throw new Error("Method not implemented.");
+		throw new Error("Moving cross-field keys during compose is currently unsupported");
 	}
 
 	private get table(): ComposeTable {
@@ -3022,3 +3105,7 @@ function getFromChangeAtomIdMap<T>(
 function setInChangeAtomIdMap<T>(map: ChangeAtomIdBTree<T>, id: ChangeAtomId, value: T): void {
 	map.set([id.revision, id.localId], value);
 }
+
+function areEqualFieldIds(a: FieldId, b: FieldId): boolean {
+	return areEqualChangeAtomIdOpts(a.nodeId, b.nodeId) && a.field === b.field;
+}

packages/dds/tree/src/feature-libraries/sequence-field/compose.ts

@@ -576,6 +576,9 @@ export class ComposeQueue {
 	private dequeueBase(length: number = Infinity): ComposeMarks {
 		const baseMark = this.baseMarks.dequeueUpTo(length);
 		const movedChanges = getMovedChangesFromMark(this.moveEffects, baseMark);
+		if (movedChanges !== undefined) {
+			this.moveEffects.onMoveIn(movedChanges);
+		}
 
 		const newMark = createNoopMark(baseMark.count, movedChanges, getOutputCellId(baseMark));
 		return { baseMark, newMark };

packages/dds/tree/src/feature-libraries/sequence-field/utils.ts

@@ -10,6 +10,7 @@ import {
 	type ChangesetLocalId,
 	type RevisionMetadataSource,
 	type RevisionTag,
+	areEqualChangeAtomIdOpts,
 	areEqualChangeAtomIds,
 	makeChangeAtomId,
 } from "../../core/index.js";
@@ -118,10 +119,7 @@ export function isActiveReattach(
 }
 
 export function areEqualCellIds(a: CellId | undefined, b: CellId | undefined): boolean {
-	if (a === undefined || b === undefined) {
-		return a === b;
-	}
-	return areEqualChangeAtomIds(a, b);
+	return areEqualChangeAtomIdOpts(a, b);
 }
 
 export function getInputCellId(mark: Mark): CellId | undefined {

packages/dds/tree/src/test/feature-libraries/modularChangeFamilyIntegration.spec.ts

@@ -546,6 +546,7 @@ describe("ModularChangeFamily integration", () => {
 			const remove = makeAnonChange(removeD);
 
 			const composed = family.compose([moves, remove]);
+			family.validateChangeset(composed);
 			const composedDelta = normalizeDelta(intoDelta(makeAnonChange(composed), fieldKinds));
 
 			const nodeAChanges: DeltaFieldMap = new Map([
@@ -637,6 +638,7 @@ describe("ModularChangeFamily integration", () => {
 				]),
 			};
 
+			family.validateChangeset(composed);
 			const delta = intoDelta(makeAnonChange(composed), family.fieldKinds);
 			assertDeltaEqual(delta, expected);
 		});
@@ -674,6 +676,8 @@ describe("ModularChangeFamily integration", () => {
 
 			const moveAndInsert = family.compose([tagChangeInline(insert, tag2), moveTagged]);
 			const composed = family.compose([returnTagged, makeAnonChange(moveAndInsert)]);
+			family.validateChangeset(composed);
+
 			const actual = intoDelta(makeAnonChange(composed), family.fieldKinds);
 			const expected: DeltaRoot = {
 				build: [
@@ -735,6 +739,7 @@ describe("ModularChangeFamily integration", () => {
 
 			const [move1, move2, expected] = getChanges();
 			const composed = family.compose([makeAnonChange(move1), makeAnonChange(move2)]);
+			family.validateChangeset(composed);
 			const actualDelta = normalizeDelta(
 				intoDelta(makeAnonChange(composed), family.fieldKinds),
 			);