@fluidframework/tree 2.3.1 → 2.4.0-297027

package/src/simple-tree/core/schemaCaching.ts

@@ -5,68 +5,31 @@
 
 import { assert } from "@fluidframework/core-utils/internal";
 
-import type { FlexTreeNodeSchema, TreeNodeSchemaBase } from "../../feature-libraries/index.js";
 import { fail } from "../../util/index.js";
 
 import type { TreeNodeSchema } from "./treeNodeSchema.js";
 import type { InnerNode } from "./treeNodeKernel.js";
+import { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
+import { SimpleContextSlot, type Context } from "./context.js";
 
 /**
- * A symbol for storing FlexTreeSchema on TreeNodeSchema.
- * Eagerly set on leaves, and lazily set for other cases.
- */
-const flexSchemaSymbol: unique symbol = Symbol(`flexSchema`);
-
-/**
- * A symbol for storing TreeNodeSchema on FlexTreeNode's schema.
+ * Gets the {@link TreeNodeSchema} for the {@link InnerNode}.
  */
-const simpleNodeSchemaSymbol: unique symbol = Symbol(`simpleNodeSchema`);
-
-export function cachedFlexSchemaFromClassSchema(
-	schema: TreeNodeSchema,
-): TreeNodeSchemaBase | undefined {
-	// eslint-disable-next-line @typescript-eslint/no-explicit-any
-	return (schema as any)[flexSchemaSymbol] as TreeNodeSchemaBase | undefined;
-}
-
-export function setFlexSchemaFromClassSchema(
-	simple: TreeNodeSchema,
-	flex: TreeNodeSchemaBase,
-): void {
-	assert(!(flexSchemaSymbol in simple), 0x91f /* simple schema already marked */);
-	assert(!(simpleNodeSchemaSymbol in flex), 0x920 /* flex schema already marked */);
-	// eslint-disable-next-line @typescript-eslint/no-explicit-any
-	(simple as any)[flexSchemaSymbol] = flex;
-	// eslint-disable-next-line @typescript-eslint/no-explicit-any
-	(flex as any)[simpleNodeSchemaSymbol] = simple;
+export function getSimpleNodeSchemaFromInnerNode(innerNode: InnerNode): TreeNodeSchema {
+	const context: Context = getSimpleContextFromInnerNode(innerNode);
+	return context.schema.get(innerNode.schema) ?? fail("missing schema from context");
 }
 
 /**
- * Gets the {@link TreeNodeSchema} cached on the provided {@link FlexTreeNodeSchema | flexSchema}.
- * Returns `undefined` if no cached value is found.
+ * Gets the {@link Context} for the {@link InnerNode}.
  */
-export function tryGetSimpleNodeSchema(
-	flexSchema: FlexTreeNodeSchema,
-): TreeNodeSchema | undefined {
-	if (simpleNodeSchemaSymbol in flexSchema) {
-		return flexSchema[simpleNodeSchemaSymbol] as TreeNodeSchema;
+export function getSimpleContextFromInnerNode(innerNode: InnerNode): Context {
+	if (innerNode instanceof UnhydratedFlexTreeNode) {
+		return innerNode.simpleContext;
 	}
-	return undefined;
-}
 
-/**
- * Gets the {@link TreeNodeSchema} cached on the provided {@link FlexTreeNodeSchema | flexSchema}.
- * Fails if no cached value is found.
- */
-export function getSimpleNodeSchema(flexSchema: FlexTreeNodeSchema): TreeNodeSchema {
-	return tryGetSimpleNodeSchema(flexSchema) ?? fail("missing simple schema");
-}
+	const context = innerNode.anchorNode.anchorSet.slots.get(SimpleContextSlot);
+	assert(context !== undefined, "missing simple tree context");
 
-/**
- * Gets the {@link TreeNodeSchema} for the {@link InnerNode}.
- */
-export function getSimpleNodeSchemaFromNode(innerNode: InnerNode): TreeNodeSchema {
-	// TODO: to make this work without depending on flex tree schema, a new caching/lookup mechanism will be required, likely leveraging the FlexTreeContext:
-	// A SimpleTreeContext could be defined and associated with the FlexTreeContext, and used to look up simple-tree schema by identifier.
-	return getSimpleNodeSchema(innerNode.flexSchema);
+	return context;
 }

package/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,
@@ -25,18 +19,18 @@ import {
 	flexTreeSlot,
 	isFlexTreeNode,
 	isFreedSymbol,
-	isMapTreeNode,
 	LazyEntity,
 	TreeStatus,
 	treeStatusFromAnchorCache,
 	type FlexTreeNode,
-	type MapTreeNode,
 } from "../../feature-libraries/index.js";
 import type { TreeNodeSchema } from "./treeNodeSchema.js";
 import { fail } from "../../util/index.js";
 // TODO: decide how to deal with dependencies on flex-tree implementation.
 // eslint-disable-next-line import/no-internal-modules
 import { makeTree } from "../../feature-libraries/flex-tree/lazyNode.js";
+import { SimpleContextSlot, type Context } from "./context.js";
+import { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
 
 const treeNodeToKernel = new WeakMap<TreeNode, TreeNodeKernel>();
 
@@ -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,43 +110,22 @@ 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}.
 	 *
+	 * @param initialContext - context from when this node was originally crated.
 	 * @param innerNode - When unhydrated/raw or marinated the MapTreeNode. FlexTreeNode when cooked.
 	 * @remarks
 	 * Exactly one kernel per TreeNode should be created.
@@ -142,13 +134,35 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 		public readonly node: TreeNode,
 		public readonly schema: TreeNodeSchema,
 		private innerNode: InnerNode,
+		private readonly initialContext: Context,
 	) {
 		assert(!treeNodeToKernel.has(node), 0xa1a /* only one kernel per node can be made */);
 		treeNodeToKernel.set(node, this);
 
-		if (isMapTreeNode(innerNode)) {
+		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(
@@ -159,26 +173,36 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 		}
 	}
 
+	public get context(): Context {
+		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.#hydrationState?.anchorNode.anchorSet.slots.get(SimpleContextSlot) ??
+				fail("missing simple-tree context")
+			);
+		}
+		return this.initialContext;
+	}
+
 	/**
 	 * Transition from {@link Unhydrated} to hydrated.
 	 * Bi-directionally associates the given hydrated TreeNode to the given anchor node.
 	 * @remarks
 	 * Happens at most once for any given node.
-	 * Cleans up mappings to {@link MapTreeNode} - it is assumed that they are no longer needed once the proxy has an anchor node.
+	 * Cleans up mappings to {@link UnhydratedFlexTreeNode} - it is assumed that they are no longer needed once the proxy has an anchor node.
 	 */
 	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 (isMapTreeNode(this.innerNode)) {
+		if (this.innerNode instanceof UnhydratedFlexTreeNode) {
 			mapTreeNodeToProxy.delete(this.innerNode);
 		}
 
 		// 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()),
@@ -190,36 +214,30 @@ 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]()) {
@@ -227,17 +245,24 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 			}
 		}
 
