tree: Reduce use of flex tree schema (microsoft#22626)

## Description

Reduce use of flex tree schema.
Work toward removal of flex tree schema.

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

@@ -159,7 +159,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.flexSchema.rootFieldSchema, cursor);
+		const field = makeField(this, this.schema.rootFieldSchema, cursor);
 		cursor.free();
 		return field;
 	}

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

@@ -20,7 +20,6 @@ import type {
 	OptionalFieldEditBuilder,
 } from "../default-schema/index.js";
 import type { FlexFieldKind } from "../modular-schema/index.js";
-import type { FlexTreeNodeSchema } from "../typed-schema/index.js";
 
 import type { FlexTreeContext } from "./context.js";
 
@@ -126,11 +125,7 @@ export enum TreeStatus {
  * the same {@link FlexTreeNode} instance will be used in the new location.
  * Similarly, edits applied to a node's sub-tree concurrently with the move of the node will still be applied to its subtree in its new location.
  *
- *
  * @remarks
- * Down-casting (via {@link FlexTreeNode#is}) is required to access Schema-Aware APIs, including editing.
- * All content in the tree is accessible without down-casting, but if the schema is known,
- * 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 {
@@ -162,11 +157,6 @@ export interface FlexTreeNode extends FlexTreeEntity {
 	 */
 	readonly parentField: { readonly parent: FlexTreeField; readonly index: number };
 
-	/**
-	 * Type guard for narrowing / down-casting to a specific schema.
-	 */
-	is(schema: FlexTreeNodeSchema): boolean;
-
 	boxedIterator(): IterableIterator<FlexTreeField>;
 
 	/**

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

@@ -43,15 +43,12 @@ 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, IDisposable
-{
+export abstract class LazyEntity<TAnchor = unknown> implements FlexTreeEntity, IDisposable {
 	readonly #lazyCursor: ITreeSubscriptionCursor;
 	public readonly [anchorSymbol]: TAnchor;
 
 	protected constructor(
 		public readonly context: Context,
-		public readonly flexSchema: TSchema,
 		cursor: ITreeSubscriptionCursor,
 		anchor: TAnchor,
 	) {

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

@@ -10,6 +10,7 @@ import {
 	type ExclusiveMapTree,
 	type FieldAnchor,
 	type FieldKey,
+	type FieldKindIdentifier,
 	type FieldUpPath,
 	type ITreeCursorSynchronous,
 	type ITreeSubscriptionCursor,
@@ -27,7 +28,6 @@ import {
 	type ValueFieldEditBuilder,
 } from "../default-schema/index.js";
 import type { FlexFieldKind } from "../modular-schema/index.js";
-import type { FlexFieldSchema } from "../typed-schema/index.js";
 
 import type { Context } from "./context.js";
 import {
@@ -75,7 +75,7 @@ const fieldCache: WeakMap<LazyTreeNode, Map<FieldKey, FlexTreeField>> = new Weak
 
 export function makeField(
 	context: Context,
-	schema: FlexFieldSchema,
+	schema: TreeFieldStoredSchema,
 	cursor: ITreeSubscriptionCursor,
 ): FlexTreeField {
 	const fieldAnchor = cursor.buildFieldAnchor();
@@ -125,19 +125,12 @@ export function makeField(
 /**
  * Base type for fields implementing {@link FlexTreeField} using cursors.
  */
-export abstract class LazyField<out TKind extends FlexFieldKind>
-	extends LazyEntity<FlexFieldSchema<TKind>, FieldAnchor>
-	implements FlexTreeField
-{
+export abstract class LazyField extends LazyEntity<FieldAnchor> implements FlexTreeField {
 	public get [flexTreeMarker](): FlexTreeEntityKind.Field {
 		return FlexTreeEntityKind.Field;
 	}
 	public readonly key: FieldKey;
 
-	public get schema(): TreeFieldStoredSchema {
-		return this.flexSchema.stored;
-	}
-
 	/**
 	 * If this field ends its lifetime before the Anchor does, this needs to be invoked to avoid a double free
 	 * if/when the Anchor is destroyed.
@@ -146,11 +139,11 @@ export abstract class LazyField<out TKind extends FlexFieldKind>
 
 	public constructor(
 		context: Context,
-		schema: FlexFieldSchema<TKind>,
+		public readonly schema: TreeFieldStoredSchema,
 		cursor: ITreeSubscriptionCursor,
 		fieldAnchor: FieldAnchor,
 	) {
-		super(context, schema, cursor, fieldAnchor);
+		super(context, cursor, fieldAnchor);
 		assert(cursor.mode === CursorLocationType.Fields, 0x77b /* must be in fields mode */);
 		this.key = cursor.getFieldKey();
 		// Fields currently live as long as their parent does.
@@ -171,7 +164,7 @@ export abstract class LazyField<out TKind extends FlexFieldKind>
 			0xa26 /* Narrowing must be done to a kind that exists in this context */,
 		);
 
-		return this.flexSchema.kind === (kind as unknown);
+		return this.schema.kind === kind.identifier;
 	}
 
 	public get parent(): FlexTreeNode | undefined {
@@ -259,19 +252,7 @@ export abstract class LazyField<out TKind extends FlexFieldKind>
 	}
 }
 
-export class LazySequence
-	extends LazyField<typeof FieldKinds.sequence>
-	implements FlexTreeSequenceField
-{
-	public constructor(
-		context: Context,
-		schema: FlexFieldSchema<typeof FieldKinds.sequence>,
-		cursor: ITreeSubscriptionCursor,
-		fieldAnchor: FieldAnchor,
-	) {
-		super(context, schema, cursor, fieldAnchor);
-	}
-
+export class LazySequence extends LazyField implements FlexTreeSequenceField {
 	public at(index: number): FlexTreeUnknownUnboxed | undefined {
 		const finalIndex = indexForAt(index, this.length);
 
@@ -302,19 +283,7 @@ export class LazySequence
 	}
 }
 
-export class ReadonlyLazyValueField
-	extends LazyField<typeof FieldKinds.required>
-	implements FlexTreeRequiredField
-{
-	public constructor(
-		context: Context,
-		schema: FlexFieldSchema<typeof FieldKinds.required>,
-		cursor: ITreeSubscriptionCursor,
-		fieldAnchor: FieldAnchor,
-	) {
-		super(context, schema, cursor, fieldAnchor);
-	}
-
+export class ReadonlyLazyValueField extends LazyField implements FlexTreeRequiredField {
 	public editor: ValueFieldEditBuilder<ExclusiveMapTree> = {
 		set: (newContent) => {
 			assert(false, 0xa0c /* Unexpected set of readonly field */);
@@ -327,15 +296,6 @@ export class ReadonlyLazyValueField
 }
 
 export class LazyValueField extends ReadonlyLazyValueField implements FlexTreeRequiredField {
-	public constructor(
-		context: Context,
-		schema: FlexFieldSchema<typeof FieldKinds.required>,
-		cursor: ITreeSubscriptionCursor,
-		fieldAnchor: FieldAnchor,
-	) {
-		super(context, schema, cursor, fieldAnchor);
-	}
-
 	public override editor: ValueFieldEditBuilder<ExclusiveMapTree> = {
 		set: (newContent) => {
 			this.valueFieldEditor().set(cursorForMapTreeNode(newContent));
@@ -353,19 +313,7 @@ export class LazyValueField extends ReadonlyLazyValueField implements FlexTreeRe
 	}
 }
 
-export class LazyOptionalField
-	extends LazyField<typeof FieldKinds.optional>
-	implements FlexTreeOptionalField
-{
-	public constructor(
-		context: Context,
-		schema: FlexFieldSchema<typeof FieldKinds.optional>,
-		cursor: ITreeSubscriptionCursor,
-		fieldAnchor: FieldAnchor,
-	) {
-		super(context, schema, cursor, fieldAnchor);
-	}
-
+export class LazyOptionalField extends LazyField implements FlexTreeOptionalField {
 	public editor: OptionalFieldEditBuilder<ExclusiveMapTree> = {
 		set: (newContent, wasEmpty) => {
 			this.optionalEditor().set(
@@ -386,25 +334,22 @@ export class LazyOptionalField
 	}
 }
 
-export class LazyForbiddenField extends LazyField<typeof FieldKinds.forbidden> {}
+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.
-	// Since this has to be done at runtime anyway, trying to use safer typing than `any` here (such as `never`, which is only slightly safer)
-	// does not seem worth it (ends up requiring type casts that are just as unsafe).
-	// eslint-disable-next-line @typescript-eslint/no-explicit-any
-	schema: FlexFieldSchema<any>,
+	schema: TreeFieldStoredSchema,
 	cursor: ITreeSubscriptionCursor,
 	fieldAnchor: FieldAnchor,
-) => LazyField<FlexFieldKind>;
-
-const builderList: [FlexFieldKind, Builder][] = [
-	[FieldKinds.forbidden, LazyForbiddenField],
-	[FieldKinds.optional, LazyOptionalField],
-	[FieldKinds.sequence, LazySequence],
-	[FieldKinds.required, LazyValueField],
-	[FieldKinds.identifier, LazyValueField],
+) => LazyField;
+
+const builderList: [FieldKindIdentifier, Builder][] = [
+	[FieldKinds.forbidden.identifier, LazyForbiddenField],
+	[FieldKinds.optional.identifier, LazyOptionalField],
+	[FieldKinds.sequence.identifier, LazySequence],
+	[FieldKinds.required.identifier, LazyValueField],
+	[FieldKinds.identifier.identifier, LazyValueField],
 ];
 
-const kindToClass: ReadonlyMap<FlexFieldKind, Builder> = new Map(builderList);
+const kindToClass: ReadonlyMap<FieldKindIdentifier, Builder> = new Map(builderList);

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

@@ -13,6 +13,7 @@ import {
 	type ITreeSubscriptionCursor,
 	type TreeNavigationResult,
 	type TreeNodeSchemaIdentifier,
+	type TreeNodeStoredSchema,
 	type Value,
 	inCursorField,
 	mapCursorFields,
@@ -52,10 +53,9 @@ export function makeTree(context: Context, cursor: ITreeSubscriptionCursor): Laz
 		context.checkout.forest.anchors.forget(anchor);
 		assert(cached.context === context, 0x782 /* contexts must match */);
 		assert(cached instanceof LazyTreeNode, 0x92c /* Expected LazyTreeNode */);
-		return cached as LazyTreeNode;
+		return cached;
 	}
-	const schema = context.flexSchema.nodeSchema.get(cursor.type) ?? fail("missing schema");
-	return new LazyTreeNode(context, schema, cursor, anchorNode, anchor);
+	return new LazyTreeNode(context, cursor.type, cursor, anchorNode, anchor);
 }
 
 function cleanupTree(anchor: AnchorNode): void {
@@ -67,47 +67,37 @@ function cleanupTree(anchor: AnchorNode): void {
 /**
  * Lazy implementation of {@link FlexTreeNode}.
  */
-export class LazyTreeNode<TSchema extends FlexTreeNodeSchema = FlexTreeNodeSchema>
-	extends LazyEntity<TSchema, Anchor>
-	implements FlexTreeNode
-{
+export class LazyTreeNode extends LazyEntity<Anchor> implements FlexTreeNode {
 	public get [flexTreeMarker](): FlexTreeEntityKind.Node {
 		return FlexTreeEntityKind.Node;
 	}
 
-	public get schema(): TreeNodeSchemaIdentifier {
-		return this.flexSchema.name;
-	}
-
 	// Using JS private here prevents it from showing up as a enumerable own property, or conflicting with struct fields.
 	readonly #removeDeleteCallback: () => void;
 
+	private readonly storedSchema: TreeNodeStoredSchema;
+	private readonly flexSchema: FlexTreeNodeSchema;
+
 	public constructor(
 		context: Context,
-		schema: TSchema,
+		public readonly schema: TreeNodeSchemaIdentifier,
 		cursor: ITreeSubscriptionCursor,
 		public readonly anchorNode: AnchorNode,
 		anchor: Anchor,
 	) {
-		super(context, schema, cursor, anchor);
+		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.flexSchema.name) !== undefined,
+			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 */,
 		);
 	}
 
-	public is(schema: FlexTreeNodeSchema): boolean {
-		assert(
-			this.context.flexSchema.nodeSchema.get(schema.name) === schema,
-			0x785 /* Narrowing must be done to a schema that exists in this context */,
-		);
-		return this.flexSchema === (schema as unknown);
-	}
-
 	protected override [tryMoveCursorToAnchorSymbol](
 		cursor: ITreeSubscriptionCursor,
 	): TreeNavigationResult {
@@ -133,20 +123,24 @@ export class LazyTreeNode<TSchema extends FlexTreeNodeSchema = FlexTreeNodeSchem
 			if (cursor.getFieldLength() === 0) {
 				return undefined;
 			}
-			return makeField(this.context, schema, cursor);
+			return makeField(this.context, schema.stored, cursor);
 		});
 	}
 
 	public getBoxed(key: FieldKey): FlexTreeField {
 		const fieldSchema = this.flexSchema.getFieldSchema(key);
 		return inCursorField(this[cursorSymbol], key, (cursor) => {
-			return makeField(this.context, fieldSchema, cursor);
+			return makeField(this.context, fieldSchema.stored, cursor);
 		});
 	}
 
 	public boxedIterator(): IterableIterator<FlexTreeField> {
 		return mapCursorFields(this[cursorSymbol], (cursor) =>
-			makeField(this.context, this.flexSchema.getFieldSchema(cursor.getFieldKey()), cursor),
+			makeField(
+				this.context,
+				this.flexSchema.getFieldSchema(cursor.getFieldKey()).stored,
+				cursor,
+			),
 		).values();
 	}
 
@@ -194,7 +188,7 @@ export class LazyTreeNode<TSchema extends FlexTreeNodeSchema = FlexTreeNodeSchem
 			fieldSchema = nodeSchema.getFieldSchema(key);
 		}
 
-		const proxifiedField = makeField(this.context, fieldSchema, cursor);
+		const proxifiedField = makeField(this.context, fieldSchema.stored, cursor);
 		cursor.enterNode(index);
 
 		return { parent: proxifiedField, index };

packages/dds/tree/src/test/feature-libraries/chunked-forest/chunkEncodingEndToEnd.spec.ts

@@ -158,7 +158,9 @@ describe("End to end chunked encoding", () => {
 			{ jsonValidator: typeboxValidator },
 		);
 		const dummyEditor = new DefaultEditBuilder(new DefaultChangeFamily(codec), changeReceiver);
-		const checkout = new MockTreeCheckout(forest, dummyEditor as unknown as ISharedTreeEditor);
+		const checkout = new MockTreeCheckout(forest, {
+			editor: dummyEditor as unknown as ISharedTreeEditor,
+		});
 		checkout.editor
 			.sequenceField({ field: rootFieldKey, parent: undefined })
 			.insert(0, chunk.cursor());