tree: Add UnhydratedContext (microsoft#22422)

## Description

Make flex trees have a context, even when unhydrated, just less context.

This UnhydratedContext is used to carry information from the simple-tree
schema to all the flex tree nodes in the subtree.

This will ensure all flex-tree entities have a way to look up stored
schema based on identifiers. This is needed to migrate flex-tree to
store schema since stored schema refers to tree node schema by
identifier instead of object reference.

packages/dds/tree/src/feature-libraries/flex-map-tree/index.ts

@@ -10,4 +10,5 @@ export {
 	isMapTreeSequenceField,
 	getOrCreateMapTreeNode,
 	tryGetMapTreeNode,
+	UnhydratedContext,
 } from "./mapTreeNode.js";

packages/dds/tree/src/feature-libraries/flex-map-tree/mapTreeNode.ts

@@ -4,6 +4,8 @@
  */
 
 import { assert, oob } from "@fluidframework/core-utils/internal";
+import { UsageError } from "@fluidframework/telemetry-utils/internal";
+
 import {
 	type AnchorNode,
 	EmptyKey,
@@ -13,10 +15,12 @@ import {
 	type MapTree,
 	type TreeFieldStoredSchema,
 	type TreeNodeSchemaIdentifier,
+	type TreeStoredSchema,
 	type Value,
 } from "../../core/index.js";
 import { brand, fail, getOrCreate, isReadonlyArray, mapIterable } from "../../util/index.js";
 import {
+	type FlexTreeContext,
 	FlexTreeEntityKind,
 	type FlexTreeField,
 	type FlexTreeNode,
@@ -27,16 +31,22 @@ import {
 	type FlexTreeUnknownUnboxed,
 	flexTreeMarker,
 	indexForAt,
+	type FlexTreeHydratedContext,
 } from "../flex-tree/index.js";
 import {
 	type FlexAllowedTypes,
 	FlexFieldSchema,
 	type FlexTreeNodeSchema,
+	type FlexTreeSchema,
+	intoStoredSchemaCollection,
 	isLazy,
 } from "../typed-schema/index.js";
 import type { FlexFieldKind } from "../modular-schema/index.js";
-import { FieldKinds, type SequenceFieldEditBuilder } from "../default-schema/index.js";
-import { UsageError } from "@fluidframework/telemetry-utils/internal";
+import {
+	defaultSchemaPolicy,
+	FieldKinds,
+	type SequenceFieldEditBuilder,
+} from "../default-schema/index.js";
 
 // #region Nodes
 
@@ -113,6 +123,7 @@ export class EagerMapTreeNode implements MapTreeNode {
 	 * Instead, it should always be acquired via {@link getOrCreateNode}.
 	 */
 	public constructor(
+		public readonly context: UnhydratedContext,
 		public readonly flexSchema: FlexTreeNodeSchema,
 		/** The underlying {@link MapTree} that this `MapTreeNode` reads its data from */
 		public readonly mapTree: ExclusiveMapTree,
@@ -201,8 +212,6 @@ export class EagerMapTreeNode implements MapTreeNode {
 		return this.mapTree.value;
 	}
 
-	public context = undefined;
-
 	public get anchorNode(): AnchorNode {
 		// This API is relevant to `LazyTreeNode`s, but not `MapTreeNode`s.
 		// TODO: Refactor the FlexTreeNode interface so that stubbing this out isn't necessary.
@@ -214,6 +223,7 @@ export class EagerMapTreeNode implements MapTreeNode {
 			const field = getOrCreateField(this, key, this.flexSchema.getFieldSchema(key));
 			for (let index = 0; index < field.length; index++) {
 				const child = getOrCreateChild(
+					this.context,
 					mapTrees[index] ?? oob(),
 					this.flexSchema.getFieldSchema(key).allowedTypes,
 					{ parent: field, index },
@@ -233,6 +243,29 @@ export class EagerMapTreeNode implements MapTreeNode {
 
 // #endregion Nodes
 
+/**
+ * Implementation of `FlexTreeContext`.
+ *
+ * @remarks An editor is required to edit the FlexTree.
+ */
+export class UnhydratedContext implements FlexTreeContext {
+	public readonly schema: TreeStoredSchema;
+
+	/**
+	 * @param flexSchema - Schema to use when working with the tree.
+	 */
+	public constructor(public readonly flexSchema: FlexTreeSchema) {
+		this.schema = {
+			rootFieldSchema: flexSchema.rootFieldSchema.stored,
+			...intoStoredSchemaCollection(flexSchema),
+		};
+	}
+
+	public isHydrated(): this is FlexTreeHydratedContext {
+		return false;
+	}
+}
+
 // #region Fields
 
 /**
@@ -244,6 +277,13 @@ interface MapTreeField extends FlexTreeField {
 	readonly mapTrees: readonly MapTree[];
 }
 
+const emptyContext = new UnhydratedContext({
+	adapters: {},
+	nodeSchema: new Map(),
+	policy: defaultSchemaPolicy,
+	rootFieldSchema: FlexFieldSchema.empty,
+});
+
 /**
  * A special singleton that is the implicit {@link LocationInField} of all un-parented {@link EagerMapTreeNode}s.
  * @remarks This exists because {@link EagerMapTreeNode.parentField} must return a field.
@@ -272,7 +312,7 @@ const unparentedLocation: LocationInField = {
 		},
 		flexSchema: FlexFieldSchema.empty,
 		schema: FlexFieldSchema.empty.stored,
-		context: undefined,
+		context: emptyContext,
 		mapTrees: [],
 		getFieldPath() {
 			fail("unsupported");
@@ -289,6 +329,7 @@ class EagerMapTreeField implements MapTreeField {
 	}
 
 	public constructor(
+		public readonly context: UnhydratedContext,
 		public readonly flexSchema: FlexFieldSchema,
 		public readonly key: FieldKey,
 		public readonly parent: EagerMapTreeNode,
@@ -330,7 +371,7 @@ class EagerMapTreeField implements MapTreeField {
 		return this.mapTrees
 			.map(
 				(m, index) =>
-					getOrCreateChild(m, this.flexSchema.allowedTypes, {
+					getOrCreateChild(this.context, m, this.flexSchema.allowedTypes, {
 						parent: this,
 						index,
 					}) as FlexTreeNode,
@@ -345,15 +386,13 @@ class EagerMapTreeField implements MapTreeField {
 		}
 		const m = this.mapTrees[i];
 		if (m !== undefined) {
-			return getOrCreateChild(m, this.flexSchema.allowedTypes, {
+			return getOrCreateChild(this.context, m, this.flexSchema.allowedTypes, {
 				parent: this,
 				index: i,
 			}) as FlexTreeNode;
 		}
 	}
 
-	public context: undefined;
-
 	/**
 	 * Mutate this field.
 	 * @param edit - A function which receives the current `MapTree`s that comprise the contents of the field so that it may be mutated.
@@ -498,14 +537,16 @@ export function tryGetMapTreeNode(mapTree: MapTree): MapTreeNode | undefined {
  * @remarks It must conform to the `nodeSchema`.
  */
 export function getOrCreateMapTreeNode(
+	context: UnhydratedContext,
 	nodeSchema: FlexTreeNodeSchema,
 	mapTree: ExclusiveMapTree,
 ): EagerMapTreeNode {
-	return nodeCache.get(mapTree) ?? createNode(nodeSchema, mapTree, undefined);
+	return nodeCache.get(mapTree) ?? createNode(context, nodeSchema, mapTree, undefined);
 }
 
 /** Helper for creating a `MapTreeNode` given the parent field (e.g. when "walking down") */
 function getOrCreateChild(
+	context: UnhydratedContext,
 	mapTree: ExclusiveMapTree,
 	allowedTypes: FlexAllowedTypes,
 	parent: LocationInField | undefined,
@@ -523,16 +564,17 @@ function getOrCreateChild(
 				return t.name === mapTree.type;
 			}) ?? fail("Unsupported node schema");
 
-	return createNode(nodeSchema, mapTree, parent);
+	return createNode(context, nodeSchema, mapTree, parent);
 }
 
 /** Always constructs a new node, therefore may not be called twice for the same `MapTree`. */
 function createNode(
+	context: UnhydratedContext,
 	nodeSchema: FlexTreeNodeSchema,
 	mapTree: ExclusiveMapTree,
 	parentField: LocationInField | undefined,
 ): EagerMapTreeNode {
-	return new EagerMapTreeNode(nodeSchema, mapTree, parentField);
+	return new EagerMapTreeNode(context, nodeSchema, mapTree, parentField);
 }
 
 /** Creates a field with the given attributes, or returns a cached field if there is one */
@@ -550,18 +592,18 @@ function getOrCreateField(
 		schema.kind.identifier === FieldKinds.required.identifier ||
 		schema.kind.identifier === FieldKinds.identifier.identifier
 	) {
-		return new EagerMapTreeRequiredField(schema, key, parent);
+		return new EagerMapTreeRequiredField(parent.context, schema, key, parent);
 	}
 
 	if (schema.kind.identifier === FieldKinds.optional.identifier) {
-		return new EagerMapTreeOptionalField(schema, key, parent);
+		return new EagerMapTreeOptionalField(parent.context, schema, key, parent);
 	}
 
 	if (schema.kind.identifier === FieldKinds.sequence.identifier) {
-		return new EagerMapTreeSequenceField(schema, key, parent);
+		return new EagerMapTreeSequenceField(parent.context, schema, key, parent);
 	}
 
-	return new EagerMapTreeField(schema, key, parent);
+	return new EagerMapTreeField(parent.context, schema, key, parent);
 }
 
 /** Unboxes leaf nodes to their values */
@@ -575,7 +617,7 @@ function unboxed(
 		return value;
 	}
 
-	return getOrCreateChild(mapTree, schema.allowedTypes, parent);
+	return getOrCreateChild(parent.parent.context, mapTree, schema.allowedTypes, parent);
 }
 
 // #endregion Caching and unboxing utilities

packages/dds/tree/src/feature-libraries/flex-tree/context.ts

@@ -25,15 +25,9 @@ import { makeField } from "./lazyField.js";
 import type { ITreeCheckout } from "../../shared-tree/index.js";
 
 /**
- * A common context of a "forest" of FlexTrees.
- * It handles group operations like transforming cursors into anchors for edits.
+ * Context for FlexTrees.
  */
-export interface FlexTreeContext extends Listenable<ForestEvents> {
-	/**
-	 * Gets the root field of the tree.
-	 */
-	get root(): FlexTreeField;
-
+export interface FlexTreeContext {
 	/**
 	 * Schema used within this context.
 	 * All data must conform to these schema.
@@ -46,9 +40,24 @@ export interface FlexTreeContext extends Listenable<ForestEvents> {
 	 */
 	readonly schema: TreeStoredSchema;
 
-	// TODO: Add more members:
-	// - transaction APIs
-	// - branching APIs
+	/**
+	 * If true, this context is the canonical context instance for a given view,
+	 * and its schema include all schema from the document.
+	 *
+	 * If false, this context was created for use in a unhydrated tree, and the full document schema is unknown.
+	 */
+	isHydrated(): this is FlexTreeHydratedContext;
+}
+
+/**
+ * A common context of a "forest" of FlexTrees.
+ * It handles group operations like transforming cursors into anchors for edits.
+ */
+export interface FlexTreeHydratedContext extends FlexTreeContext, Listenable<ForestEvents> {
+	/**
+	 * Gets the root field of the tree.
+	 */
+	get root(): FlexTreeField;
 
 	readonly nodeKeyManager: NodeKeyManager;
 
@@ -72,7 +81,7 @@ export const ContextSlot = anchorSlot<Context>();
  *
  * @remarks An editor is required to edit the FlexTree.
  */
-export class Context implements FlexTreeContext, IDisposable {
+export class Context implements FlexTreeHydratedContext, IDisposable {
 	public readonly withCursors: Set<LazyEntity> = new Set();
 	public readonly withAnchors: Set<LazyEntity> = new Set();
 
@@ -102,6 +111,10 @@ export class Context implements FlexTreeContext, IDisposable {
 		this.checkout.forest.anchors.slots.set(ContextSlot, this);
 	}
 
+	public isHydrated(): this is FlexTreeHydratedContext {
+		return true;
+	}
+
 	public get schema(): TreeStoredSchema {
 		return this.checkout.storedSchema;
 	}

packages/dds/tree/src/feature-libraries/flex-tree/flexTreeTypes.ts

@@ -70,24 +70,17 @@ export enum FlexTreeEntityKind {
  * Design and document iterator invalidation rules and ordering rules.
  * Providing a custom iterator type with place anchor semantics would be a good approach.
  */
-export interface FlexTreeEntity<out TSchema = unknown> {
+export interface FlexTreeEntity {
 	/**
 	 * Indicates that an object is a specific kind of flex tree FlexTreeEntity.
 	 * This makes it possible to both down cast FlexTreeEntities safely as well as validate if an object is or is not a FlexTreeEntity.
 	 */
 	readonly [flexTreeMarker]: FlexTreeEntityKind;
 
 	/**
-	 * Schema for this entity.
-	 * If well-formed, it must follow this schema.
-	 */
-	readonly flexSchema: TSchema;
-
-	/**
-	 * A common context of a "forest" of FlexTrees.
-	 * @remarks This is `undefined` for unhydrated nodes or fields that have not yet been inserted into the tree.
+	 * A common context of FlexTrees.
 	 */
-	readonly context?: FlexTreeContext;
+	readonly context: FlexTreeContext;
 }
 
 /**
@@ -132,7 +125,7 @@ export enum TreeStatus {
  * the schema aware API may be more ergonomic.
  * All editing is actually done via {@link FlexTreeField}s: the nodes are immutable other than that they contain mutable fields.
  */
-export interface FlexTreeNode extends FlexTreeEntity<FlexTreeNodeSchema> {
+export interface FlexTreeNode extends FlexTreeEntity {
 	readonly [flexTreeMarker]: FlexTreeEntityKind.Node;
 
 	/**
@@ -193,6 +186,12 @@ export interface FlexTreeNode extends FlexTreeEntity<FlexTreeNodeSchema> {
 	 * If well-formed, it must follow this schema.
 	 */
 	readonly schema: TreeNodeSchemaIdentifier;
+
+	/**
+	 * Schema for this entity.
+	 * If well-formed, it must follow this schema.
+	 */
+	readonly flexSchema: FlexTreeNodeSchema;
 }
 
 /**
@@ -213,7 +212,7 @@ export interface FlexTreeNode extends FlexTreeEntity<FlexTreeNodeSchema> {
  * All content in the tree is accessible without down-casting, but if the schema is known,
  * the schema aware API may be more ergonomic.
  */
-export interface FlexTreeField extends FlexTreeEntity<FlexFieldSchema> {
+export interface FlexTreeField extends FlexTreeEntity {
 	readonly [flexTreeMarker]: FlexTreeEntityKind.Field;
 
 	/**
@@ -267,6 +266,12 @@ export interface FlexTreeField extends FlexTreeEntity<FlexFieldSchema> {
 	 * If well-formed, it must follow this schema.
 	 */
 	readonly schema: TreeFieldStoredSchema;
+
+	/**
+	 * Schema for this entity.
+	 * If well-formed, it must follow this schema.
+	 */
+	readonly flexSchema: FlexFieldSchema;
 }
 
 // #region Node Kinds

packages/dds/tree/src/feature-libraries/flex-tree/index.ts

@@ -29,7 +29,13 @@ export {
 	visitIterableTreeWithState,
 } from "./navigation.js";
 
-export { getTreeContext, type FlexTreeContext, Context, ContextSlot } from "./context.js";
+export {
+	getTreeContext,
+	type FlexTreeContext,
+	type FlexTreeHydratedContext,
+	Context,
+	ContextSlot,
+} from "./context.js";
 
 export { type FlexTreeNodeEvents } from "./treeEvents.js";
 

packages/dds/tree/src/feature-libraries/flex-tree/lazyEntity.ts

@@ -44,7 +44,7 @@ export function assertFlexTreeEntityNotFreed(entity: FlexTreeEntity): void {
  * This is a base class for lazy (cursor based) UntypedEntity implementations, which uniformly handles cursors and anchors.
  */
 export abstract class LazyEntity<TSchema = unknown, TAnchor = unknown>
-	implements FlexTreeEntity<TSchema>, IDisposable
+	implements FlexTreeEntity, IDisposable
 {
 	readonly #lazyCursor: ITreeSubscriptionCursor;
 	public readonly [anchorSymbol]: TAnchor;