Shared Tree: Add allowed types object with staged schema upgrade information to SimpleSchema

packages/dds/tree/api-report/tree.alpha.api.md

@@ -274,6 +274,8 @@ export class FieldSchemaAlpha<Kind extends FieldKind = FieldKind, Types extends
     get allowedTypesIdentifiers(): ReadonlySet<string>;
     // (undocumented)
     get persistedMetadata(): JsonCompatibleReadOnlyObject | undefined;
+    // (undocumented)
+    get simpleAllowedTypes(): ReadonlyMap<string, SimpleAllowedTypeAttributes>;
 }
 
 // @alpha @sealed @system
@@ -1004,17 +1006,22 @@ export type SharedTreeOptions = Partial<CodecWriteOptions> & Partial<SharedTreeF
 // @beta @input
 export type SharedTreeOptionsBeta = ForestOptions;
 
+// @alpha @sealed
+export interface SimpleAllowedTypeAttributes {
+    readonly isStaged: boolean | undefined;
+}
+
 // @alpha @sealed
 export interface SimpleArrayNodeSchema<out TCustomMetadata = unknown> extends SimpleNodeSchemaBaseAlpha<NodeKind.Array, TCustomMetadata> {
-    readonly allowedTypesIdentifiers: ReadonlySet<string>;
+    readonly simpleAllowedTypes: ReadonlyMap<string, SimpleAllowedTypeAttributes>;
 }
 
 // @alpha @sealed
 export interface SimpleFieldSchema {
-    readonly allowedTypesIdentifiers: ReadonlySet<string>;
     readonly kind: FieldKind;
     readonly metadata: FieldSchemaMetadata;
     readonly persistedMetadata?: JsonCompatibleReadOnlyObject | undefined;
+    readonly simpleAllowedTypes: ReadonlyMap<string, SimpleAllowedTypeAttributes>;
 }
 
 // @alpha @sealed
@@ -1024,7 +1031,7 @@ export interface SimpleLeafNodeSchema extends SimpleNodeSchemaBaseAlpha<NodeKind
 
 // @alpha @sealed
 export interface SimpleMapNodeSchema<out TCustomMetadata = unknown> extends SimpleNodeSchemaBaseAlpha<NodeKind.Map, TCustomMetadata> {
-    readonly allowedTypesIdentifiers: ReadonlySet<string>;
+    readonly simpleAllowedTypes: ReadonlyMap<string, SimpleAllowedTypeAttributes>;
 }
 
 // @alpha
@@ -1053,7 +1060,7 @@ export interface SimpleObjectNodeSchema<out TCustomMetadata = unknown> extends S
 
 // @alpha @sealed
 export interface SimpleRecordNodeSchema<out TCustomMetadata = unknown> extends SimpleNodeSchemaBaseAlpha<NodeKind.Record, TCustomMetadata> {
-    readonly allowedTypesIdentifiers: ReadonlySet<string>;
+    readonly simpleAllowedTypes: ReadonlyMap<string, SimpleAllowedTypeAttributes>;
 }
 
 // @alpha

packages/dds/tree/src/index.ts

