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

packages/dds/tree/src/index.ts

@@ -277,6 +277,7 @@ export {
 	type TreeParsingOptions,
 	type SchemaFactory_base,
 	type NumberKeys,
+	type AllowedTypeInfo,
 } 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 AllowedTypeInfo,
 } from "../simple-tree/index.js";
 
 import { SchematizingSimpleTreeView } from "./schematizingTreeView.js";
@@ -950,6 +952,16 @@ export const defaultSharedTreeOptions: Required<SharedTreeOptionsInternal> = {
 	shouldEncodeIncrementally: defaultIncrementalEncodingPolicy,
 };
 
+function buildSimpleAllowedTypes(types: TreeTypeSet): ReadonlyMap<string, AllowedTypeInfo> {
+	const allowedTypesInfo = new Map<string, AllowedTypeInfo>();
+	for (const type of types) {
+		allowedTypesInfo.set(type, {
+			isStaged: false,
+		});
+	}
+	return allowedTypesInfo;
+}
+
 function exportSimpleFieldSchemaStored(schema: TreeFieldStoredSchema): SimpleFieldSchema {
 	let kind: FieldKind;
 	switch (schema.kind) {
@@ -971,7 +983,8 @@ function exportSimpleFieldSchemaStored(schema: TreeFieldStoredSchema): SimpleFie
 	}
 	return {
 		kind,
-		allowedTypesIdentifiers: schema.types,
+		// TODO: Refactor
+		simpleAllowedTypes: buildSimpleAllowedTypes(schema.types),
 		metadata: {},
 		persistedMetadata: schema.persistedMetadata,
 	};
@@ -987,7 +1000,7 @@ function exportSimpleNodeSchemaStored(schema: TreeNodeStoredSchema): SimpleNodeS
 	if (arrayTypes !== undefined) {
 		return {
 			kind: NodeKind.Array,
-			allowedTypesIdentifiers: arrayTypes,
+			simpleAllowedTypes: buildSimpleAllowedTypes(arrayTypes),
 			metadata: {},
 			persistedMetadata: schema.metadata,
 		};
@@ -1006,7 +1019,7 @@ function exportSimpleNodeSchemaStored(schema: TreeNodeStoredSchema): SimpleNodeS
 		);
 		return {
 			kind: NodeKind.Map,
-			allowedTypesIdentifiers: schema.mapFields.types,
+			simpleAllowedTypes: buildSimpleAllowedTypes(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 {
+	SimpleAllowedTypes,
 	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,13 @@ 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, SimpleAllowedTypes>,
+	context: Context,
+): AllowedTypes {
+	return [...new Set(allowed.keys())].map(
+		(id) => context.get(id) ?? fail(0xb5a /* Missing schema */),
+	);
 }
 
 function generateNode(
@@ -104,21 +110,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

@@ -67,7 +67,7 @@ export function toSimpleTreeSchema(
 	return {
 		root: copySchemaObjects
 			? ({
-					allowedTypesIdentifiers: normalizedSchema.allowedTypesIdentifiers,
+					simpleAllowedTypes: normalizedSchema.simpleAllowedTypes,
 					kind: normalizedSchema.kind,
 					metadata: normalizedSchema.metadata,
 					persistedMetadata: normalizedSchema.persistedMetadata,
@@ -112,7 +112,7 @@ function copySimpleSchemaWithAllowedTypes(
 ): SimpleMapNodeSchema | SimpleArrayNodeSchema | SimpleRecordNodeSchema {
 	return {
 		kind: schema.kind,
-		allowedTypesIdentifiers: schema.allowedTypesIdentifiers,
+		simpleAllowedTypes: schema.simpleAllowedTypes,
 		metadata: schema.metadata,
 		persistedMetadata: schema.persistedMetadata,
 	};
@@ -124,7 +124,7 @@ function copySimpleObjectSchema(schema: SimpleObjectNodeSchema): SimpleObjectNod
 		// 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: field.simpleAllowedTypes,
 			metadata: field.metadata,
 			persistedMetadata: field.persistedMetadata,
 			storedKey: field.storedKey,

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

@@ -31,7 +31,7 @@ import type {
 } from "./core/index.js";
 import { normalizeAllowedTypes } from "./core/index.js";
 
-import type { SimpleFieldSchema } from "./simpleSchema.js";
+import type { SimpleAllowedTypes, SimpleFieldSchema } from "./simpleSchema.js";
 import type { UnsafeUnknownSchema } from "./unsafeUnknownSchema.js";
 import type { InsertableContent } from "./unhydratedFlexTreeFromInsertable.js";
 
@@ -412,6 +412,19 @@ export class FieldSchemaAlpha<
 		return this.allowedTypesFull.evaluateIdentifiers();
 	}
 
+	public get simpleAllowedTypes(): Map<string, SimpleAllowedTypes> {
+		const types = this.allowedTypesFull.evaluate().types;
+		const info = new Map<string, SimpleAllowedTypes>();
+
+		for (const type of types) {
+			info.set(type.type.identifier, {
+				isStaged: type.metadata.stagedSchemaUpgrade !== undefined,
+			});
+		}
+
+		return info;
+	}
+
 	protected constructor(
 		kind: Kind,
 		types: Types,

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

@@ -193,6 +193,7 @@ export type {
 	SimpleNodeSchemaBaseAlpha,
 	SimpleObjectFieldSchema,
 	SimpleRecordNodeSchema,
+	SimpleAllowedTypes as AllowedTypeInfo,
 } from "./simpleSchema.js";
 export {
 	type ImplicitFieldSchema,

packages/dds/tree/src/simple-tree/node-kinds/array/arrayNode.ts

@@ -68,6 +68,7 @@ import type {
 import { brand, type JsonCompatibleReadOnlyObject } from "../../../util/index.js";
 import { nullSchema } from "../../leafNodeSchema.js";
 import { arrayNodeStoredSchema } from "../../toStoredSchema.js";
+import type { SimpleAllowedTypes } from "../../simpleSchema.js";
 
 /**
  * A covariant base type for {@link (TreeArrayNode:interface)}.
@@ -1168,6 +1169,15 @@ export function arraySchema<
 	const lazyAllowedTypesIdentifiers = new Lazy(
 		() => new Set(normalizedTypes.evaluate().map((type) => type.identifier)),
 	);
+	const lazySimpleAllowedTypes = new Lazy(() => {
+		const map = new Map<string, SimpleAllowedTypes>();
+		for (const type of normalizedTypes.evaluate()) {
+			map.set(type.identifier, {
+				isStaged: false,
+			});
+		}
+		return map;
+	});
 
 	let privateData: TreeNodeSchemaPrivateData | undefined;
 
@@ -1206,6 +1216,10 @@ export function arraySchema<
 			return lazyAllowedTypesIdentifiers.value;
 		}
 
+		public static get simpleAllowedTypes(): ReadonlyMap<string, SimpleAllowedTypes> {
+			return lazySimpleAllowedTypes.value;
+		}
+
 		protected static override constructorCached: MostDerivedData | undefined = undefined;
 
 		protected static override oneTimeSetup(): TreeNodeSchemaInitializedData {

packages/dds/tree/src/simple-tree/node-kinds/map/mapNode.ts

@@ -66,6 +66,7 @@ import type {
 import { recordLikeDataToFlexContent } from "../common.js";
 import { MapNodeStoredSchema } from "../../../core/index.js";
 import type { NodeSchemaOptionsAlpha } from "../../api/index.js";
+import type { SimpleAllowedTypes } from "../../simpleSchema.js";
 
 /**
  * A map of string keys to tree objects.
@@ -275,6 +276,15 @@ export function mapSchema<
 	const lazyAllowedTypesIdentifiers = new Lazy(
 		() => new Set(normalizedTypes.evaluate().map((type) => type.identifier)),
 	);
+	const lazySimpleAllowedTypes = new Lazy(() => {
+		const map = new Map<string, SimpleAllowedTypes>();
+		for (const type of normalizedTypes.evaluate()) {
+			map.set(type.identifier, {
+				isStaged: false,
+			});
+		}
+		return map;
+	});
 
 	let privateData: TreeNodeSchemaPrivateData | undefined;
 	const persistedMetadata = nodeOptions.persistedMetadata;
@@ -303,6 +313,10 @@ export function mapSchema<
 			return lazyAllowedTypesIdentifiers.value;
 		}
 
+		public static get simpleAllowedTypes(): ReadonlyMap<string, SimpleAllowedTypes> {
+			return lazySimpleAllowedTypes.value;
+		}
+
 		protected static override constructorCached: MostDerivedData | undefined = undefined;
 
 		protected static override oneTimeSetup(): TreeNodeSchemaInitializedData {

packages/dds/tree/src/simple-tree/node-kinds/record/recordNode.ts

@@ -58,6 +58,7 @@ import { prepareForInsertion } from "../../prepareForInsertion.js";
 import { recordLikeDataToFlexContent } from "../common.js";
 import { MapNodeStoredSchema } from "../../../core/index.js";
 import type { NodeSchemaOptionsAlpha } from "../../api/index.js";
+import type { SimpleAllowedTypes } from "../../simpleSchema.js";
 
 /**
  * Create a proxy which implements the {@link TreeRecordNode} API.
@@ -258,6 +259,15 @@ export function recordSchema<
 	const lazyAllowedTypesIdentifiers = new Lazy(
 		() => new Set(normalizedTypes.evaluate().map((type) => type.identifier)),
 	);
+	const lazySimpleAllowedTypes = new Lazy(() => {
+		const map = new Map<string, SimpleAllowedTypes>();
+		for (const type of normalizedTypes.evaluate()) {
+			map.set(type.identifier, {
+				isStaged: false,
+			});
+		}
+		return map;
+	});
 
 	let privateData: TreeNodeSchemaPrivateData | undefined;
 
@@ -347,6 +357,10 @@ export function recordSchema<
 			return lazyAllowedTypesIdentifiers.value;
 		}
 
+		public static get simpleAllowedTypes(): ReadonlyMap<string, SimpleAllowedTypes> {
+			return lazySimpleAllowedTypes.value;
+		}
+
 		protected static override constructorCached: MostDerivedData | undefined = undefined;
 
 		public static readonly identifier = identifier;