Tree: Remove use of toFlexSchema (microsoft#22638)

## Description

Remove use of toFlexSchema.

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

@@ -18,7 +18,6 @@ import type { Listenable } from "../../events/index.js";
 import { type IDisposable, disposeSymbol } from "../../util/index.js";
 import type { FieldGenerator } from "../fieldGenerator.js";
 import type { NodeKeyManager } from "../node-key/index.js";
-import type { FlexTreeSchema } from "../typed-schema/index.js";
 
 import type { FlexTreeField } from "./flexTreeTypes.js";
 import { type LazyEntity, prepareForEditSymbol } from "./lazyEntity.js";
@@ -29,12 +28,6 @@ import type { ITreeCheckout } from "../../shared-tree/index.js";
  * Context for FlexTrees.
  */
 export interface FlexTreeContext {
-	/**
-	 * Schema used within this context.
-	 * All data must conform to these schema.
-	 */
-	readonly flexSchema: FlexTreeSchema;
-
 	/**
 	 * Schema used within this context.
 	 * All data must conform to these schema.
@@ -97,7 +90,7 @@ export class Context implements FlexTreeHydratedContext, IDisposable {
 	 * @param nodeKeyManager - An object which handles node key generation and conversion
 	 */
 	public constructor(
-		public readonly flexSchema: FlexTreeSchema,
+		public readonly schemaPolicy: SchemaPolicy,
 		public readonly checkout: ITreeCheckout,
 		public readonly nodeKeyManager: NodeKeyManager,
 	) {
@@ -114,10 +107,6 @@ 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;
 	}
@@ -201,7 +190,7 @@ export class Context implements FlexTreeHydratedContext, IDisposable {
  * This is necessary for supporting using this tree across edits to the forest, and not leaking memory.
  */
 export function getTreeContext(
-	schema: FlexTreeSchema,
+	schema: SchemaPolicy,
 	checkout: ITreeCheckout,
 	nodeKeyManager: NodeKeyManager,
 ): Context {

packages/dds/tree/src/shared-tree/checkoutFlexTreeView.ts

@@ -7,10 +7,10 @@ import { assert } from "@fluidframework/core-utils/internal";
 import {
 	type Context,
 	type FlexTreeField,
-	type FlexTreeSchema,
 	type NodeKeyManager,
 	getTreeContext,
 	type FlexTreeHydratedContext,
+	type FullSchemaPolicy,
 } from "../feature-libraries/index.js";
 import { tryDisposeTreeNode } from "../simple-tree/index.js";
 import { disposeSymbol } from "../util/index.js";
@@ -42,7 +42,7 @@ export class CheckoutFlexTreeView<out TCheckout extends ITreeCheckout = ITreeChe
 		 * This is a non-owning reference: disposing of this view does not impact the branch.
 		 */
 		public readonly checkout: TCheckout,
-		public readonly schema: FlexTreeSchema,
+		public readonly schema: FullSchemaPolicy,
 		public readonly nodeKeyManager: NodeKeyManager,
 		private readonly onDispose?: () => void,
 	) {

packages/dds/tree/src/shared-tree/schematizingTreeView.ts

@@ -6,7 +6,12 @@
 import { assert } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
-import { AllowedUpdateType, anchorSlot, Compatibility } from "../core/index.js";
+import {
+	AllowedUpdateType,
+	anchorSlot,
+	Compatibility,
+	type SchemaPolicy,
+} from "../core/index.js";
 import {
 	type HasListeners,
 	type IEmitter,
@@ -19,7 +24,7 @@ import {
 	defaultSchemaPolicy,
 	ContextSlot,
 	cursorForMapTreeNode,
-	type FlexTreeSchema,
+	type FullSchemaPolicy,
 } from "../feature-libraries/index.js";
 import {
 	type FieldSchema,
@@ -30,7 +35,6 @@ import {
 	type TreeView,
 	type TreeViewEvents,
 	getTreeNodeForField,
-	toFlexSchema,
 	setField,
 	normalizeFieldSchema,
 	type InsertableContent,
@@ -69,7 +73,7 @@ export class SchematizingSimpleTreeView<in out TRootSchema extends ImplicitField
 	 * Undefined iff uninitialized or disposed.
 	 */
 	private currentCompatibility: SchemaCompatibilityStatus | undefined;
-	private readonly flexSchema: FlexTreeSchema;
+	private readonly schemaPolicy: SchemaPolicy;
 	public readonly events: Listenable<TreeViewEvents> &
 		IEmitter<TreeViewEvents> &
 		HasListeners<TreeViewEvents> = createEmitter();
@@ -106,7 +110,7 @@ export class SchematizingSimpleTreeView<in out TRootSchema extends ImplicitField
 			validateSchema: config.enableSchemaValidation,
 		};
 		this.rootFieldSchema = normalizeFieldSchema(config.schema);
-		this.flexSchema = toFlexSchema(config.schema);
+		this.schemaPolicy = defaultSchemaPolicy;
 
 		this.viewSchema = new ViewSchema(policy, {}, toStoredSchema(this.rootFieldSchema));
 		// This must be initialized before `update` can be called.
@@ -273,7 +277,7 @@ export class SchematizingSimpleTreeView<in out TRootSchema extends ImplicitField
 				this.viewSchema,
 				onViewDispose,
 				this.nodeKeyManager,
-				this.flexSchema,
+				this.schemaPolicy,
 			);
 			this.view = view;
 			assert(
@@ -377,7 +381,7 @@ export function requireSchema(
 	viewSchema: ViewSchema,
 	onDispose: () => void,
 	nodeKeyManager: NodeKeyManager,
-	flexTreeSchema: FlexTreeSchema,
+	schemaPolicy: FullSchemaPolicy,
 ): CheckoutFlexTreeView {
 	const slots = checkout.forest.anchors.slots;
 	assert(!slots.has(ContextSlot), 0x8c2 /* Cannot create second view from checkout */);
@@ -391,7 +395,7 @@ export function requireSchema(
 		);
 	}
 
-	const view = new CheckoutFlexTreeView(checkout, flexTreeSchema, nodeKeyManager, onDispose);
+	const view = new CheckoutFlexTreeView(checkout, schemaPolicy, nodeKeyManager, onDispose);
 	assert(slots.has(ContextSlot), 0x90d /* Context should be tracked in slot */);
 
 	return view;

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

@@ -24,12 +24,11 @@ import {
 	cursorForMapTreeNode,
 	defaultSchemaPolicy,
 	FieldKinds,
-	intoStoredSchema,
 	mapTreeFromCursor,
 	type NodeKeyManager,
 } from "../../feature-libraries/index.js";
 import { isFieldInSchema } from "../../feature-libraries/index.js";
-import { toFlexSchema } from "../toFlexSchema.js";
+import { toStoredSchema } from "../toFlexSchema.js";
 import { inSchemaOrThrow, mapTreeFromNodeData, type InsertableContent } from "../toMapTree.js";
 import {
 	applySchemaToParserOptions,
@@ -80,11 +79,11 @@ export function cursorFromInsertable<TSchema extends ImplicitFieldSchema>(
 ):
 	| ITreeCursorSynchronous
 	| (TSchema extends FieldSchema<FieldKind.Optional> ? undefined : never) {
-	const flexSchema = toFlexSchema(schema);
+	const storedSchema = toStoredSchema(schema);
 	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: storedSchema,
 	};
 
 	const mapTree = mapTreeFromNodeData(
@@ -95,7 +94,7 @@ export function cursorFromInsertable<TSchema extends ImplicitFieldSchema>(
 	);
 	if (mapTree === undefined) {
 		assert(
-			flexSchema.rootFieldSchema.kind === FieldKinds.optional,
+			storedSchema.rootFieldSchema.kind === FieldKinds.optional.identifier,
 			0xa10 /* missing non-optional field */,
 		);
 		return undefined as TSchema extends FieldSchema<FieldKind.Optional> ? undefined : never;

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

@@ -16,7 +16,7 @@ import {
 	FieldKind,
 } from "../schemaTypes.js";
 import { NodeKind, type TreeNodeSchema } from "../core/index.js";
-import { toFlexSchema } from "../toFlexSchema.js";
+import { toStoredSchema } from "../toFlexSchema.js";
 import { LeafNodeSchema } from "../leafNodeSchema.js";
 import { assert } from "@fluidframework/core-utils/internal";
 import { isObjectNodeSchema, type ObjectNodeSchema } from "../objectNodeTypes.js";
@@ -242,7 +242,7 @@ export class TreeViewConfiguration<TSchema extends ImplicitFieldSchema = Implici
 		}
 
 		// Eagerly perform this conversion to surface errors sooner.
-		toFlexSchema(config.schema);
+		toStoredSchema(config.schema);
 	}
 }
 

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

@@ -25,7 +25,7 @@ import type {
 	ImplicitFieldSchema,
 	ImplicitAllowedTypes,
 } from "../schemaTypes.js";
-import { getSimpleNodeSchema, NodeKind, type TreeNodeSchema } from "../core/index.js";
+import { NodeKind, type TreeNodeSchema } from "../core/index.js";
 import {
 	isTreeValue,
 	stackTreeFieldCursor,
@@ -39,9 +39,9 @@ import {
 	numberSchema,
 	stringSchema,
 } from "../leafNodeSchema.js";
-import { toFlexSchema } from "../toFlexSchema.js";
 import { isObjectNodeSchema } from "../objectNodeTypes.js";
 import { walkFieldSchema } from "../walkFieldSchema.js";
+import { getUnhydratedContext } from "../createContext.js";
 
 /**
  * Verbose encoding of a {@link TreeNode} or {@link TreeValue}.
@@ -174,9 +174,7 @@ export function applySchemaToParserOptions<TCustom>(
 		...options,
 	};
 
-	// TODO: should provide a way to look up schema by name efficiently without converting to flex tree schema and back.
-	// Maybe cache identifier->schema map on simple tree schema lazily.
-	const flexSchema = toFlexSchema(schema);
+	const context = getUnhydratedContext(schema);
 
 	return {
 		valueConverter: config.valueConverter,
@@ -185,9 +183,7 @@ export function applySchemaToParserOptions<TCustom>(
 			: {
 					encode: (type, key: FieldKey): string => {
 						// translate stored key into property key.
-						const flexNodeSchema =
-							flexSchema.nodeSchema.get(brand(type)) ?? fail("missing schema");
-						const simpleNodeSchema = getSimpleNodeSchema(flexNodeSchema);
+						const simpleNodeSchema = context.schema.get(brand(type)) ?? fail("missing schema");
 						if (isObjectNodeSchema(simpleNodeSchema)) {
 							const propertyKey = simpleNodeSchema.storedKeyToPropertyKey.get(key);
 							if (propertyKey !== undefined) {
@@ -208,9 +204,7 @@ export function applySchemaToParserOptions<TCustom>(
 						return key;
 					},
 					parse: (type, inputKey): FieldKey => {
-						const flexNodeSchema =
-							flexSchema.nodeSchema.get(brand(type)) ?? fail("missing schema");
-						const simpleNodeSchema = getSimpleNodeSchema(flexNodeSchema);
+						const simpleNodeSchema = context.schema.get(brand(type)) ?? fail("missing schema");
 						if (isObjectNodeSchema(simpleNodeSchema)) {
 							const info =
 								simpleNodeSchema.flexKeyMap.get(inputKey) ?? fail("missing field info");

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

@@ -35,8 +35,6 @@ import {
 	indexForAt,
 	type FlexTreeHydratedContext,
 	FlexFieldSchema,
-	type FlexTreeSchema,
-	intoStoredSchemaCollection,
 	type FlexFieldKind,
 	FieldKinds,
 	type SequenceFieldEditBuilder,
@@ -92,7 +90,7 @@ export class UnhydratedFlexTreeNode implements UnhydratedFlexTreeNode {
 		return nodeCache.get(mapTree) ?? new UnhydratedFlexTreeNode(context, mapTree, undefined);
 	}
 
-	public get context(): UnhydratedContext {
+	public get context(): FlexTreeContext {
 		return this.simpleContext.flexContext;
 	}
 
@@ -222,25 +220,17 @@ export class UnhydratedFlexTreeNode implements UnhydratedFlexTreeNode {
  * @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 constructor(
+		public readonly schemaPolicy: SchemaPolicy,
+		public readonly schema: TreeStoredSchema,
+	) {}
 
 	public isHydrated(): this is FlexTreeHydratedContext {
 		return false;
 	}
-
-	public get schemaPolicy(): SchemaPolicy {
-		return this.flexSchema.policy;
-	}
 }
 
 // #region Fields
@@ -282,7 +272,7 @@ const unparentedLocation: LocationInField = {
 class UnhydratedFlexTreeField implements FlexTreeField {
 	public [flexTreeMarker] = FlexTreeEntityKind.Field as const;
 
-	public get context(): UnhydratedContext {
+	public get context(): FlexTreeContext {
 		return this.simpleContext.flexContext;
 	}
 

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

@@ -3,10 +3,11 @@
  * Licensed under the MIT License.
  */
 
+import { defaultSchemaPolicy } from "../feature-libraries/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";
+import { toStoredSchema } from "./toFlexSchema.js";
 
 const contextCache: WeakMap<ImplicitFieldSchema, Context> = new WeakMap();
 
@@ -17,7 +18,7 @@ export function getUnhydratedContext(schema: ImplicitFieldSchema): Context {
 	return getOrCreate(contextCache, schema, (s) => {
 		const normalized = normalizeFieldSchema(schema);
 
-		const flexContext = new UnhydratedContext(toFlexSchema(normalized));
+		const flexContext = new UnhydratedContext(defaultSchemaPolicy, toStoredSchema(schema));
 		return new Context(normalized.allowedTypeSet, flexContext);
 	});
 }

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

@@ -90,7 +90,6 @@ export {
 	type FieldSchemaMetadata,
 } from "./schemaTypes.js";
 export { getOrCreateInnerNode } from "./proxyBinding.js";
-export { toFlexSchema } from "./toFlexSchema.js";
 export type {
 	FieldHasDefaultUnsafe,
 	ObjectFromSchemaRecordUnsafe,

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

@@ -57,7 +57,7 @@ type SchemaMap = Map<TreeNodeSchemaIdentifier, SchemaInfo>;
  *
  * This also has the side effect of populating the cached view schema on the class-based schema.
  */
-export function toFlexSchema(root: ImplicitFieldSchema): FlexTreeSchema {
+function toFlexSchema(root: ImplicitFieldSchema): FlexTreeSchema {
 	const schemaMap: SchemaMap = new Map();
 	const field = convertField(schemaMap, root);
 	const nodeSchema = new Map(

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

@@ -40,7 +40,6 @@ import {
 	TreeCompressionStrategy,
 	cursorForJsonableTreeField,
 	defaultSchemaPolicy,
-	intoStoredSchema,
 } from "../../../../feature-libraries/index.js";
 import { type JsonCompatibleReadOnly, brand } from "../../../../util/index.js";
 import { ajvValidator } from "../../../codec/index.js";
@@ -336,7 +335,7 @@ describe("schemaBasedEncoding", () => {
 				const idCompressor = createIdCompressor(
 					assertIsSessionId("00000000-0000-4000-b000-000000000000"),
 				);
-				const storedSchema = intoStoredSchema(schemaData);
+				const storedSchema = schemaData;
 				const tree = treeFactory(idCompressor);
 				// Check with checkFieldEncode
 				const cache = buildCache(storedSchema, defaultSchemaPolicy, idCompressor);

packages/dds/tree/src/test/feature-libraries/default-schema/schemaChecker.spec.ts

@@ -15,8 +15,8 @@ import {
 } from "../../../feature-libraries/default-schema/schemaChecker.js";
 import {
 	cursorForJsonableTreeNode,
+	defaultSchemaPolicy,
 	FieldKinds,
-	intoStoredSchema,
 	mapTreeFromCursor,
 } from "../../../feature-libraries/index.js";
 import {
@@ -504,11 +504,11 @@ describe("schema validation", () => {
 				const mapTrees = testTree
 					.treeFactory(testIdCompressor)
 					.map((j) => mapTreeFromCursor(cursorForJsonableTreeNode(j)));
-				const schema = intoStoredSchema(testTree.schemaData);
+				const schema = testTree.schemaData;
 				assert.equal(
 					isFieldInSchema(mapTrees, schema.rootFieldSchema, {
 						schema,
-						policy: testTree.schemaData.policy,
+						policy: defaultSchemaPolicy,
 					}),
 					SchemaValidationErrors.NoError,
 				);