@@ -280,6 +280,7 @@ export {
 	type TreeParsingOptions,
 	type SchemaFactory_base,
 	type NumberKeys,
+	type SimpleAllowedTypeAttributes,
 } from "./simple-tree/index.js";
 export {
 	SharedTree,

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

@@ -42,6 +42,7 @@ import {
 	type TreeStoredSchema,
 	TreeStoredSchemaRepository,
 	type TreeStoredSchemaSubscription,
+	type TreeTypeSet,
 	getCodecTreeForDetachedFieldIndexFormat,
 	makeDetachedFieldIndex,
 	moveToDetachedField,
@@ -98,6 +99,7 @@ import {
 	FieldKind,
 	type ITreeAlpha,
 	type SimpleObjectFieldSchema,
+	type SimpleAllowedTypeAttributes,
 } from "../simple-tree/index.js";
 
 import { SchematizingSimpleTreeView } from "./schematizingTreeView.js";
@@ -952,6 +954,24 @@ export const defaultSharedTreeOptions: Required<SharedTreeOptionsInternal> = {
 	shouldEncodeIncrementally: defaultIncrementalEncodingPolicy,
 };
 
+/**
+ * Build the allowed types for a Stored Schema.
+ *
+ * @remarks Staged upgrades do not apply to stored schemas, so we omit the {@link SimpleAllowedTypeAttributes.isStaged | staging flag } when building {@link SimpleAllowedTypeAttributes}.
+ * @param types - The types to create allowed types for.
+ * @returns The allowed types.
+ */
+function buildSimpleAllowedTypeAttributesForStoredSchema(
+	types: TreeTypeSet,
+): ReadonlyMap<string, SimpleAllowedTypeAttributes> {
+	const allowedTypesInfo = new Map<string, SimpleAllowedTypeAttributes>();
+	for (const type of types) {
+		// Stored schemas do not have staged upgrades
+		allowedTypesInfo.set(type, { isStaged: undefined });
+	}
+	return allowedTypesInfo;
+}
+
 function exportSimpleFieldSchemaStored(schema: TreeFieldStoredSchema): SimpleFieldSchema {
 	let kind: FieldKind;
 	switch (schema.kind) {
@@ -973,7 +993,7 @@ function exportSimpleFieldSchemaStored(schema: TreeFieldStoredSchema): SimpleFie
 	}
 	return {
 		kind,
-		allowedTypesIdentifiers: schema.types,
+		simpleAllowedTypes: buildSimpleAllowedTypeAttributesForStoredSchema(schema.types),
 		metadata: {},
 		persistedMetadata: schema.persistedMetadata,
 	};
@@ -989,7 +1009,7 @@ function exportSimpleNodeSchemaStored(schema: TreeNodeStoredSchema): SimpleNodeS
 	if (arrayTypes !== undefined) {
 		return {
 			kind: NodeKind.Array,
-			allowedTypesIdentifiers: arrayTypes,
+			simpleAllowedTypes: buildSimpleAllowedTypeAttributesForStoredSchema(arrayTypes),
 			metadata: {},
 			persistedMetadata: schema.metadata,
 		};
@@ -1008,7 +1028,9 @@ function exportSimpleNodeSchemaStored(schema: TreeNodeStoredSchema): SimpleNodeS
 		);
 		return {
 			kind: NodeKind.Map,
-			allowedTypesIdentifiers: schema.mapFields.types,
+			simpleAllowedTypes: buildSimpleAllowedTypeAttributesForStoredSchema(
+				schema.mapFields.types,
+			),
 			metadata: {},
 			persistedMetadata: schema.metadata,
 		};

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

@@ -13,6 +13,7 @@ import {
 	type FieldProps,
 } from "../fieldSchema.js";
 import type {
+	SimpleAllowedTypeAttributes,
 	SimpleFieldSchema,
 	SimpleNodeSchema,
 	SimpleTreeSchema,
@@ -65,7 +66,7 @@ function generateFieldSchema(
 	context: Context,
 	storedKey: string | undefined,
 ): FieldSchemaAlpha {
-	const allowed = generateAllowedTypes(simple.allowedTypesIdentifiers, context);
+	const allowed = generateAllowedTypes(simple.simpleAllowedTypes, context);
 	const props: Omit<FieldProps, "defaultProvider"> = {
 		metadata: simple.metadata,
 		key: storedKey,
@@ -84,8 +85,14 @@ function generateFieldSchema(
 	}
 }
 
-function generateAllowedTypes(allowed: ReadonlySet<string>, context: Context): AllowedTypes {
-	return [...allowed].map((id) => context.get(id) ?? fail(0xb5a /* Missing schema */));
+function generateAllowedTypes(
+	allowed: ReadonlyMap<string, SimpleAllowedTypeAttributes>,
+	context: Context,
+): AllowedTypes {
+	return Array.from(
+		allowed.keys(),
+		(id) => context.get(id) ?? fail(0xb5a /* Missing schema */),
+	);
 }
 
 function generateNode(
@@ -104,21 +111,17 @@ function generateNode(
 			return factory.objectAlpha(id, fields, { metadata: schema.metadata });
 		}
 		case NodeKind.Array:
-			return factory.arrayAlpha(
-				id,
-				generateAllowedTypes(schema.allowedTypesIdentifiers, context),
-				{ metadata: schema.metadata },
-			);
+			return factory.arrayAlpha(id, generateAllowedTypes(schema.simpleAllowedTypes, context), {
+				metadata: schema.metadata,
+			});
 		case NodeKind.Map:
-			return factory.mapAlpha(
-				id,
-				generateAllowedTypes(schema.allowedTypesIdentifiers, context),
-				{ metadata: schema.metadata },
-			);
+			return factory.mapAlpha(id, generateAllowedTypes(schema.simpleAllowedTypes, context), {
+				metadata: schema.metadata,
+			});
 		case NodeKind.Record:
 			return factory.recordAlpha(
 				id,
-				generateAllowedTypes(schema.allowedTypesIdentifiers, context),
+				generateAllowedTypes(schema.simpleAllowedTypes, context),
 				{ metadata: schema.metadata },
 			);
 		case NodeKind.Leaf:

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

@@ -107,7 +107,10 @@ function convertNodeSchema(
 
 function convertArrayNodeSchema(schema: SimpleArrayNodeSchema): JsonArrayNodeSchema {
 	const allowedTypes: JsonSchemaRef[] = [];
-	schema.allowedTypesIdentifiers.forEach((type) => {
+	const allowedTypesIdentifiers: ReadonlySet<string> = new Set(
+		schema.simpleAllowedTypes.keys(),
+	);
+	allowedTypesIdentifiers.forEach((type) => {
 		allowedTypes.push(createSchemaRef(type));
 	});
 
@@ -206,7 +209,10 @@ function convertRecordLikeNodeSchema(
 	schema: SimpleRecordNodeSchema | SimpleMapNodeSchema,
 ): JsonMapNodeSchema | JsonRecordNodeSchema {
 	const allowedTypes: JsonSchemaRef[] = [];
-	schema.allowedTypesIdentifiers.forEach((type) => {
+	const allowedTypesIdentifiers: ReadonlySet<string> = new Set(
+		schema.simpleAllowedTypes.keys(),
+	);
+	allowedTypesIdentifiers.forEach((type) => {
 		allowedTypes.push(createSchemaRef(type));
 	});
 

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

@@ -6,6 +6,7 @@
 import { assert, unreachableCase } from "@fluidframework/core-utils/internal";
 import { normalizeFieldSchema, type ImplicitFieldSchema } from "../fieldSchema.js";
 import type {
+	SimpleAllowedTypeAttributes,
 	SimpleArrayNodeSchema,
 	SimpleFieldSchema,
 	SimpleLeafNodeSchema,
@@ -31,6 +32,7 @@ import { LeafNodeSchema } from "../leafNodeSchema.js";
  *
  * @param schema - The schema to convert
  * @param copySchemaObjects - If true, TreeNodeSchema and FieldSchema are copied into plain JavaScript objects. Either way, custom metadata is referenced and not copied.
+ * @param preserveViewSchemaProperties - If true (default), properties used by view schema but not part of stored schema (for example, `isStaged` on allowed types) are preserved in the output.
  *
  * @remarks
  * Given that the Schema types used in {@link ImplicitFieldSchema} already implement the {@link SimpleNodeSchema} interfaces, there are limited use-cases for this function.
@@ -45,6 +47,7 @@ import { LeafNodeSchema } from "../leafNodeSchema.js";
 export function toSimpleTreeSchema(
 	schema: ImplicitFieldSchema,
 	copySchemaObjects: boolean,
+	preserveViewSchemaProperties: boolean = true,
 ): SimpleTreeSchema {
 	const normalizedSchema = normalizeFieldSchema(schema);
 	const definitions = new Map<string, SimpleNodeSchema>();
@@ -59,15 +62,20 @@ export function toSimpleTreeSchema(
 					nodeSchema instanceof RecordNodeSchema,
 				0xb60 /* Invalid schema */,
 			);
-			const outSchema = copySchemaObjects ? copySimpleNodeSchema(nodeSchema) : nodeSchema;
+			const outSchema = copySchemaObjects
+				? copySimpleNodeSchema(nodeSchema, preserveViewSchemaProperties)
+				: nodeSchema;
 			definitions.set(nodeSchema.identifier, outSchema);
 		},
 	});
 
 	return {
 		root: copySchemaObjects
 			? ({
-					allowedTypesIdentifiers: normalizedSchema.allowedTypesIdentifiers,
+					simpleAllowedTypes: normalizeSimpleAllowedTypes(
+						normalizedSchema.simpleAllowedTypes,
+						preserveViewSchemaProperties,
+					),
 					kind: normalizedSchema.kind,
 					metadata: normalizedSchema.metadata,
 					persistedMetadata: normalizedSchema.persistedMetadata,
@@ -77,22 +85,46 @@ export function toSimpleTreeSchema(
 	};
 }
 
+/**
+ * Normalizes the {@link SimpleAllowedTypeAttributes} by either preserving or omitting view-specific schema properties.
+ * @param simpleAllowedTypes - The simple allowed types to normalize.
+ * @param preserveViewSchemaProperties - If true, properties used by view schema but not part of stored schema (for example, `isStaged` on allowed types) are preserved in the output.
+ * @returns The normalized simple allowed types.
+ */
+function normalizeSimpleAllowedTypes(
+	simpleAllowedTypes: ReadonlyMap<string, SimpleAllowedTypeAttributes>,
+	preserveViewSchemaProperties: boolean,
+): ReadonlyMap<string, SimpleAllowedTypeAttributes> {
+	if (preserveViewSchemaProperties) {
+		return simpleAllowedTypes;
+	} else {
+		const normalized = new Map<string, SimpleAllowedTypeAttributes>();
+		for (const [identifier, attributes] of simpleAllowedTypes.entries()) {
+			normalized.set(identifier, { ...attributes, isStaged: undefined });
+		}
+		return normalized;
+	}
+}
+
 /**
  * Copies a {@link SimpleNodeSchema} into a new plain JavaScript object.
  *
  * @remarks Caches the result on the input schema for future calls.
  */
-function copySimpleNodeSchema(schema: SimpleNodeSchema): SimpleNodeSchema {
+function copySimpleNodeSchema(
+	schema: SimpleNodeSchema,
+	preserveViewSchemaProperties: boolean,
+): SimpleNodeSchema {
 	const kind = schema.kind;
 	switch (kind) {
 		case NodeKind.Leaf:
 			return copySimpleLeafSchema(schema);
 		case NodeKind.Array:
 		case NodeKind.Map:
 		case NodeKind.Record:
-			return copySimpleSchemaWithAllowedTypes(schema);
+			return copySimpleSchemaWithAllowedTypes(schema, preserveViewSchemaProperties);
 		case NodeKind.Object:
-			return copySimpleObjectSchema(schema);
+			return copySimpleObjectSchema(schema, preserveViewSchemaProperties);
 		default:
 			unreachableCase(kind);
 	}
@@ -109,22 +141,32 @@ function copySimpleLeafSchema(schema: SimpleLeafNodeSchema): SimpleLeafNodeSchem
 
 function copySimpleSchemaWithAllowedTypes(
 	schema: SimpleMapNodeSchema | SimpleArrayNodeSchema | SimpleRecordNodeSchema,
+	preserveViewSchemaProperties: boolean,
 ): SimpleMapNodeSchema | SimpleArrayNodeSchema | SimpleRecordNodeSchema {
 	return {
 		kind: schema.kind,
-		allowedTypesIdentifiers: schema.allowedTypesIdentifiers,
+		simpleAllowedTypes: normalizeSimpleAllowedTypes(
+			schema.simpleAllowedTypes,
+			preserveViewSchemaProperties,
+		),
 		metadata: schema.metadata,
 		persistedMetadata: schema.persistedMetadata,
 	};
 }
 
-function copySimpleObjectSchema(schema: SimpleObjectNodeSchema): SimpleObjectNodeSchema {
+function copySimpleObjectSchema(
+	schema: SimpleObjectNodeSchema,
+	preserveViewSchemaProperties: boolean,
+): SimpleObjectNodeSchema {
 	const fields: Map<string, SimpleObjectFieldSchema> = new Map();
 	for (const [propertyKey, field] of schema.fields) {
 		// field already is a SimpleObjectFieldSchema, but copy the subset of the properties needed by this interface to get a clean simple object.
 		fields.set(propertyKey, {
 			kind: field.kind,
-			allowedTypesIdentifiers: field.allowedTypesIdentifiers,
+			simpleAllowedTypes: normalizeSimpleAllowedTypes(
+				field.simpleAllowedTypes,
+				preserveViewSchemaProperties,
+			),
 			metadata: field.metadata,
 			persistedMetadata: field.persistedMetadata,
 			storedKey: field.storedKey,

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

@@ -25,6 +25,7 @@ import {
 	type TreeNodeSchema,
 } from "./treeNodeSchema.js";
 import { schemaAsTreeNodeValid } from "./treeNodeValid.js";
+import type { SimpleAllowedTypeAttributes } from "../simpleSchema.js";
 
 /**
  * Schema for types allowed in some location in a tree (like a field, map entry or array).
@@ -252,6 +253,21 @@ export class AnnotatedAllowedTypesInternal<
 		return this.lazyEvaluate.value.identifiers;
 	}
 
+	/**
+	 * Get the {@link SimpleAllowedTypeAttributes} version of the allowed types set.
+	 */
+	public static evaluateSimpleAllowedTypes(
+		annotatedAllowedTypes: AnnotatedAllowedTypes,
+	): ReadonlyMap<string, SimpleAllowedTypeAttributes> {
+		const simpleAllowedTypes = new Map<string, SimpleAllowedTypeAttributes>();
+		for (const type of annotatedAllowedTypes.evaluate().types) {
+			simpleAllowedTypes.set(type.type.identifier, {
+				isStaged: type.metadata.stagedSchemaUpgrade !== undefined,
+			});
+		}
+		return simpleAllowedTypes;
+	}
+
 	public static override [Symbol.hasInstance]<TThis extends { prototype: object }>(
 		this: TThis,
 		value: unknown,