tree: Remove FlexTreeNode.flexSchema and its usage (microsoft#22624)

## Description

Remove FlexTreeNode.flexSchema and its usage.

This required getting simple-tree schema from the simple tree context,
which got its implementation finished as part of this.

Additionally fixed a bug in POJO mode object nodes' symbol support.

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

@@ -194,12 +194,6 @@ export interface FlexTreeNode extends FlexTreeEntity {
 	 * 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;
 }
 
 /**

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

@@ -19,7 +19,6 @@ import {
 	getOrCreateNodeFromInnerNode,
 	UnhydratedFlexTreeNode,
 	type Unhydrated,
-	UnhydratedContext,
 } from "../core/index.js";
 import {
 	cursorForMapTreeNode,
@@ -39,6 +38,7 @@ import {
 	type VerboseTree,
 	type VerboseTreeNode,
 } from "./verboseTree.js";
+import { getUnhydratedContext } from "../createContext.js";
 
 /**
  * Construct tree content that is compatible with the field defined by the provided `schema`.
@@ -151,12 +151,12 @@ export function createFromCursor<TSchema extends ImplicitFieldSchema>(
 	cursor: ITreeCursorSynchronous | undefined,
 ): Unhydrated<TreeFieldFromImplicitField<TSchema>> {
 	const mapTrees = cursor === undefined ? [] : [mapTreeFromCursor(cursor)];
-	const flexSchema = toFlexSchema(schema);
+	const context = getUnhydratedContext(schema);
+	const flexSchema = context.flexContext.flexSchema;
 
 	const schemaValidationPolicy: SchemaAndPolicy = {
 		policy: defaultSchemaPolicy,
-		// TODO: optimize: This isn't the most efficient operation since its not cached, and has to convert all the schema.
-		schema: intoStoredSchema(flexSchema),
+		schema: context.flexContext.schema,
 	};
 
 	const maybeError = isFieldInSchema(
@@ -174,8 +174,7 @@ export function createFromCursor<TSchema extends ImplicitFieldSchema>(
 	// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 	const mapTree = mapTrees[0]!;
 	const mapTreeNode = UnhydratedFlexTreeNode.getOrCreate(
-		// TODO: Provide a way to get simple-tree context here, then make UnhydratedFlexTreeNode's hold simple-tree contexts. Use this for InnerNode -> TreeSchemaSchema
-		new UnhydratedContext(flexSchema),
+		getUnhydratedContext(schema),
 		mapTree,
 	);
 

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

@@ -11,7 +11,6 @@ import {
 	type TreeStatus,
 	isLazy,
 	isTreeValue,
-	FlexObjectNodeSchema,
 	FieldKinds,
 } from "../../feature-libraries/index.js";
 import { fail, extractFromOpaque, isReadonlyArray } from "../../util/index.js";
@@ -43,6 +42,7 @@ import {
 	tryGetTreeNodeSchema,
 	getOrCreateNodeFromInnerNode,
 	UnhydratedFlexTreeNode,
+	typeSchemaSymbol,
 } from "../core/index.js";
 import { isObjectNodeSchema } from "../objectNodeTypes.js";
 
@@ -230,10 +230,13 @@ export const treeNodeApi: TreeNodeApi = {
 		return tryGetSchema(node) ?? fail("Not a tree node");
 	},
 	shortId(node: TreeNode): number | string | undefined {
+		const schema = node[typeSchemaSymbol];
+		if (!isObjectNodeSchema(schema)) {
+			return undefined;
+		}
+
 		const flexNode = getOrCreateInnerNode(node);
-		const flexSchema = flexNode.flexSchema;
-		const identifierFieldKeys =
-			flexSchema instanceof FlexObjectNodeSchema ? flexSchema.identifierFieldKeys : [];
+		const identifierFieldKeys = schema.identifierFieldKeys;
 
 		switch (identifierFieldKeys.length) {
 			case 0:

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

@@ -42,7 +42,7 @@ import {
 	UnhydratedTreeSequenceField,
 } from "./core/index.js";
 import { TreeNodeValid, type MostDerivedData } from "./treeNodeValid.js";
-import { createUnhydratedContext } from "./createContext.js";
+import { getUnhydratedContext } from "./createContext.js";
 
 /**
  * A generic array type, used to defined types like {@link (TreeArrayNode:interface)}.
@@ -956,7 +956,7 @@ export function arraySchema<
 			input: T2,
 		): UnhydratedFlexTreeNode {
 			return UnhydratedFlexTreeNode.getOrCreate(
-				unhydratedContext.flexContext,
+				unhydratedContext,
 				mapTreeFromNodeData(input as object, this as unknown as ImplicitAllowedTypes),
 			);
 		}
@@ -965,7 +965,7 @@ export function arraySchema<
 
 		protected static override oneTimeSetup<T2>(this: typeof TreeNodeValid<T2>): Context {
 			const schema = this as unknown as TreeNodeSchema;
-			unhydratedContext = createUnhydratedContext(schema);
+			unhydratedContext = getUnhydratedContext(schema);
 
 			// First run, do extra validation.
 			// TODO: provide a way for TreeConfiguration to trigger this same validation to ensure it gets run early.

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

@@ -10,6 +10,8 @@ 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.
@@ -66,7 +68,20 @@ export function getSimpleNodeSchema(flexSchema: FlexTreeNodeSchema): TreeNodeSch
  * Gets the {@link TreeNodeSchema} for the {@link InnerNode}.
  */
 export function getSimpleNodeSchemaFromInnerNode(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);
+	const context: Context = getSimpleContextFromInnerNode(innerNode);
+	return context.schema.get(innerNode.schema) ?? fail("missing schema from context");
+}
+
+/**
+ * Gets the {@link Context} for the {@link InnerNode}.
+ */
+export function getSimpleContextFromInnerNode(innerNode: InnerNode): Context {
+	if (innerNode instanceof UnhydratedFlexTreeNode) {
+		return innerNode.simpleContext;
+	}
+
+	const context = innerNode.anchorNode.anchorSet.slots.get(SimpleContextSlot);
+	assert(context !== undefined, "missing simple tree context");
+
+	return context;
 }

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

