tree: Remove more use of flex tree schema (microsoft#22632)

## Description

Remove more use of flex tree schema.

packages/dds/tree/src/core/schema-stored/schema.ts

@@ -166,6 +166,11 @@ export abstract class TreeNodeStoredSchema {
 	 * and is runtime validated by the codec.
 	 */
 	public abstract encode(): ErasedTreeNodeSchemaDataFormat;
+
+	/**
+	 * Returns the schema for the provided field.
+	 */
+	public abstract getFieldSchema(field: FieldKey): TreeFieldStoredSchema;
 }
 
 /**
@@ -202,6 +207,10 @@ export class ObjectNodeStoredSchema extends TreeNodeStoredSchema {
 			object: fieldsObject,
 		});
 	}
+
+	public override getFieldSchema(field: FieldKey): TreeFieldStoredSchema {
+		return this.objectNodeFields.get(field) ?? storedEmptyFieldSchema;
+	}
 }
 
 /**
@@ -225,6 +234,10 @@ export class MapNodeStoredSchema extends TreeNodeStoredSchema {
 			map: encodeFieldSchema(this.mapFields),
 		});
 	}
+
+	public override getFieldSchema(field: FieldKey): TreeFieldStoredSchema {
+		return this.mapFields;
+	}
 }
 
 /**
@@ -252,6 +265,10 @@ export class LeafNodeStoredSchema extends TreeNodeStoredSchema {
 			leaf: encodeValueSchema(this.leafValue),
 		});
 	}
+
+	public override getFieldSchema(field: FieldKey): TreeFieldStoredSchema {
+		return storedEmptyFieldSchema;
+	}
 }
 
 export const storedSchemaDecodeDispatcher: DiscriminatedUnionDispatcher<

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

@@ -8,6 +8,7 @@ import { assert } from "@fluidframework/core-utils/internal";
 import {
 	type FieldKey,
 	type ForestEvents,
+	type SchemaPolicy,
 	type TreeFieldStoredSchema,
 	type TreeStoredSchema,
 	anchorSlot,
@@ -40,6 +41,11 @@ export interface FlexTreeContext {
 	 */
 	readonly schema: TreeStoredSchema;
 
+	/**
+	 * SchemaPolicy used within this context.
+	 */
+	readonly schemaPolicy: SchemaPolicy;
+
 	/**
 	 * If true, this context is the canonical context instance for a given view,
 	 * and its schema include all schema from the document.
@@ -108,6 +114,10 @@ export class Context implements FlexTreeHydratedContext, IDisposable {
 		this.checkout.forest.anchors.slots.set(ContextSlot, this);
 	}
 
+	public get schemaPolicy(): SchemaPolicy {
+		return this.flexSchema.policy;
+	}
+
 	public isHydrated(): this is FlexTreeHydratedContext {
 		return true;
 	}
@@ -159,7 +169,7 @@ export class Context implements FlexTreeHydratedContext, IDisposable {
 		assert(this.disposed === false, 0x804 /* use after dispose */);
 		const cursor = this.checkout.forest.allocateCursor("root");
 		moveToDetachedField(this.checkout.forest, cursor);
-		const field = makeField(this, this.schema.rootFieldSchema, cursor);
+		const field = makeField(this, this.schema.rootFieldSchema.kind, cursor);
 		cursor.free();
 		return field;
 	}

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

