tree: Refactor tree node kernel (#22369)

## Description

Refactor TreeNodeKernel.

Makes its events in terms of the internal anchor node events instead of
the public API surface, allowing for the public API surface to be
tweaked in node kind specific ways.

Track the disposed state more clearly.

Move generation number logic from array nodes to node kernel where other
kinds could use it (ex: to error when iterating map keys if the map was
modified) and implementation is simpler.

Changes split out from
#22229 for separate
review.

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

@@ -167,7 +167,19 @@ export const treeNodeApi: TreeNodeApi = {
 		eventName: K,
 		listener: TreeChangeEvents[K],
 	): Off {
-		return getKernel(node).on(eventName, listener);
+		const kernel = getKernel(node);
+		switch (eventName) {
+			case "nodeChanged": {
+				return kernel.on("childrenChangedAfterBatch", () => {
+					listener();
+				});
+			}
+			case "treeChanged": {
+				return kernel.on("subtreeChangedAfterBatch", () => listener());
+			}
+			default:
+				throw new UsageError(`No event named ${JSON.stringify(eventName)}.`);
+		}
 	},
 	status(node: TreeNode): TreeStatus {
 		return getKernel(node).getStatus();

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

@@ -668,19 +668,10 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 
 	protected abstract get simpleSchema(): T;
 
-	/**
-	 * Generation number which is incremented any time we have an edit on the node.
-	 * Used during iteration to make sure there has been no edits that were concurrently made.
-	 */
-	#generationNumber: number = 0;
-
 	public constructor(
 		input: Iterable<InsertableTreeNodeFromImplicitAllowedTypes<T>> | InternalTreeNode,
 	) {
 		super(input);
-		getKernel(this).on("nodeChanged", () => {
-			this.#generationNumber += 1;
-		});
 	}
 
 	#mapTreesFromFieldData(value: Insertable<T>): ExclusiveMapTree[] {
@@ -882,17 +873,18 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 	}
 
 	public values(): IterableIterator<TreeNodeFromImplicitAllowedTypes<T>> {
-		return this.generateValues(this.#generationNumber);
+		return this.generateValues(getKernel(this).generationNumber);
 	}
 	private *generateValues(
 		initialLastUpdatedStamp: number,
 	): Generator<TreeNodeFromImplicitAllowedTypes<T>> {
-		if (initialLastUpdatedStamp !== this.#generationNumber) {
+		const kernel = getKernel(this);
+		if (initialLastUpdatedStamp !== kernel.generationNumber) {
 			throw new UsageError(`Concurrent editing and iteration is not allowed.`);
 		}
 		for (let i = 0; i < this.length; i++) {
 			yield this.at(i) ?? fail("Index is out of bounds");
-			if (initialLastUpdatedStamp !== this.#generationNumber) {
+			if (initialLastUpdatedStamp !== kernel.generationNumber) {
 				throw new UsageError(`Concurrent editing and iteration is not allowed.`);
 			}
 		}

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

@@ -4,9 +4,15 @@
  */
 
 import { assert } from "@fluidframework/core-utils/internal";
-import { createEmitter, type Listenable, type Off } from "../../events/index.js";
-import type { TreeChangeEvents, TreeNode, Unhydrated } from "./types.js";
-import type { AnchorNode } from "../../core/index.js";
+import {
+	createEmitter,
+	type HasListeners,
+	type IEmitter,
+	type Listenable,
+	type Off,
+} from "../../events/index.js";
+import type { TreeNode, Unhydrated } from "./types.js";
+import type { AnchorEvents, AnchorNode } from "../../core/index.js";
 import {
 	flexTreeSlot,
 	isFreedSymbol,
@@ -60,12 +66,54 @@ export function tryGetTreeNodeSchema(value: unknown): undefined | TreeNodeSchema
  * The kernel has the same lifetime as the node and spans both its unhydrated and hydrated states.
  * When hydration occurs, the kernel is notified via the {@link TreeNodeKernel.hydrate | hydrate} method.
  */
-export class TreeNodeKernel implements Listenable<TreeChangeEvents> {
+export class TreeNodeKernel implements Listenable<KernelEvents> {
+	private disposed = false;
+
+	/**
+	 * Generation number which is incremented any time we have an edit on the node.
+	 * Used during iteration to make sure there has been no edits that were concurrently made.
+	 * @remarks
+	 * This is updated monotonically by this class when edits are applied.
+	 * TODO: update this when applying edits to unhydrated trees.
+	 *
+	 * If TypeScript supported making this immutable from outside the class without making it readonly from inside, that would be used here,
+	 * but they only way to do that is add a separate public accessor and make it private, which was deemed not worth the boilerplate, runtime overhead and bundle size.
+	 */
+	public generationNumber: number = 0;
+
 	#hydrated?: {
 		anchorNode: AnchorNode;
-		offAnchorNode: Off;
+		offAnchorNode: Set<Off>;
 	};
-	#events = createEmitter<TreeChangeEvents>();
+
+	/**
+	 * Events registered before hydration.
+	 * @remarks
+	 * As an optimization these are allocated lazily as they are usually unused.
+	 */
+	#preHydrationEvents?: Listenable<KernelEvents> &
+		IEmitter<KernelEvents> &
+		HasListeners<KernelEvents>;
+
+	/**
+	 * Get the listener.
+	 * @remarks
+	 * If before hydration, allocates and uses `#preHydrationEvents`, otherwise the anchorNode.
+	 * This design avoids allocating `#preHydrationEvents` if unneeded.
+	 *
+	 * This design also avoids extra forwarding overhead for events from anchorNode and also
+	 * avoids registering for events that are unneeded.
+	 * This means optimizations like skipping processing data in subtrees where no subtreeChanged events are subscribed to would be able to work,
+	 * since this code does not unconditionally subscribe to those events (like a design simply forwarding all events would).
+	 */
+	get #events(): Listenable<KernelEvents> {
+		if (this.#hydrated === undefined) {
+			this.#preHydrationEvents ??= createEmitter<KernelEvents>();
+			return this.#preHydrationEvents;
+		} else {
+			return this.#hydrated.anchorNode;
+		}
+	}
 
 	/**
 	 * Create a TreeNodeKernel which can be looked up with {@link getKernel}.
@@ -80,37 +128,50 @@ export class TreeNodeKernel implements Listenable<TreeChangeEvents> {
 		treeNodeToKernel.set(node, this);
 	}
 
+	/**
+	 * Transition from {@link Unhydrated} to hydrated.
+	 * @remarks
+	 * Happens at most once for any given node.
+	 */
 	public hydrate(anchorNode: AnchorNode): void {
-		const offChildrenChanged = anchorNode.on("childrenChangedAfterBatch", () => {
-			this.#events.emit("nodeChanged");
-		});
-
-		const offSubtreeChanged = anchorNode.on("subtreeChangedAfterBatch", () => {
-			this.#events.emit("treeChanged");
-		});
-
-		const offAfterDestroy = anchorNode.on("afterDestroy", () => this.dispose());
+		assert(!this.disposed, "cannot hydrate a disposed node");
 
 		this.#hydrated = {
 			anchorNode,
-			offAnchorNode: () => {
-				offChildrenChanged();
-				offSubtreeChanged();
-				offAfterDestroy();
-			},
+			offAnchorNode: new Set([
+				anchorNode.on("afterDestroy", () => this.dispose()),
+				// TODO: this should be triggered on change even for unhydrated nodes.
+				anchorNode.on("childrenChanging", () => {
+					this.generationNumber += 1;
+				}),
+			]),
 		};
-	}
 
-	public dehydrate(): void {
-		this.#hydrated?.offAnchorNode?.();
-		this.#hydrated = undefined;
+		// If needed, register forwarding emitters for events from before hydration
+		if (this.#preHydrationEvents !== undefined) {
+			for (const eventName of kernelEvents) {
+				if (this.#preHydrationEvents.hasListeners(eventName)) {
+					this.#hydrated.offAnchorNode.add(
+						// Argument is forwarded between matching events, so the type should be correct.
+						// eslint-disable-next-line @typescript-eslint/no-explicit-any
+						anchorNode.on(eventName, (arg: any) =>
+							this.#preHydrationEvents?.emit(eventName, arg),
+						),
+					);
+				}
+			}
+		}
 	}
 
 	public isHydrated(): boolean {
+		assert(!this.disposed, "cannot use a disposed node");
 		return this.#hydrated !== undefined;
 	}
 
 	public getStatus(): TreeStatus {
+		if (this.disposed) {
+			return TreeStatus.Deleted;
+		}
 		if (this.#hydrated?.anchorNode === undefined) {
 			return TreeStatus.New;
 		}
@@ -127,15 +188,19 @@ export class TreeNodeKernel implements Listenable<TreeChangeEvents> {
 		return treeStatusFromAnchorCache(this.#hydrated.anchorNode);
 	}
 
-	public on<K extends keyof TreeChangeEvents>(
-		eventName: K,
-		listener: TreeChangeEvents[K],
-	): Off {
+	public on<K extends keyof KernelEvents>(eventName: K, listener: KernelEvents[K]): Off {
 		return this.#events.on(eventName, listener);
 	}
 
 	public dispose(): void {
-		this.dehydrate();
+		this.disposed = true;
+		for (const off of this.#hydrated?.offAnchorNode ?? []) {
+			off();
+		}
 		// TODO: go to the context and remove myself from withAnchors
 	}
 }
+
+const kernelEvents = ["childrenChangedAfterBatch", "subtreeChangedAfterBatch"] as const;
+
+type KernelEvents = Pick<AnchorEvents, (typeof kernelEvents)[number]>;