@@ -37,10 +37,10 @@ import {
 	type FlexTreeSchema,
 	intoStoredSchemaCollection,
 	type FlexFieldKind,
-	defaultSchemaPolicy,
 	FieldKinds,
 	type SequenceFieldEditBuilder,
 } from "../../feature-libraries/index.js";
+import type { Context } from "./context.js";
 
 interface UnhydratedTreeSequenceFieldEditBuilder
 	extends SequenceFieldEditBuilder<ExclusiveMapTree[]> {
@@ -85,12 +85,16 @@ export class UnhydratedFlexTreeNode implements UnhydratedFlexTreeNode {
 	 * @remarks It must conform to the `nodeSchema`.
 	 */
 	public static getOrCreate(
-		context: UnhydratedContext,
+		context: Context,
 		mapTree: ExclusiveMapTree,
 	): UnhydratedFlexTreeNode {
 		return nodeCache.get(mapTree) ?? new UnhydratedFlexTreeNode(context, mapTree, undefined);
 	}
 
+	public get context(): UnhydratedContext {
+		return this.simpleContext.flexContext;
+	}
+
 	/**
 	 * Create a new UnhydratedFlexTreeNode.
 	 * @param location - the parentage of this node, if it is being created underneath an existing node and field, or undefined if not
@@ -100,7 +104,7 @@ export class UnhydratedFlexTreeNode implements UnhydratedFlexTreeNode {
 	 * Instead, it should always be acquired via {@link getOrCreateNode}.
 	 */
 	public constructor(
-		public readonly context: UnhydratedContext,
+		public readonly simpleContext: Context,
 		/** The underlying {@link MapTree} that this `UnhydratedFlexTreeNode` reads its data from */
 		public readonly mapTree: ExclusiveMapTree,
 		private location = unparentedLocation,
@@ -198,7 +202,7 @@ export class UnhydratedFlexTreeNode implements UnhydratedFlexTreeNode {
 		for (const [key, mapTrees] of this.mapTree.fields) {
 			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(), {
+				const child = getOrCreateChild(this.simpleContext, mapTrees[index] ?? oob(), {
 					parent: field,
 					index,
 				});
@@ -240,13 +244,6 @@ export class UnhydratedContext implements FlexTreeContext {
 
 // #region Fields
 
-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 UnhydratedFlexTreeNode}s.
  * @remarks This exists because {@link UnhydratedFlexTreeNode.parentField} must return a field.
@@ -271,7 +268,9 @@ const unparentedLocation: LocationInField = {
 			return undefined;
 		},
 		schema: FlexFieldSchema.empty.stored,
-		context: emptyContext,
+		get context(): never {
+			return fail("unsupported");
+		},
 		getFieldPath() {
 			fail("unsupported");
 		},
@@ -286,8 +285,12 @@ class UnhydratedFlexTreeField implements FlexTreeField {
 		return this.flexSchema.stored;
 	}
 
+	public get context(): UnhydratedContext {
+		return this.simpleContext.flexContext;
+	}
+
 	public constructor(
-		public readonly context: UnhydratedContext,
+		public readonly simpleContext: Context,
 		public readonly flexSchema: FlexFieldSchema,
 		public readonly key: FieldKey,
 		public readonly parent: UnhydratedFlexTreeNode,
@@ -325,7 +328,7 @@ class UnhydratedFlexTreeField implements FlexTreeField {
 		return this.mapTrees
 			.map(
 				(m, index) =>
-					getOrCreateChild(this.context, m, {
+					getOrCreateChild(this.simpleContext, m, {
 						parent: this,
 						index,
 					}) as FlexTreeNode,
@@ -340,7 +343,7 @@ class UnhydratedFlexTreeField implements FlexTreeField {
 		}
 		const m = this.mapTrees[i];
 		if (m !== undefined) {
-			return getOrCreateChild(this.context, m, {
+			return getOrCreateChild(this.simpleContext, m, {
 				parent: this,
 				index: i,
 			}) as FlexTreeNode;
@@ -370,16 +373,14 @@ class UnhydratedFlexTreeField implements FlexTreeField {
 	}
 
 	/** Unboxes leaf nodes to their values */
-	protected unboxed(
-		mapTree: ExclusiveMapTree,
-		parent: LocationInField,
-	): FlexTreeUnknownUnboxed {
+	protected unboxed(index: number): FlexTreeUnknownUnboxed {
+		const mapTree: ExclusiveMapTree = this.mapTrees[index] ?? oob();
 		const value = mapTree.value;
 		if (value !== undefined) {
 			return value;
 		}
 
-		return getOrCreateChild(parent.parent.context, mapTree, parent);
+		return getOrCreateChild(this.simpleContext, mapTree, { parent: this, index });
 	}
 }
 
@@ -412,10 +413,7 @@ class EagerMapTreeOptionalField
 	public get content(): FlexTreeUnknownUnboxed | undefined {
 		const value = this.mapTrees[0];
 		if (value !== undefined) {
-			return this.unboxed(value, {
-				parent: this,
-				index: 0,
-			});
+			return this.unboxed(0);
 		}
 
 		return undefined;
@@ -471,15 +469,15 @@ export class UnhydratedTreeSequenceField
 		if (i === undefined) {
 			return undefined;
 		}
-		return this.unboxed(this.mapTrees[i] ?? oob(), { parent: this, index: i });
+		return this.unboxed(i);
 	}
 	public map<U>(callbackfn: (value: FlexTreeUnknownUnboxed, index: number) => U): U[] {
 		return Array.from(this, callbackfn);
 	}
 
 	public *[Symbol.iterator](): IterableIterator<FlexTreeUnknownUnboxed> {
-		for (const [i, mapTree] of this.mapTrees.entries()) {
-			yield this.unboxed(mapTree, { parent: this, index: i });
+		for (const [i] of this.mapTrees.entries()) {
+			yield this.unboxed(i);
 		}
 	}
 }
@@ -512,7 +510,7 @@ export function tryUnhydratedFlexTreeNode(
 
 /** Helper for creating a `UnhydratedFlexTreeNode` given the parent field (e.g. when "walking down") */
 function getOrCreateChild(
-	context: UnhydratedContext,
+	context: Context,
 	mapTree: ExclusiveMapTree,
 	parent: LocationInField | undefined,
 ): UnhydratedFlexTreeNode {
@@ -539,18 +537,18 @@ function getOrCreateField(
 		schema.kind.identifier === FieldKinds.required.identifier ||
 		schema.kind.identifier === FieldKinds.identifier.identifier
 	) {
-		return new EagerMapTreeRequiredField(parent.context, schema, key, parent);
+		return new EagerMapTreeRequiredField(parent.simpleContext, schema, key, parent);
 	}
 
 	if (schema.kind.identifier === FieldKinds.optional.identifier) {
-		return new EagerMapTreeOptionalField(parent.context, schema, key, parent);
+		return new EagerMapTreeOptionalField(parent.simpleContext, schema, key, parent);
 	}
 
 	if (schema.kind.identifier === FieldKinds.sequence.identifier) {
-		return new UnhydratedTreeSequenceField(parent.context, schema, key, parent);
+		return new UnhydratedTreeSequenceField(parent.simpleContext, schema, key, parent);
 	}
 
-	return new UnhydratedFlexTreeField(parent.context, schema, key, parent);
+	return new UnhydratedFlexTreeField(parent.simpleContext, schema, key, parent);
 }
 
 // #endregion Caching and unboxing utilities

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

@@ -3,13 +3,21 @@
  * Licensed under the MIT License.
  */
 
-import { Context, type TreeNodeSchema, UnhydratedContext } from "./core/index.js";
+import { getOrCreate } from "../util/index.js";
+import { Context, UnhydratedContext } from "./core/index.js";
+import { normalizeFieldSchema, type ImplicitFieldSchema } from "./schemaTypes.js";
 import { toFlexSchema } from "./toFlexSchema.js";
 
+const contextCache: WeakMap<ImplicitFieldSchema, Context> = new WeakMap();
+
 /**
  * Utility for creating {@link Context}s for unhydrated nodes.
  */
-export function createUnhydratedContext(schema: TreeNodeSchema): Context {
-	const flexContext = new UnhydratedContext(toFlexSchema(schema));
-	return new Context([schema], flexContext);
+export function getUnhydratedContext(schema: ImplicitFieldSchema): Context {
+	return getOrCreate(contextCache, schema, (s) => {
+		const normalized = normalizeFieldSchema(schema);
+
+		const flexContext = new UnhydratedContext(toFlexSchema(normalized));
+		return new Context(normalized.allowedTypeSet, flexContext);
+	});
 }

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

@@ -42,7 +42,7 @@ import {
 import { brand, count, type RestrictiveStringRecord } from "../util/index.js";
 import { TreeNodeValid, type MostDerivedData } from "./treeNodeValid.js";
 import type { ExclusiveMapTree } from "../core/index.js";
-import { createUnhydratedContext } from "./createContext.js";
+import { getUnhydratedContext } from "./createContext.js";
 
 /**
  * A map of string keys to tree objects.
@@ -255,7 +255,7 @@ export function mapSchema<
 			input: T2,
 		): UnhydratedFlexTreeNode {
 			return UnhydratedFlexTreeNode.getOrCreate(
-				unhydratedContext.flexContext,
+				unhydratedContext,
 				mapTreeFromNodeData(input as FactoryContent, this as unknown as ImplicitAllowedTypes),
 			);
 		}
@@ -264,7 +264,7 @@ export function mapSchema<
 
 		protected static override oneTimeSetup<T2>(this: typeof TreeNodeValid<T2>): Context {
 			const schema = this as unknown as TreeNodeSchema;
-			unhydratedContext = createUnhydratedContext(schema);
+			unhydratedContext = getUnhydratedContext(schema);
 			return unhydratedContext;
 		}