@@ -7,8 +7,8 @@ import {
 	type AnchorNode,
 	type ExclusiveMapTree,
 	type FieldKey,
+	type FieldKindIdentifier,
 	type FieldUpPath,
-	type TreeFieldStoredSchema,
 	type TreeNodeSchemaIdentifier,
 	type TreeValue,
 	anchorSlot,
@@ -252,7 +252,7 @@ export interface FlexTreeField extends FlexTreeEntity {
 	 * Schema for this entity.
 	 * If well-formed, it must follow this schema.
 	 */
-	readonly schema: TreeFieldStoredSchema;
+	readonly schema: FieldKindIdentifier;
 }
 
 // #region Field Kinds

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

@@ -14,7 +14,6 @@ import {
 	type FieldUpPath,
 	type ITreeCursorSynchronous,
 	type ITreeSubscriptionCursor,
-	type TreeFieldStoredSchema,
 	type TreeNavigationResult,
 	inCursorNode,
 	iterateCursorField,
@@ -75,15 +74,15 @@ const fieldCache: WeakMap<LazyTreeNode, Map<FieldKey, FlexTreeField>> = new Weak
 
 export function makeField(
 	context: Context,
-	schema: TreeFieldStoredSchema,
+	schema: FieldKindIdentifier,
 	cursor: ITreeSubscriptionCursor,
 ): FlexTreeField {
 	const fieldAnchor = cursor.buildFieldAnchor();
 	let usedAnchor = false;
 
 	const makeFlexTreeField = (): FlexTreeField => {
 		usedAnchor = true;
-		const field = new (kindToClass.get(schema.kind) ?? fail("missing field implementation"))(
+		const field = new (kindToClass.get(schema) ?? fail("missing field implementation"))(
 			context,
 			schema,
 			cursor,
@@ -139,7 +138,7 @@ export abstract class LazyField extends LazyEntity<FieldAnchor> implements FlexT
 
 	public constructor(
 		context: Context,
-		public readonly schema: TreeFieldStoredSchema,
+		public readonly schema: FieldKindIdentifier,
 		cursor: ITreeSubscriptionCursor,
 		fieldAnchor: FieldAnchor,
 	) {
@@ -160,11 +159,11 @@ export abstract class LazyField extends LazyEntity<FieldAnchor> implements FlexT
 
 	public is<TKind2 extends FlexFieldKind>(kind: TKind2): this is FlexTreeTypedField<TKind2> {
 		assert(
-			this.context.flexSchema.policy.fieldKinds.get(kind.identifier) === kind,
+			this.context.schemaPolicy.fieldKinds.get(kind.identifier) === kind,
 			0xa26 /* Narrowing must be done to a kind that exists in this context */,
 		);
 
-		return this.schema.kind === kind.identifier;
+		return this.schema === kind.identifier;
 	}
 
 	public get parent(): FlexTreeNode | undefined {
@@ -339,7 +338,7 @@ export class LazyForbiddenField extends LazyField {}
 type Builder = new (
 	context: Context,
 	// Correct use of these builders requires the builder of the matching type to be used.
-	schema: TreeFieldStoredSchema,
+	schema: FieldKindIdentifier,
 	cursor: ITreeSubscriptionCursor,
 	fieldAnchor: FieldAnchor,
 ) => LazyField;

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

@@ -10,6 +10,7 @@ import {
 	type AnchorNode,
 	CursorLocationType,
 	type FieldKey,
+	type FieldKindIdentifier,
 	type ITreeSubscriptionCursor,
 	type TreeNavigationResult,
 	type TreeNodeSchemaIdentifier,
@@ -21,7 +22,6 @@ import {
 } from "../../core/index.js";
 import { disposeSymbol, fail } from "../../util/index.js";
 import { FieldKinds } from "../default-schema/index.js";
-import { FlexFieldSchema, type FlexTreeNodeSchema } from "../typed-schema/index.js";
 
 import type { Context } from "./context.js";
 import {
@@ -76,7 +76,6 @@ export class LazyTreeNode extends LazyEntity<Anchor> implements FlexTreeNode {
 	readonly #removeDeleteCallback: () => void;
 
 	private readonly storedSchema: TreeNodeStoredSchema;
-	private readonly flexSchema: FlexTreeNodeSchema;
 
 	public constructor(
 		context: Context,
@@ -87,15 +86,9 @@ export class LazyTreeNode extends LazyEntity<Anchor> implements FlexTreeNode {
 	) {
 		super(context, cursor, anchor);
 		this.storedSchema = context.schema.nodeSchema.get(this.schema) ?? fail("missing schema");
-		this.flexSchema = context.flexSchema.nodeSchema.get(this.schema) ?? fail("missing schema");
 		assert(cursor.mode === CursorLocationType.Nodes, 0x783 /* must be in nodes mode */);
 		anchorNode.slots.set(flexTreeSlot, this);
 		this.#removeDeleteCallback = anchorNode.on("afterDestroy", cleanupTree);
-
-		assert(
-			this.context.flexSchema.nodeSchema.get(this.schema) !== undefined,
-			0x784 /* There is no explicit schema for this node type. Ensure that the type is correct and the schema for it was added to the TreeStoredSchema */,
-		);
 	}
 
 	protected override [tryMoveCursorToAnchorSymbol](
@@ -118,27 +111,27 @@ export class LazyTreeNode extends LazyEntity<Anchor> implements FlexTreeNode {
 	}
 
 	public tryGetField(fieldKey: FieldKey): FlexTreeField | undefined {
-		const schema = this.flexSchema.getFieldSchema(fieldKey);
+		const schema = this.storedSchema.getFieldSchema(fieldKey);
 		return inCursorField(this[cursorSymbol], fieldKey, (cursor) => {
 			if (cursor.getFieldLength() === 0) {
 				return undefined;
 			}
-			return makeField(this.context, schema.stored, cursor);
+			return makeField(this.context, schema.kind, cursor);
 		});
 	}
 
 	public getBoxed(key: FieldKey): FlexTreeField {
-		const fieldSchema = this.flexSchema.getFieldSchema(key);
+		const fieldSchema = this.storedSchema.getFieldSchema(key);
 		return inCursorField(this[cursorSymbol], key, (cursor) => {
-			return makeField(this.context, fieldSchema.stored, cursor);
+			return makeField(this.context, fieldSchema.kind, cursor);
 		});
 	}
 
 	public boxedIterator(): IterableIterator<FlexTreeField> {
 		return mapCursorFields(this[cursorSymbol], (cursor) =>
 			makeField(
 				this.context,
-				this.flexSchema.getFieldSchema(cursor.getFieldKey()).stored,
+				this.storedSchema.getFieldSchema(cursor.getFieldKey()).kind,
 				cursor,
 			),
 		).values();
@@ -152,14 +145,14 @@ export class LazyTreeNode extends LazyEntity<Anchor> implements FlexTreeNode {
 
 		cursor.exitNode();
 		assert(key === cursor.getFieldKey(), 0x787 /* mismatched keys */);
-		let fieldSchema: FlexFieldSchema;
+		let fieldSchema: FieldKindIdentifier;
 
 		// Check if the current node is in a detached sequence.
 		if (this.anchorNode.parent === undefined) {
 			// Parent field is a detached sequence, and thus needs special handling for its schema.
 			// eslint-disable-next-line unicorn/prefer-ternary
 			if (key === rootFieldKey) {
-				fieldSchema = this.context.flexSchema.rootFieldSchema;
+				fieldSchema = this.context.schema.rootFieldSchema.kind;
 			} else {
 				// All fields (in the flex tree API) have a schema.
 				// Since currently there is no known schema for detached field other than the special default root:
@@ -174,21 +167,19 @@ export class LazyTreeNode extends LazyEntity<Anchor> implements FlexTreeNode {
 				// Additionally this approach makes it possible for a user to take a FlexTree node, get its parent, check its schema, down cast based on that, then edit that detached field (ex: removing the node in it).
 				// This MIGHT work properly with existing merge resolution logic (it must keep client in sync and be unable to violate schema), but this either needs robust testing or to be explicitly banned (error before sending the op).
 				// Issues like replacing a node in the a removed sequenced then undoing the remove could easily violate schema if not everything works exactly right!
-				fieldSchema = FlexFieldSchema.create(FieldKinds.sequence, [
-					...this.context.flexSchema.nodeSchema.values(),
-				]);
+				fieldSchema = FieldKinds.sequence.identifier;
 			}
 		} else {
 			cursor.exitField();
 			const parentType = cursor.type;
 			cursor.enterField(key);
 			const nodeSchema =
-				this.context.flexSchema.nodeSchema.get(parentType) ??
+				this.context.schema.nodeSchema.get(parentType) ??
 				fail("requested schema that does not exist");
-			fieldSchema = nodeSchema.getFieldSchema(key);
+			fieldSchema = nodeSchema.getFieldSchema(key).kind;
 		}
 
-		const proxifiedField = makeField(this.context, fieldSchema.stored, cursor);
+		const proxifiedField = makeField(this.context, fieldSchema, cursor);
 		cursor.enterNode(index);
 
 		return { parent: proxifiedField, index };

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

@@ -83,7 +83,7 @@ export interface DetachedFieldCache {
 export function getSchemaAndPolicy(nodeOrField: FlexTreeEntity): SchemaAndPolicy {
 	return {
 		schema: nodeOrField.context.schema,
-		policy: nodeOrField.context.flexSchema.policy,
+		policy: nodeOrField.context.schemaPolicy,
 	};
 }
 

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

@@ -152,7 +152,7 @@ export function createFromCursor<TSchema extends ImplicitFieldSchema>(
 ): Unhydrated<TreeFieldFromImplicitField<TSchema>> {
 	const mapTrees = cursor === undefined ? [] : [mapTreeFromCursor(cursor)];
 	const context = getUnhydratedContext(schema);
-	const flexSchema = context.flexContext.flexSchema;
+	const flexSchema = context.flexContext.schema;
 
 	const schemaValidationPolicy: SchemaAndPolicy = {
 		policy: defaultSchemaPolicy,
@@ -161,7 +161,7 @@ export function createFromCursor<TSchema extends ImplicitFieldSchema>(
 
 	const maybeError = isFieldInSchema(
 		mapTrees,
-		flexSchema.rootFieldSchema.stored,
+		flexSchema.rootFieldSchema,
 		schemaValidationPolicy,
 	);
 	inSchemaOrThrow(maybeError);

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

@@ -305,7 +305,7 @@ function getStoredKey(node: TreeNode): string | number {
 	// Note: the flex domain strictly works with "stored keys", and knows nothing about the developer-facing
 	// "property keys".
 	const parentField = getOrCreateInnerNode(node).parentField;
-	if (parentField.parent.schema.kind === FieldKinds.sequence.identifier) {
+	if (parentField.parent.schema === FieldKinds.sequence.identifier) {
 		// The parent of `node` is an array node
 		assert(
 			parentField.parent.key === EmptyKey,

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

@@ -33,6 +33,7 @@ import {
 	type Context,
 	getOrCreateNodeFromInnerNode,
 	type TreeNodeSchemaBoth,
+	getSimpleNodeSchemaFromInnerNode,
 } from "./core/index.js";
 import { type InsertableContent, mapTreeFromNodeData } from "./toMapTree.js";
 import { fail } from "../util/index.js";
@@ -663,6 +664,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 	public static readonly kind = NodeKind.Array;
 
 	protected abstract get simpleSchema(): T;
+	protected abstract get allowedTypes(): ReadonlySet<TreeNodeSchema>;
 
 	public constructor(
 		input: Iterable<InsertableTreeNodeFromImplicitAllowedTypes<T>> | InternalTreeNode,
@@ -819,6 +821,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		source?: TreeArrayNode,
 	): void {
 		const destinationField = getSequenceField(this);
+		const destinationSchema = this.allowedTypes;
 		const sourceField = source !== undefined ? getSequenceField(source) : destinationField;
 
 		validateIndex(destinationIndex, destinationField, "moveRangeToIndex", true);
@@ -828,7 +831,8 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		if (sourceField !== destinationField) {
 			for (let i = sourceStart; i < sourceEnd; i++) {
 				const sourceNode = sourceField.boxedAt(i) ?? oob();
-				if (!destinationField.schema.types.has(sourceNode.schema)) {
+				const sourceSchema = getSimpleNodeSchemaFromInnerNode(sourceNode);
+				if (!destinationSchema.has(sourceSchema)) {
 					throw new UsageError("Type in source sequence is not allowed in destination.");
 				}
 			}
@@ -1014,6 +1018,9 @@ export function arraySchema<
 		protected get simpleSchema(): T {
 			return info;
 		}
+		protected get allowedTypes(): ReadonlySet<TreeNodeSchema> {
+			return lazyChildTypes.value;
+		}
 	}
 
 	const output: Output = Schema;