-		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
 	}
@@ -251,18 +276,17 @@ export class TreeNodeKernel implements Listenable<KernelEvents> {
 	 * Note that for "marinated" nodes, this FlexTreeNode exists and returns it: it does not return the MapTreeNode which is the current InnerNode.
 	 */
 	public getOrCreateInnerNode(allowFreed = false): InnerNode {
-		if (!isMapTreeNode(this.innerNode)) {
+		if (!(this.innerNode instanceof UnhydratedFlexTreeNode)) {
 			// Cooked case
 			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) {
@@ -324,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);
 	}
@@ -357,12 +380,12 @@ type KernelEvents = Pick<AnchorEvents, (typeof kernelEvents)[number]>;
  * TODO: The inconsistent handling of "marinated" cases should be cleaned up.
  * Maybe getOrCreateInnerNode should cook marinated nodes so they have a proper InnerNode?
  */
-export type InnerNode = FlexTreeNode | MapTreeNode;
+export type InnerNode = FlexTreeNode | UnhydratedFlexTreeNode;
 
 /**
  * {@inheritdoc proxyToMapTreeNode}
  */
-const mapTreeNodeToProxy = new WeakMap<MapTreeNode, TreeNode>();
+export const mapTreeNodeToProxy = new WeakMap<UnhydratedFlexTreeNode, TreeNode>();
 
 /**
  * An anchor slot which associates an anchor with its corresponding TreeNode, if there is one.
@@ -370,18 +393,15 @@ const mapTreeNodeToProxy = new WeakMap<MapTreeNode, TreeNode>();
  * For this to work, we have to require that there is at most a single view using a given AnchorSet.
  * FlexTree already has this assumption, and we also assume there is a single simple-tree per FlexTree, so this is valid.
  */
-const proxySlot = anchorSlot<TreeNode>();
+export const proxySlot = anchorSlot<TreeNode>();
 
 /**
- * Retrieves the proxy associated with the given flex node via {@link setInnerNode}, if any.
+ * Retrieves the node associated with the given MapTreeNode node if any.
  */
-export function tryGetCachedTreeNode(flexNode: InnerNode): TreeNode | undefined {
-	if (isMapTreeNode(flexNode)) {
-		// Unhydrated case
-		return mapTreeNodeToProxy.get(flexNode);
-	}
-	// Hydrated case
-	return flexNode.anchorNode.slots.get(proxySlot);
+export function tryGetTreeNodeFromMapNode(
+	flexNode: UnhydratedFlexTreeNode,
+): TreeNode | undefined {
+	return mapTreeNodeToProxy.get(flexNode);
 }
 
 export function tryDisposeTreeNode(anchorNode: AnchorNode): void {
@@ -391,3 +411,22 @@ export function tryDisposeTreeNode(anchorNode: AnchorNode): void {
 		kernel.dispose();
 	}
 }
+
+/**
+ * Lookup a TreeNodeSchema from a Hydrated FlexTreeNode.
+ * @privateRemarks
+ * This provides a way to access simple tree schema from the flex tree without depending on {@link FlexTreeSchema} which is in the process of being removed.
+ * This is currently limited to hydrated nodes: this limitation will have to be fixed before {@link FlexTreeSchema} can be fully removed.
+ */
+export function getTreeNodeSchemaFromHydratedFlexNode(flexNode: FlexTreeNode): TreeNodeSchema {
+	assert(
+		flexNode.context.isHydrated(),
+		"getTreeNodeSchemaFromHydratedFlexNode only allows hydrated flex tree nodes",
+	);
+
+	const context =
+		flexNode.anchorNode.anchorSet.slots.get(SimpleContextSlot) ??
+		fail("Missing SimpleContextSlot");
+
+	return context.schema.get(flexNode.schema) ?? fail("Missing schema");
+}

package/src/simple-tree/core/treeNodeSchema.ts

@@ -108,13 +108,26 @@ export interface TreeNodeSchemaClass<
 	 */
 	new (data: TInsertable | InternalTreeNode): Unhydrated<TNode>;
 }
+
+/**
+ * Internal helper for utilities that return schema which can be used in class and non class formats depending on the API exposing it.
+ */
+export type TreeNodeSchemaBoth<
+	Name extends string = string,
+	Kind extends NodeKind = NodeKind,
+	TNode = unknown,
+	TInsertable = never,
+	ImplicitlyConstructable extends boolean = boolean,
+	Info = unknown,
+> = TreeNodeSchemaClass<Name, Kind, TNode, TInsertable, ImplicitlyConstructable, Info> &
+	TreeNodeSchemaNonClass<Name, Kind, TNode, TInsertable, ImplicitlyConstructable, Info>;
+
 /**
  * Data common to all tree node schema.
  * @remarks
  * Implementation detail of {@link TreeNodeSchema} which should be accessed instead of referring to this type directly.
  * @sealed @public
  */
-
 export interface TreeNodeSchemaCore<
 	out Name extends string,
 	out Kind extends NodeKind,
@@ -154,6 +167,27 @@ export interface TreeNodeSchemaCore<
 	 * Setting this to false adjusts the insertable types to disallow cases which could be impacted by these inconsistencies.
 	 */
 	readonly implicitlyConstructable: ImplicitlyConstructable;
+
+	/**
+	 * All possible schema that a direct child of a node with this schema could have.
+	 *
+	 * Equivalently, this is also all schema directly referenced when defining this schema's allowed child types,
+	 * which is also the same as the set of schema referenced directly by the `Info` type parameter and the `info` property.
+	 * This property is simply re-exposing that information in an easier to traverse format consistent across all node kinds.
+	 * @remarks
+	 * Some kinds of nodes may have additional restrictions on children:
+	 * this set simply enumerates all directly referenced schema, and can be use to walk over all referenced schema types.
+	 *
+	 * This set cannot be used before the schema in it have been defined:
+	 * more specifically, when using lazy schema references (for example to make foreword references to schema which have not yet been defined),
+	 * users must wait until after the schema are defined to access this set.
+	 * @privateRemarks
+	 * Currently there isn't much use for this in the public API,
+	 * and it's possible this will want to be tweaked or renamed as part of a larger schema reflection API surface that might be added later.
+	 * To keep options option, this is marked `@system` for now.
+	 * @system
+	 */
+	readonly childTypes: ReadonlySet<TreeNodeSchema>;
 }
 
 /**

package/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.