@fluidframework/tree 2.51.0 → 2.53.0-350190

package/src/simple-tree/core/allowedTypes.ts

@@ -5,7 +5,12 @@
 
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
-import { getOrCreate, isReadonlyArray, type IsUnion } from "../../util/index.js";
+import {
+	getOrCreate,
+	isReadonlyArray,
+	type IsUnion,
+	type MakeNominal,
+} from "../../util/index.js";
 import { isLazy, type FlexListToUnion, type LazyItem } from "./flexList.js";
 import {
 	NodeKind,
@@ -76,6 +81,7 @@ export interface NormalizedAnnotatedAllowedTypes {
 export function isAnnotatedAllowedTypes(
 	allowedTypes: ImplicitAnnotatedAllowedTypes,
 ): allowedTypes is AnnotatedAllowedTypes {
+	checkForUninitializedSchema(allowedTypes);
 	return (
 		// Class based schema, and lazy schema references report type "function": filtering them out with typeof makes narrowing based on members mostly safe
 		typeof allowedTypes === "object" && "metadata" in allowedTypes && "types" in allowedTypes
@@ -116,7 +122,9 @@ export interface AllowedTypesMetadata {
 export function isAnnotatedAllowedType(
 	allowedType: AnnotatedAllowedType | LazyItem<TreeNodeSchema>,
 ): allowedType is AnnotatedAllowedType {
-	return "metadata" in allowedType && "type" in allowedType;
+	checkForUninitializedSchema(allowedType);
+	// Class based schema, and lazy schema references report type "function": filtering them out with typeof makes narrowing based on members mostly safe
+	return typeof allowedType === "object" && "metadata" in allowedType && "type" in allowedType;
 }
 
 /**
@@ -131,7 +139,34 @@ export interface AllowedTypeMetadata {
 	 */
 	readonly custom?: unknown;
 
-	// TODO metadata for enablable types will be added here
+	/**
+	 * If defined, indicates that an allowed type is {@link SchemaStaticsAlpha.staged | staged}.
+	 */
+	readonly stagedSchemaUpgrade?: SchemaUpgrade;
+}
+
+/**
+ * Package internal {@link SchemaUpgrade} construction API.
+ */
+export let createSchemaUpgrade: () => SchemaUpgrade;
+
+/**
+ * Unique token used to upgrade schemas and determine if a particular upgrade has been completed.
+ * @remarks
+ * Create using {@link SchemaStaticsAlpha.staged}.
+ * @privateRemarks
+ * TODO:#38722 implement runtime schema upgrades.
+ * Until then, the class purely behaves mostly as a placeholder.
+ * TODO: Consider allowing users to store a name for the upgrade to use in error messages.
+ * @sealed @alpha
+ */
+export class SchemaUpgrade {
+	protected _typeCheck!: MakeNominal;
+	static {
+		createSchemaUpgrade = () => new SchemaUpgrade();
+	}
+
+	private constructor() {}
 }
 
 /**
@@ -188,7 +223,7 @@ export type UnannotateImplicitAllowedTypes<T extends ImplicitAnnotatedAllowedTyp
 	T extends AnnotatedAllowedTypes
 		? UnannotateAllowedTypes<T>
 		: T extends AnnotatedAllowedType
-			? UnannotateAllowedType<T>
+			? UnannotateAllowedTypesList<[T]>
 			: T extends readonly (AnnotatedAllowedType | LazyItem<TreeNodeSchema>)[]
 				? UnannotateAllowedTypesList<T>
 				: T extends TreeNodeSchema
@@ -202,17 +237,9 @@ export type UnannotateImplicitAllowedTypes<T extends ImplicitAnnotatedAllowedTyp
 export type UnannotateAllowedTypesList<
 	T extends readonly (AnnotatedAllowedType | LazyItem<TreeNodeSchema>)[],
 > = {
-	[I in keyof T]: UnannotateAllowedTypeOrLazyItem<T[I]>;
+	[I in keyof T]: UnannotateAllowedType<T[I]>;
 };
 
-/**
- * Removes annotations from an allowed type that may contain annotations.
- * @system @alpha
- */
-export type UnannotateAllowedTypeOrLazyItem<
-	T extends AnnotatedAllowedType | LazyItem<TreeNodeSchema>,
-> = T extends AnnotatedAllowedType<infer X> ? X : T;
-
 /**
  * Removes all annotations from a set of allowed types.
  * @system @alpha
@@ -222,10 +249,15 @@ export type UnannotateAllowedTypes<T extends AnnotatedAllowedTypes> =
 
 /**
  * Removes annotations from an allowed type.
+ * @remarks
+ * If the input could be lazy
+ * (is a LazyItem or AnnotatedAllowedType<LazyItem> instead of just a TreeNodeSchema | AnnotatedAllowedType<TreeNodeSchema>)
+ * then the output of this will be a LazyItem and thus is not valid as an ImplicitAllowedTypes without wrapping it in an array.
+ * This can however be used on items within an AllowedTypes array.
  * @system @alpha
  */
-export type UnannotateAllowedType<T extends AnnotatedAllowedType> =
-	T extends AnnotatedAllowedType<infer X> ? [X] : T;
+export type UnannotateAllowedType<T extends AnnotatedAllowedType | LazyItem<TreeNodeSchema>> =
+	T extends AnnotatedAllowedType<infer X> ? X : T;
 
 /**
  * Normalizes a {@link ImplicitAllowedTypes} to a set of {@link TreeNodeSchema}s, by eagerly evaluating any
@@ -237,17 +269,19 @@ export type UnannotateAllowedType<T extends AnnotatedAllowedType> =
  * @internal
  */
 export function normalizeAllowedTypes(
-	types: ImplicitAllowedTypes,
+	types: ImplicitAnnotatedAllowedTypes,
 ): ReadonlySet<TreeNodeSchema> {
+	// remove annotations before normalizing
+	const unannotated = unannotateImplicitAllowedTypes(types);
 	const normalized = new Set<TreeNodeSchema>();
-	if (isReadonlyArray(types)) {
+	if (isReadonlyArray(unannotated)) {
 		// Types array must not be modified after it is normalized since that would result in the user of the normalized data having wrong (out of date) content.
-		Object.freeze(types);
-		for (const lazyType of types) {
+		Object.freeze(unannotated);
+		for (const lazyType of unannotated) {
 			normalized.add(evaluateLazySchema(lazyType));
 		}
 	} else {
-		normalized.add(evaluateLazySchema(types));
+		normalized.add(evaluateLazySchema(unannotated));
 	}
 	return normalized;
 }
@@ -255,9 +289,9 @@ export function normalizeAllowedTypes(
 /**
  * Normalizes an allowed type to an {@link AnnotatedAllowedType}, by adding empty annotations if they don't already exist.
  */
-export function normalizeToAnnotatedAllowedType<T extends TreeNodeSchema>(
-	type: T | AnnotatedAllowedType<T> | AnnotatedAllowedType<LazyItem<T>>,
-): AnnotatedAllowedType<T> | AnnotatedAllowedType<LazyItem<T>> {
+export function normalizeToAnnotatedAllowedType<T extends LazyItem<TreeNodeSchema>>(
+	type: T | AnnotatedAllowedType<T>,
+): AnnotatedAllowedType<T> {
 	return isAnnotatedAllowedType(type)
 		? type
 		: {
@@ -317,19 +351,28 @@ export function unannotateImplicitAllowedTypes<Types extends ImplicitAnnotatedAl
 ): UnannotateImplicitAllowedTypes<Types> {
 	return (
 		isAnnotatedAllowedTypes(types)
-			? types.types.map((allowedType) =>
-					isAnnotatedAllowedType(allowedType) ? allowedType.type : allowedType,
-				)
+			? types.types.map(unannotateAllowedType)
 			: isReadonlyArray(types)
-				? types.map((allowedType) =>
-						isAnnotatedAllowedType(allowedType) ? allowedType.type : allowedType,
-					)
+				? types.map(unannotateAllowedType)
 				: isAnnotatedAllowedType(types)
-					? (types.type as UnannotateImplicitAllowedTypes<Types>)
+					? types.type
 					: types
 	) as UnannotateImplicitAllowedTypes<Types>;
 }
 
+/**
+ * Converts an {@link AnnotatedAllowedType} to an {@link LazyItem} by removing any annotations.
+ * @remarks
+ * This does not evaluate any lazy schemas.
+ */
+export function unannotateAllowedType<
+	Type extends AnnotatedAllowedType | LazyItem<TreeNodeSchema>,
+>(allowedType: Type): UnannotateAllowedType<Type> {
+	return isAnnotatedAllowedType(allowedType)
+		? (allowedType.type as UnannotateAllowedType<Type>)
+		: (allowedType as UnannotateAllowedType<Type>);
+}
+
 const cachedLazyItem = new WeakMap<() => unknown, unknown>();
 
 /**
@@ -342,13 +385,23 @@ export function evaluateLazySchema<T extends TreeNodeSchema>(value: LazyItem<T>)
 	const evaluatedSchema = isLazy(value)
 		? (getOrCreate(cachedLazyItem, value, value) as T)
 		: value;
-	if (evaluatedSchema === undefined) {
+
+	checkForUninitializedSchema(evaluatedSchema);
+	markSchemaMostDerived(evaluatedSchema);
+	return evaluatedSchema;
+}
+
+/**
+ * Throws a UsageError if the provided schema is undefined, most likely due to being used before it was initialized.
+ */
+export function checkForUninitializedSchema(
+	schema: ImplicitAnnotatedAllowedTypes | LazyItem<TreeNodeSchema>,
+): void {
+	if (schema === undefined) {
 		throw new UsageError(
-			`Encountered an undefined schema. This could indicate that some referenced schema has not yet been instantiated.`,
+			`Encountered an undefined schema. This could indicate that some referenced schema has not yet been instantiated. Consider using a lazy schema reference (like "() => schema") or delaying the evaluation of the lazy reference if one is already being used.`,
 		);
 	}
-	markSchemaMostDerived(evaluatedSchema);
-	return evaluatedSchema;
 }
 
 /**

package/src/simple-tree/core/context.ts

@@ -28,34 +28,36 @@ import { walkAllowedTypes } from "./walkSchema.js";
  * This design is the same as {@link FlexTreeContext} with its base type and {@link FlexTreeHydratedContext} extending it.
  */
 export class Context {
-	/**
-	 * All schema which could transitively be used under the associated node.
-	 * @remarks
-	 * While generally {@link TreeNodeSchema} are referenced as objects and thus do not need to be looked up by identifier,
-	 * there are a few cases (mainly constructing new TreeNodes from existing tree data) where such a lookup is useful.
-	 * Having this map in the context addresses this use-case.
-	 * @privateRemarks
-	 * This design mirrors how {@link FlexTreeSchema} are accessed off the {@link FlexTreeContext}, making the migration away from them simpler.
-	 */
-	public readonly schema: ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema>;
-
-	/**
-	 * Builds the context.
-	 * @remarks
-	 * Since this walks the schema, it must not be invoked during schema declaration or schema forward references could fail to be resolved.
-	 */
-	public constructor(
+	public static schemaMapFromRootSchema(
 		rootSchema: NormalizedAnnotatedAllowedTypes,
-		public readonly flexContext: FlexTreeContext,
-	) {
+	): ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema> {
 		const schema: Map<TreeNodeSchemaIdentifier, TreeNodeSchema> = new Map();
 		walkAllowedTypes(rootSchema, {
 			node(nodeSchema) {
 				schema.set(brand(nodeSchema.identifier), nodeSchema);
 			},
 		});
-		this.schema = schema;
+		return schema;
 	}
+
+	/**
+	 * Builds the context.
+	 * @remarks
+	 * Since this walks the schema, it must not be invoked during schema declaration or schema forward references could fail to be resolved.
+	 */
+	public constructor(
+		public readonly flexContext: FlexTreeContext,
+		/**
+		 * All schema which could transitively be used under the associated node.
+		 * @remarks
+		 * While generally {@link TreeNodeSchema} are referenced as objects and thus do not need to be looked up by identifier,
+		 * there are a few cases (mainly constructing new TreeNodes from existing tree data) where such a lookup is useful.
+		 * Having this map in the context addresses this use-case.
+		 * @privateRemarks
+		 * This design mirrors how {@link FlexTreeSchema} are accessed off the {@link FlexTreeContext}, making the migration away from them simpler.
+		 */
+		public readonly schema: ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema>,
+	) {}
 }
 
 /**
@@ -64,9 +66,9 @@ export class Context {
  */
 export class HydratedContext extends Context {
 	public constructor(
-		rootSchema: NormalizedAnnotatedAllowedTypes,
 		public override readonly flexContext: FlexTreeHydratedContext,
+		schema: ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema>,
 	) {
-		super(rootSchema, flexContext);
+		super(flexContext, schema);
 	}
 }

package/src/simple-tree/core/index.ts

@@ -53,9 +53,12 @@ export {
 	isAnnotatedAllowedType,
 	normalizeAllowedTypes,
 	normalizeAnnotatedAllowedTypes,
+	normalizeToAnnotatedAllowedType,
 	unannotateImplicitAllowedTypes,
 	markSchemaMostDerived,
 	evaluateLazySchema,
+	createSchemaUpgrade,
+	type SchemaUpgrade,
 } from "./allowedTypes.js";
 export type {
 	AnnotatedAllowedType,
@@ -72,11 +75,10 @@ export type {
 	UnannotateAllowedTypes,
 	UnannotateAllowedType,
 	UnannotateAllowedTypesList,
-	UnannotateAllowedTypeOrLazyItem,
 	AllowedTypeMetadata,
 	AnnotatedAllowedTypes,
 } from "./allowedTypes.js";
-export { walkAllowedTypes, type SchemaVisitor } from "./walkSchema.js";
+export { walkAllowedTypes, walkNodeSchema, type SchemaVisitor } from "./walkSchema.js";
 export { Context, HydratedContext } from "./context.js";
 export {
 	getOrCreateNodeFromInnerNode,
@@ -100,4 +102,11 @@ export {
 	TreeNodeValid,
 	type MostDerivedData,
 	createTreeNodeSchemaPrivateData,
+	isClassBasedSchema,
 } from "./treeNodeValid.js";
+export type { SimpleNodeSchemaBase } from "./simpleNodeSchemaBase.js";
+export {
+	type StoredSchemaGenerationOptions,
+	convertAllowedTypes,
+	allowedTypeFilter,
+} from "./toStored.js";

package/src/simple-tree/core/simpleNodeSchemaBase.ts

@@ -0,0 +1,30 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import type { NodeKind, NodeSchemaMetadata } from "./treeNodeSchema.js";
+
+/**
+ * Base interface for {@link TreeNodeSchema} and {@link SimpleNodeSchema} types.
+ * Once simple schema is stable this doesn't have a reason to be kept `@system`, but it could be.
+ * @system
+ * @public
+ * @sealed
+ */
+export interface SimpleNodeSchemaBase<
+	out TNodeKind extends NodeKind,
+	out TCustomMetadata = unknown,
+> {
+	/**
+	 * The {@link NodeKind}.
+	 *
+	 * @remarks can be used to type-switch between implementations.
+	 */
+	readonly kind: TNodeKind;
+
+	/**
+	 * User-provided {@link NodeSchemaMetadata} for this schema.
+	 */
+	readonly metadata: NodeSchemaMetadata<TCustomMetadata>;
+}

package/src/simple-tree/core/toStored.ts

@@ -0,0 +1,58 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { brand } from "../../util/index.js";
+import {
+	normalizeAnnotatedAllowedTypes,
+	type AnnotatedAllowedType,
+	type ImplicitAnnotatedAllowedTypes,
+	type SchemaUpgrade,
+} from "./allowedTypes.js";
+import type { TreeNodeSchemaIdentifier, TreeTypeSet } from "../../core/index.js";
+
+/**
+ * Options for generating a {@link TreeStoredSchema} from view schema.
+ */
+export interface StoredSchemaGenerationOptions {
+	/**
+	 * Determines whether to include staged schema in the resulting stored schema.
+	 * @remarks
+	 * Due to caching, the behavior of this function must be pure.
+	 */
+	includeStaged(upgrade: SchemaUpgrade): boolean;
+}
+
+/**
+ * Filters an allowed type based on the provided options.
+ * @param allowedType - The allowed type to filter.
+ * @param options - The options to use for filtering.
+ * @returns Whether the allowed type passes the filter.
+ */
+export function allowedTypeFilter(
+	allowedType: AnnotatedAllowedType,
+	options: StoredSchemaGenerationOptions,
+): boolean {
+	// If the allowed type is staged, only include it if the options allow it.
+	if (allowedType.metadata.stagedSchemaUpgrade !== undefined) {
+		return options.includeStaged(allowedType.metadata.stagedSchemaUpgrade);
+	}
+	return true;
+}
+
+/**
+ * Converts an ImplicitAnnotatedAllowedTypes to a stored schema.
+ * @param schema - The schema to convert.
+ * @param options - The options to use for filtering.
+ * @returns The converted stored schema.
+ */
+export function convertAllowedTypes(
+	schema: ImplicitAnnotatedAllowedTypes,
+	options: StoredSchemaGenerationOptions,
+): TreeTypeSet {
+	const filtered: TreeNodeSchemaIdentifier[] = normalizeAnnotatedAllowedTypes(schema)
+		.types.filter((allowedType) => allowedTypeFilter(allowedType, options))
+		.map((a) => brand(a.type.identifier));
+	return new Set(filtered);
+}

package/src/simple-tree/core/treeNodeKernel.ts

@@ -131,7 +131,7 @@ export class TreeNodeKernel {
 	/**
 	 * Create a TreeNodeKernel which can be looked up with {@link getKernel}.
 	 *
-	 * @param initialContext - context from when this node was originally crated.
+	 * @param initialContext - context from when this node was originally created. Only used when unhydrated.
 	 * @param innerNode - When unhydrated/raw or marinated the MapTreeNode. FlexTreeNode when cooked.
 	 * @remarks
 	 * Exactly one kernel per TreeNode should be created.

package/src/simple-tree/core/treeNodeSchema.ts

@@ -6,7 +6,7 @@
 import { assert } from "@fluidframework/core-utils/internal";
 import type { IFluidHandle } from "@fluidframework/core-interfaces";
 
-import type { SimpleNodeSchemaBase } from "../simpleSchema.js";
+import type { SimpleNodeSchemaBase } from "./simpleNodeSchemaBase.js";
 import type { TreeNode } from "./treeNode.js";
 import type { InternalTreeNode, Unhydrated } from "./types.js";
 import type { UnionToIntersection } from "../../util/index.js";
@@ -15,9 +15,10 @@ import type {
 	NormalizedAnnotatedAllowedTypes,
 } from "./allowedTypes.js";
 import type { Context } from "./context.js";
-import type { FieldKey, NodeData } from "../../core/index.js";
+import type { FieldKey, NodeData, TreeNodeStoredSchema } from "../../core/index.js";
 import type { UnhydratedFlexTreeField } from "./unhydratedFlexTree.js";
 import type { FactoryContent } from "../unhydratedFlexTreeFromInsertable.js";
+import type { StoredSchemaGenerationOptions } from "./toStored.js";
 
 /**
  * Schema for a {@link TreeNode} or {@link TreeLeafValue}.
@@ -396,6 +397,11 @@ export interface TreeNodeSchemaPrivateData {
 	 * Idempotent initialization function that pre-caches data and can dereference lazy schema references.
 	 */
 	idempotentInitialize(): TreeNodeSchemaInitializedData;
+
+	/**
+	 * Converts a the schema into a {@link TreeNodeStoredSchema}.
+	 */
+	toStored(options: StoredSchemaGenerationOptions): TreeNodeStoredSchema;
 }
 
 /**

package/src/simple-tree/core/treeNodeValid.ts

@@ -10,7 +10,7 @@ import { type FlexTreeNode, isFlexTreeNode } from "../../feature-libraries/index
 
 import { markEager } from "./flexList.js";
 import { inPrototypeChain, privateToken, TreeNode } from "./treeNode.js";
-import type { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
+import { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
 import {
 	NodeKind,
 	type TreeNodeSchema,
@@ -27,6 +27,7 @@ import {
 import type { InternalTreeNode } from "./types.js";
 import { typeSchemaSymbol } from "./withType.js";
 import type { ImplicitAnnotatedAllowedTypes } from "./allowedTypes.js";
+import type { SimpleNodeSchemaBase } from "./simpleNodeSchemaBase.js";
 
 /**
  * Class which all {@link TreeNode}s must extend.
@@ -193,10 +194,17 @@ export abstract class TreeNodeValid<TInput> extends TreeNode {
 		);
 
 		const result = schema.prepareInstance(this, node);
+		// For unhydrated nodes, grab the context from the `simpleContext` (which likely is from cache.oneTimeInitialized.context,
+		// but might be customized for better schema evolution support like unknown optional fields like in clone)
+		// For hydrated nodes, this context is unused, so using the default from cache.oneTimeInitialized.context is always fine.
+		const context =
+			node instanceof UnhydratedFlexTreeNode
+				? node.simpleContext
+				: cache.oneTimeInitialized.context;
 		// The TreeNodeKernel associates itself the TreeNode (result here, not node) so it can be looked up later via getKernel.
 		// If desired this could be put in a non-enumerable symbol property for lookup instead, but that gets messy going through proxies,
 		// so just relying on the WeakMap seems like the cleanest approach.
-		new TreeNodeKernel(result, schema, node, cache.oneTimeInitialized.context);
+		new TreeNodeKernel(result, schema, node, context);
 
 		return result;
 	}
@@ -226,7 +234,7 @@ export interface MostDerivedData {
 export function schemaAsTreeNodeValid(
 	schema: TreeNodeSchemaCore<string, NodeKind, boolean>,
 ): typeof TreeNodeValid & TreeNodeSchema {
-	if (!inPrototypeChain(schema, TreeNodeValid)) {
+	if (!isClassBasedSchema(schema)) {
 		// Use JSON.stringify to quote and escape identifier string.
 		throw new UsageError(
 			`Schema for ${JSON.stringify(
@@ -235,7 +243,16 @@ export function schemaAsTreeNodeValid(
 		);
 	}
 
-	return schema as typeof TreeNodeValid & TreeNodeSchema;
+	return schema;
+}
+
+/**
+ * Check if a schema is a {@link TreeNodeValid}.
+ */
+export function isClassBasedSchema(
+	schema: SimpleNodeSchemaBase<NodeKind>,
+): schema is typeof TreeNodeValid & TreeNodeSchema {
+	return inPrototypeChain(schema, TreeNodeValid);
 }
 
 /**
@@ -246,6 +263,7 @@ export function schemaAsTreeNodeValid(
 export function createTreeNodeSchemaPrivateData(
 	schema: TreeNodeSchemaCore<string, NodeKind, boolean>,
 	childAnnotatedAllowedTypes: readonly ImplicitAnnotatedAllowedTypes[],
+	toStored: TreeNodeSchemaPrivateData["toStored"],
 ): TreeNodeSchemaPrivateData {
 	const schemaValid = schemaAsTreeNodeValid(schema);
 	// Since this closes over the schema, ensure this schema is marked as most derived
@@ -255,6 +273,7 @@ export function createTreeNodeSchemaPrivateData(
 	return {
 		idempotentInitialize: () => schemaValid.oneTimeInitialize().oneTimeInitialized,
 		childAnnotatedAllowedTypes,
+		toStored,
 	};
 }
 

package/src/simple-tree/core/walkSchema.ts

@@ -5,12 +5,14 @@
 
 import {
 	normalizeAnnotatedAllowedTypes,
+	type AnnotatedAllowedType,
 	type NormalizedAnnotatedAllowedTypes,
 } from "./allowedTypes.js";
 import { getTreeNodeSchemaPrivateData, type TreeNodeSchema } from "./treeNodeSchema.js";
 
 /**
  * Traverses all {@link TreeNodeSchema} schema reachable from `schema`, applying the visitor pattern.
+ * @internal
  */
 export function walkNodeSchema(
 	schema: TreeNodeSchema,
@@ -42,31 +44,49 @@ export function walkNodeSchema(
 
 /**
  * Traverses all {@link TreeNodeSchema} schema reachable from `allowedTypes`, applying the visitor pattern.
+ * @internal
  */
 export function walkAllowedTypes(
 	annotatedAllowedTypes: NormalizedAnnotatedAllowedTypes,
 	visitor: SchemaVisitor,
 	visitedSet: Set<TreeNodeSchema> = new Set(),
 ): void {
-	for (const { type } of annotatedAllowedTypes.types) {
-		walkNodeSchema(type, visitor, visitedSet);
+	for (const allowedType of annotatedAllowedTypes.types) {
+		if ((visitor.allowedTypeFilter ?? (() => true))(allowedType)) {
+			walkNodeSchema(allowedType.type, visitor, visitedSet);
+		}
 	}
 	visitor.allowedTypes?.(annotatedAllowedTypes);
 }
 
 /**
- * Callbacks for use in {@link walkFieldSchema} / {@link walkAllowedTypes} / {@link walkNodeSchema}.
+ * Callbacks and options for use in {@link walkFieldSchema} / {@link walkAllowedTypes} / {@link walkNodeSchema}.
+ * @internal
  */
 export interface SchemaVisitor {
 	/**
-	 * Called once for each node schema.
+	 * Called once for each node schema reached.
 	 */
 	node?: (schema: TreeNodeSchema) => void;
 	/**
 	 * Called once for each set of allowed types.
-	 * Includes implicit allowed types (when a single type was used instead of an array).
+	 * @remarks
+	 * This includes every field, as well as the allowed types for maps and arrays nodes and the root if starting at {@link walkAllowedTypes}.
 	 *
-	 * This includes every field, but also the allowed types array for maps and arrays and the root if starting at {@link walkAllowedTypes}.
+	 * Each allowed types in the schema is visited as it was in the original schema except for normalization.
+	 *
+	 * After this is called {@link SchemaVisitor.allowedTypeFilter} is applied to each allowed type in the schema to determine which of them are walked into.
 	 */
 	allowedTypes?: (allowedTypes: NormalizedAnnotatedAllowedTypes) => void;
+	/**
+	 * If true, will walk into this `allowedType`.
+	 * If false, the `allowedType` will not be walked into.
+	 *
+	 * If not provided, all allowedTypes will be walked into.
+	 * @remarks
+	 * Called after {@link SchemaVisitor.allowedTypes}.
+	 * @privateRemarks
+	 * It would be possible to combine this with `allowedTypes` into a single callback, but for the current usage this separation is more convenient.
+	 */
+	allowedTypeFilter?: (allowedType: AnnotatedAllowedType<TreeNodeSchema>) => boolean;
 }

package/src/simple-tree/createContext.ts

@@ -15,19 +15,27 @@ import {
 	type TreeNodeSchemaInitializedData,
 } from "./core/index.js";
 import { normalizeFieldSchema, type ImplicitFieldSchema } from "./fieldSchema.js";
-import { toStoredSchema } from "./toStoredSchema.js";
+import { toStoredSchema, toUnhydratedSchema } from "./toStoredSchema.js";
 
 const contextCache: WeakMap<ImplicitFieldSchema, Context> = new WeakMap();
 
 /**
  * Utility for creating {@link Context}s for unhydrated nodes.
+ * @remarks
+ * The resulting context will not allow any unknown optional fields.
  */
 export function getUnhydratedContext(schema: ImplicitFieldSchema): Context {
 	return getOrCreate(contextCache, schema, (s) => {
 		const normalized = normalizeFieldSchema(schema);
 
-		const flexContext = new UnhydratedContext(defaultSchemaPolicy, toStoredSchema(schema));
-		return new Context(normalized.annotatedAllowedTypesNormalized, flexContext);
+		const flexContext = new UnhydratedContext(
+			defaultSchemaPolicy,
+			toStoredSchema(schema, toUnhydratedSchema),
+		);
+		return new Context(
+			flexContext,
+			Context.schemaMapFromRootSchema(normalized.annotatedAllowedTypesNormalized),
+		);
 	});
 }
 

package/src/simple-tree/fieldSchema.ts

@@ -468,7 +468,8 @@ export class FieldSchemaAlpha<
 			normalizeAnnotatedAllowedTypes(annotatedAllowedTypes),
 		);
 		this.lazyIdentifiers = new Lazy(
-			() => new Set([...this.allowedTypeSet].map((t) => t.identifier)),
+			() =>
+				new Set(this.annotatedAllowedTypesNormalized.types.map(({ type }) => type.identifier)),
 		);
 		this.propsAlpha = props;
 	}
@@ -488,6 +489,7 @@ export class FieldSchemaAlpha<
 
 /**
  * Normalizes a {@link ImplicitFieldSchema} or {@link ImplicitAnnotatedFieldSchema} to a {@link FieldSchema}.
+ * @internal
  */
 export function normalizeFieldSchema(
 	schema: ImplicitFieldSchema | ImplicitAnnotatedFieldSchema,