Emit events from Unhydrated SharedTree Nodes (#22661)

## Description

Currently, unhydrated nodes can be edited but they do not emit any
change events. This PR fixes that by properly emitting events. Changes
include:

* Removing the `anchor` parameter from the relevant events on
`AnchorEvents`. This allows those events to be implemented by things
that don't have anchors (like `UnhydratedFlexTreeNode`) and things that
only sometimes have anchors (like `TreeNodeKernel`). Nobody uses that
parameter currently anyway because it is redundant (it's the same as the
object that the event is being registered on).
* Emitting a change event from `UnhydratedFlexTreeNode` when one of its
fields is edited. This is listened to by the `TreeNodeKernel`, which can
then emit its own corresponding event.
* Adding distinct types for the "unhydrated" and "hydrated" versions of
the state in `TreeNodeKernel`. This makes it more obvious what the state
transition is, and makes the type checking and safety more explicit.
* Adding a `lazy` function helper to better encapsulate the lazy-getting
of the events in TreeNodeKernel.

.changeset/poor-islands-battle.md

@@ -0,0 +1,19 @@
+---
+"@fluidframework/tree": minor
+---
+---
+"section": tree
+---
+
+Unhydrated SharedTree nodes emit change events when edited
+
+Newly-created SharedTree nodes which have not yet been inserted into the tree will now emit `nodeChanged` and `treeChanged` events when they are mutated via editing operations.
+
+```ts
+const node = new Foo({ foo: 3 });
+Tree.on(node, "nodeChanged", () => {
+	console.log("This will fire even before node is inserted!");
+});
+
+node.foo = 4; // log: "This will fire even before node is inserted!";
+```

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

@@ -128,7 +128,6 @@ export interface AnchorEvents {
 	 * Compare to {@link AnchorEvents.childrenChanged} which is emitted in the middle of the batch/delta-visit.
 	 */
 	childrenChangedAfterBatch(arg: {
-		anchor: AnchorNode;
 		changedFields: ReadonlySet<FieldKey>;
 	}): void;
 
@@ -179,7 +178,7 @@ export interface AnchorEvents {
 	 * subtree changed, compared to {@link AnchorEvents.subtreeChanged} or {@link AnchorEvents.subtreeChanging} which
 	 * fire when something _may_ have changed or _may_ be about to change.
 	 */
-	subtreeChangedAfterBatch(anchor: AnchorNode): void;
+	subtreeChangedAfterBatch(): void;
 }
 
 /**
@@ -775,9 +774,9 @@ export class AnchorSet implements Listenable<AnchorSetRootEvents>, AnchorLocator
 									e.changedField ??
 									fail("childrenChangedAfterBatch events should have a changedField"),
 							);
-						node.events.emit(event, { anchor: node, changedFields: new Set(fieldKeys) });
+						node.events.emit(event, { changedFields: new Set(fieldKeys) });
 					} else {
-						node.events.emit(event, node);
+						node.events.emit(event);
 					}
 				}
 			},

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

@@ -3,14 +3,8 @@
  * Licensed under the MIT License.
  */
 
-import { assert } from "@fluidframework/core-utils/internal";
-import {
-	createEmitter,
-	type HasListeners,
-	type IEmitter,
-	type Listenable,
-	type Off,
-} from "../../events/index.js";
+import { assert, Lazy } from "@fluidframework/core-utils/internal";
+import { createEmitter, type Listenable, type Off } from "../../events/index.js";
 import type { TreeNode, Unhydrated } from "./types.js";
 import {
 	anchorSlot,
@@ -76,6 +70,25 @@ export function tryGetTreeNodeSchema(value: unknown): undefined | TreeNodeSchema
 	return kernel?.schema;
 }
 
+/** The {@link HydrationState} of a {@link TreeNodeKernel} before the kernel is hydrated */
+type UnhydratedState = Off;
+
+/** The {@link HydrationState} of a {@link TreeNodeKernel} after the kernel is hydrated */
+interface HydratedState {
+	/** The {@link AnchorNode} that this node is associated with. */
+	anchorNode: AnchorNode;
+	/** All {@link Off | event deregistration functions} that should be run when the kernel is disposed. */
+	offAnchorNode: Set<Off>;
+}
+
+/** State within a {@link TreeNodeKernel} that is related to the hydration process */
+type HydrationState = UnhydratedState | HydratedState;
+
+/** True if and only if the given {@link HydrationState} is post-hydration */
+function isHydrated(state: HydrationState): state is HydratedState {
+	return typeof state === "object";
+}
+
 /**
  * Contains state and an internal API for managing {@link TreeNode}s.
  * @remarks All {@link TreeNode}s have an associated kernel object.
@@ -97,39 +110,17 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 	 */
 	public generationNumber: number = 0;
 
-	#hydrated?: {
-		anchorNode: AnchorNode;
-		offAnchorNode: Set<Off>;
-	};
+	#hydrationState: HydrationState = () => {};
 
 	/**
 	 * 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.
+	 * Since these are usually not used, they are allocated lazily as an optimization.
+	 * The laziness also avoids extra forwarding overhead for events from this kernel's anchor node 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).
+	 * since the kernel does not unconditionally subscribe to those events (like a design which simply forwards all events would).
 	 */
-	get #events(): Listenable<KernelEvents> {
-		if (this.#hydrated === undefined) {
-			this.#preHydrationEvents ??= createEmitter<KernelEvents>();
-			return this.#preHydrationEvents;
-		} else {
-			return this.#hydrated.anchorNode;
-		}
-	}
+	readonly #unhydratedEvents = new Lazy(createEmitter<KernelEvents>);
 
 	/**
 	 * Create a TreeNodeKernel which can be looked up with {@link getKernel}.
@@ -151,6 +142,27 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 		if (innerNode instanceof UnhydratedFlexTreeNode) {
 			// Unhydrated case
 			mapTreeNodeToProxy.set(innerNode, node);
+			// Register for change events from the unhydrated flex node.
+			// These will be fired if the unhydrated node is edited, and will also be forwarded later to the hydrated node.
+			this.#hydrationState = innerNode.events.on(
+				"childrenChangedAfterBatch",
+				({ changedFields }) => {
+					this.#unhydratedEvents.value.emit("childrenChangedAfterBatch", {
+						changedFields,
+					});
+
+					let n: UnhydratedFlexTreeNode | undefined = innerNode;
+					while (n !== undefined) {
+						const treeNode = mapTreeNodeToProxy.get(n);
+						if (treeNode !== undefined) {
+							const kernel = getKernel(treeNode);
+							kernel.#unhydratedEvents.value.emit("subtreeChangedAfterBatch");
+						}
+						// This cast is safe because the parent (if it exists) of an unhydrated flex node is always another unhydrated flex node.
+						n = n.parentField.parent.parent as UnhydratedFlexTreeNode | undefined;
+					}
+				},
+			);
 		} else {
 			// Hydrated case
 			assert(
@@ -162,10 +174,10 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 	}
 
 	public get context(): Context {
-		if (this.#hydrated !== undefined) {
+		if (isHydrated(this.#hydrationState)) {
 			// This can't be cached on this.#hydrated during hydration since initial tree is hydrated before the context is cached on the anchorSet.
 			return (
-				this.#hydrated?.anchorNode.anchorSet.slots.get(SimpleContextSlot) ??
+				this.#hydrationState?.anchorNode.anchorSet.slots.get(SimpleContextSlot) ??
 				fail("missing simple-tree context")
 			);
 		}
@@ -181,7 +193,7 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 	 */
 	public hydrate(anchorNode: AnchorNode): void {
 		assert(!this.disposed, 0xa2a /* cannot hydrate a disposed node */);
-		assert(this.#hydrated === undefined, 0xa2b /* hydration should only happen once */);
+		assert(!isHydrated(this.#hydrationState), 0xa2b /* hydration should only happen once */);
 
 		// If the this node is raw and thus has a MapTreeNode, forget it:
 		if (this.innerNode instanceof UnhydratedFlexTreeNode) {
@@ -190,7 +202,7 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 
 		// However, it's fine for an anchor node to rotate through different proxies when the content at that place in the tree is replaced.
 		anchorNode.slots.set(proxySlot, this.node);
-		this.#hydrated = {
+		this.#hydrationState = {
 			anchorNode,
 			offAnchorNode: new Set([
 				anchorNode.on("afterDestroy", () => this.dispose()),
@@ -202,54 +214,55 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 		};
 
 		// If needed, register forwarding emitters for events from before hydration
-		if (this.#preHydrationEvents !== undefined) {
+		if (this.#unhydratedEvents.evaluated) {
+			const events = this.#unhydratedEvents.value;
 			for (const eventName of kernelEvents) {
-				if (this.#preHydrationEvents.hasListeners(eventName)) {
-					this.#hydrated.offAnchorNode.add(
+				if (events.hasListeners(eventName)) {
+					this.#hydrationState.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),
-						),
+						anchorNode.on(eventName, (arg: any) => events.emit(eventName, arg)),
 					);
 				}
 			}
 		}
 	}
 
-	public isHydrated(): boolean {
-		assert(!this.disposed, 0xa2c /* cannot use a disposed node */);
-		return this.#hydrated !== undefined;
-	}
-
 	public getStatus(): TreeStatus {
 		if (this.disposed) {
 			return TreeStatus.Deleted;
 		}
-		if (this.#hydrated?.anchorNode === undefined) {
+		if (!isHydrated(this.#hydrationState)) {
 			return TreeStatus.New;
 		}
 
 		// TODO: Replace this check with the proper check against the cursor state when the cursor becomes part of the kernel
-		const flex = this.#hydrated.anchorNode.slots.get(flexTreeSlot);
+		const flex = this.#hydrationState.anchorNode.slots.get(flexTreeSlot);
 		if (flex !== undefined) {
 			assert(flex instanceof LazyEntity, 0x9b4 /* Unexpected flex node implementation */);
 			if (flex[isFreedSymbol]()) {
 				return TreeStatus.Deleted;
 			}
 		}
 
-		return treeStatusFromAnchorCache(this.#hydrated.anchorNode);
+		return treeStatusFromAnchorCache(this.#hydrationState.anchorNode);
 	}
 
 	public on<K extends keyof KernelEvents>(eventName: K, listener: KernelEvents[K]): Off {
-		return this.#events.on(eventName, listener);
+		// Retrieve the correct events object based on whether this node is pre or post hydration.
+		const events: Listenable<KernelEvents> = isHydrated(this.#hydrationState)
+			? this.#hydrationState.anchorNode
+			: this.#unhydratedEvents.value;
+
+		return events.on(eventName, listener);
 	}
 
 	public dispose(): void {
 		this.disposed = true;
-		for (const off of this.#hydrated?.offAnchorNode ?? []) {
-			off();
+		if (isHydrated(this.#hydrationState)) {
+			for (const off of this.#hydrationState.offAnchorNode) {
+				off();
+			}
 		}
 		// TODO: go to the context and remove myself from withAnchors
 	}
@@ -268,13 +281,12 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 			return this.innerNode;
 		}
 
-		if (this.#hydrated === undefined) {
-			// Unhydrated case
+		if (!isHydrated(this.#hydrationState)) {
 			return this.innerNode;
 		}
 
 		// Marinated case -> cooked
-		const anchorNode = this.#hydrated.anchorNode;
+		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) {
@@ -336,13 +348,12 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 			return this.innerNode;
 		}
 
-		if (this.#hydrated === undefined) {
-			// Unhydrated case
+		if (!isHydrated(this.#hydrationState)) {
 			return this.innerNode;
 		}
 
 		// Marinated case -> cooked
-		const anchorNode = this.#hydrated.anchorNode;
+		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
 		return anchorNode.slots.get(flexTreeSlot);
 	}

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

@@ -32,11 +32,6 @@ export type Unhydrated<T> = T;
 /**
  * A collection of events that can be emitted by a {@link TreeNode}.
  *
- * @remarks
- * Currently, events can be subscribed to for {@link Unhydrated} nodes, however no events will be triggered for the nodes until after they are hydrated.
- * This is considered a known issue, and should be fixed in future versions.
- * Do not rely on the fact that editing unhydrated nodes does not trigger their events.
- *
  * @privateRemarks
  * TODO: add a way to subscribe to a specific field (for nodeChanged and treeChanged).
  * Probably have object node and map node specific APIs for this.