@fluidframework/tree 2.43.0-343119 → 2.50.0-345060

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

@@ -3,7 +3,15 @@
  * Licensed under the MIT License.
  */
 
-import type { TreeLeafValue } from "../schemaTypes.js";
+import { assert } from "@fluidframework/core-utils/internal";
+import type { LazyItem } from "../flexList.js";
+import type {
+	AllowedTypeMetadata,
+	AllowedTypesMetadata,
+	AnnotatedAllowedTypes,
+	ImplicitAnnotatedAllowedTypes,
+	TreeLeafValue,
+} from "../schemaTypes.js";
 import type { SimpleNodeSchemaBase } from "../simpleSchema.js";
 
 import type { TreeNode } from "./treeNode.js";
@@ -62,6 +70,48 @@ export type TreeNodeSchema<
 			TCustomMetadata
 	  >;
 
+/**
+ * Stores annotations for an individual allowed type.
+ * @alpha
+ */
+export interface AnnotatedAllowedType<T = LazyItem<TreeNodeSchema>> {
+	/**
+	 * Annotations for the allowed type.
+	 */
+	readonly metadata: AllowedTypeMetadata;
+	/**
+	 * The allowed type the annotations apply to in a particular schema.
+	 */
+	readonly type: T;
+}
+
+/**
+ * Stores annotations for a set of evaluated annotated allowed types.
+ * @alpha
+ */
+export interface NormalizedAnnotatedAllowedTypes {
+	/**
+	 * Annotations that apply to a set of allowed types.
+	 */
+	readonly metadata: AllowedTypesMetadata;
+	/**
+	 * All the evaluated allowed types that the annotations apply to. The types themselves are also individually annotated.
+	 */
+	readonly types: readonly AnnotatedAllowedType<TreeNodeSchema>[];
+}
+
+/**
+ * Checks if the input is an {@link AnnotatedAllowedTypes}.
+ */
+export function isAnnotatedAllowedTypes(
+	allowedTypes: ImplicitAnnotatedAllowedTypes,
+): allowedTypes is AnnotatedAllowedTypes {
+	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
+	);
+}
+
 /**
  * Schema which is not a class.
  * @remarks
@@ -231,8 +281,13 @@ export type TreeNodeSchemaBoth<
 
 /**
  * Data common to all tree node schema.
+ *
  * @remarks
  * Implementation detail of {@link TreeNodeSchema} which should be accessed instead of referring to this type directly.
+ *
+ * @privateRemarks
+ * All implementations must implement {@link TreeNodeSchemaCorePrivate} as well.
+ *
  * @sealed @public
  */
 export interface TreeNodeSchemaCore<
@@ -316,6 +371,63 @@ export interface TreeNodeSchemaCore<
 	createFromInsertable(data: TInsertable): Unhydrated<TreeNode | TreeLeafValue>;
 }
 
+/**
+ * {@link TreeNodeSchemaCore} extended with some non-exported APIs.
+ */
+export interface TreeNodeSchemaCorePrivate<
+	Name extends string = string,
+	Kind extends NodeKind = NodeKind,
+	TInsertable = never,
+	ImplicitlyConstructable extends boolean = boolean,
+	Info = unknown,
+	TCustomMetadata = unknown,
+> extends TreeNodeSchemaCore<
+		Name,
+		Kind,
+		ImplicitlyConstructable,
+		Info,
+		TInsertable,
+		TCustomMetadata
+	> {
+	/**
+	 * All possible annotated allowed types that a field under a node with this schema could have.
+	 * @remarks
+	 * In this case "field" includes anything that is a field in the internal (flex-tree) abstraction layer.
+	 * This includes the content field for arrays, and all the fields for map nodes.
+	 * If this node does not have fields (and thus is a leaf), the array will be empty.
+	 *
+	 * This set cannot be used before the schema in it have been defined:
+	 * more specifically, when using lazy schema references (for example to make foreword references to schema which have not yet been defined),
+	 * users must wait until after the schema are defined to access this array.
+	 *
+	 * @privateRemarks
+	 * If this is stabilized, it will live alongside the childTypes property on {@link TreeNodeSchemaCore}.
+	 * @system
+	 */
+	readonly childAnnotatedAllowedTypes: readonly NormalizedAnnotatedAllowedTypes[];
+}
+
+/**
+ * Downcasts a {@link TreeNodeSchemaCore} to {@link TreeNodeSchemaCorePrivate} if it is one.
+ *
+ * @remarks
+ * This function should only be used internally. The result should not be exposed publicly
+ * in any exported types or API return values.
+ */
+export function asTreeNodeSchemaCorePrivate(
+	schema: TreeNodeSchemaCore<string, NodeKind, boolean>,
+): TreeNodeSchemaCorePrivate {
+	assert(
+		"childAnnotatedAllowedTypes" in schema,
+		0xbc9 /* All implementations of TreeNodeSchemaCore must also implement TreeNodeSchemaCorePrivate */,
+	);
+	assert(
+		Array.isArray((schema as TreeNodeSchemaCorePrivate).childAnnotatedAllowedTypes),
+		0xbca /* All implementations of TreeNodeSchemaCore must also implement TreeNodeSchemaCorePrivate */,
+	);
+	return schema as TreeNodeSchemaCorePrivate;
+}
+
 /**
  * Kind of tree node.
  * @remarks
@@ -326,19 +438,23 @@ export enum NodeKind {
 	/**
 	 * A node which serves as a map, storing children under string keys.
 	 */
-	Map,
+	Map = 0,
 	/**
 	 * A node which serves as an array, storing children in an ordered sequence.
 	 */
-	Array,
+	Array = 1,
 	/**
 	 * A node which stores a heterogenous collection of children in named fields.
 	 * @remarks
 	 * Each field gets its own schema.
 	 */
-	Object,
+	Object = 2,
 	/**
 	 * A node which stores a single leaf value.
 	 */
-	Leaf,
+	Leaf = 3,
+	/**
+	 * A node which serves as a record, storing children under string keys.
+	 */
+	Record = 4,
 }

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

@@ -3,7 +3,11 @@
  * Licensed under the MIT License.
  */
 
-import type { TreeNodeSchema } from "./treeNodeSchema.js";
+import {
+	asTreeNodeSchemaCorePrivate,
+	type NormalizedAnnotatedAllowedTypes,
+	type TreeNodeSchema,
+} from "./treeNodeSchema.js";
 
 /**
  * Traverses all {@link TreeNodeSchema} schema reachable from `schema`, applying the visitor pattern.
@@ -16,9 +20,14 @@ export function walkNodeSchema(
 	if (visitedSet.has(schema)) {
 		return;
 	}
+
 	visitedSet.add(schema);
 
-	walkAllowedTypes(schema.childTypes, visitor, visitedSet);
+	const annotatedAllowedTypes = asTreeNodeSchemaCorePrivate(schema).childAnnotatedAllowedTypes;
+
+	for (const fieldAllowedTypes of annotatedAllowedTypes) {
+		walkAllowedTypes(fieldAllowedTypes, visitor, visitedSet);
+	}
 
 	// This visit is done at the end so the traversal order is most inner types first.
 	// This was picked since when fixing errors,
@@ -31,14 +40,14 @@ export function walkNodeSchema(
  * Traverses all {@link TreeNodeSchema} schema reachable from `allowedTypes`, applying the visitor pattern.
  */
 export function walkAllowedTypes(
-	allowedTypes: Iterable<TreeNodeSchema>,
+	annotatedAllowedTypes: NormalizedAnnotatedAllowedTypes,
 	visitor: SchemaVisitor,
 	visitedSet: Set<TreeNodeSchema> = new Set(),
 ): void {
-	for (const childType of allowedTypes) {
-		walkNodeSchema(childType, visitor, visitedSet);
+	for (const { type } of annotatedAllowedTypes.types) {
+		walkNodeSchema(type, visitor, visitedSet);
 	}
-	visitor.allowedTypes?.(allowedTypes);
+	visitor.allowedTypes?.(annotatedAllowedTypes);
 }
 
 /**
@@ -55,5 +64,5 @@ export interface SchemaVisitor {
 	 *
 	 * This includes every field, but also the allowed types array for maps and arrays and the root if starting at {@link walkAllowedTypes}.
 	 */
-	allowedTypes?: (allowedTypes: Iterable<TreeNodeSchema>) => void;
+	allowedTypes?: (allowedTypes: NormalizedAnnotatedAllowedTypes) => void;
 }

package/src/simple-tree/createContext.ts

@@ -20,6 +20,6 @@ export function getUnhydratedContext(schema: ImplicitFieldSchema): Context {
 		const normalized = normalizeFieldSchema(schema);
 
 		const flexContext = new UnhydratedContext(defaultSchemaPolicy, toStoredSchema(schema));
-		return new Context(normalized.allowedTypeSet, flexContext);
+		return new Context(normalized.annotatedAllowedTypesNormalized, flexContext);
 	});
 }

package/src/simple-tree/getTreeNodeForField.ts

@@ -19,7 +19,9 @@ import { type TreeNode, getOrCreateNodeFromInnerNode } from "./core/index.js";
 /**
  * Retrieve the associated {@link TreeNode} for the given field's content.
  */
-export function getTreeNodeForField(field: FlexTreeField): TreeNode | TreeValue | undefined {
+export function tryGetTreeNodeForField(
+	field: FlexTreeField,
+): TreeNode | TreeValue | undefined {
 	function tryToUnboxLeaves(
 		flexField: FlexTreeOptionalField | FlexTreeRequiredField,
 	): TreeNode | TreeValue | undefined {

package/src/simple-tree/index.ts

@@ -8,6 +8,8 @@ export {
 	typeSchemaSymbol,
 	type WithType,
 	type TreeNodeSchema,
+	type AnnotatedAllowedType,
+	type NormalizedAnnotatedAllowedTypes,
 	NodeKind,
 	type TreeNodeSchemaClass,
 	type TreeNodeSchemaNonClass,
@@ -23,6 +25,7 @@ export {
 	SimpleContextSlot,
 	getOrCreateInnerNode,
 	getOrCreateNodeFromInnerNode,
+	getOrCreateNodeFromInnerUnboxedNode,
 	getKernel,
 } from "./core/index.js";
 export {
@@ -66,6 +69,8 @@ export {
 	type JsonTreeSchema,
 	type JsonFieldSchema,
 	type JsonLeafSchemaType,
+	type JsonRecordNodeSchema,
+	type JsonStringKeyPatternProperties,
 	type TreeSchemaEncodingOptions,
 	getJsonSchema,
 	getSimpleSchema,
@@ -80,12 +85,13 @@ export {
 	type System_Unsafe,
 	type ArrayNodeCustomizableSchemaUnsafe,
 	type MapNodeCustomizableSchemaUnsafe,
+	type TreeRecordNodeUnsafe,
 	type TreeViewAlpha,
 	type TreeBranch,
 	type TreeBranchEvents,
-	tryGetSchema,
-	getStoredKey,
 	getPropertyKeyFromStoredKey,
+	getStoredKey,
+	tryGetSchema,
 	applySchemaToParserOptions,
 	cursorFromVerbose,
 	verboseFromCursor,
@@ -133,12 +139,11 @@ export type {
 	SimpleNodeSchemaBase,
 	SimpleNodeSchemaBaseAlpha,
 	SimpleObjectFieldSchema,
+	SimpleRecordNodeSchema,
 } from "./simpleSchema.js";
 export {
 	type NodeFromSchema,
 	isTreeNodeSchemaClass,
-	type AnnotatedAllowedType,
-	type AnnotatedAllowedTypes,
 	type ImplicitFieldSchema,
 	type ImplicitAnnotatedFieldSchema,
 	type TreeFieldFromImplicitField,
@@ -157,6 +162,7 @@ export {
 	type AllowedTypes,
 	type AllowedTypeMetadata,
 	type AllowedTypesMetadata,
+	type AnnotatedAllowedTypes,
 	FieldKind,
 	FieldSchema,
 	type FieldSchemaAlpha,
@@ -185,11 +191,12 @@ export {
 	type NodeSchemaMetadata,
 	evaluateLazySchema,
 } from "./schemaTypes.js";
-export { getTreeNodeForField } from "./getTreeNodeForField.js";
+export { tryGetTreeNodeForField } from "./getTreeNodeForField.js";
 export {
 	type ArrayNodeCustomizableSchema,
 	type ArrayNodePojoEmulationSchema,
 	ArrayNodeSchema,
+	asIndex,
 	isArrayNodeSchema,
 	IterableTreeArrayContent,
 	type ReadonlyArrayNode,
@@ -204,10 +211,18 @@ export {
 	type InsertableObjectFromSchemaRecord,
 	type ObjectFromSchemaRecord,
 	ObjectNodeSchema,
+	type ObjectNodeSchemaPrivate,
+	isObjectNodeSchema,
 	type InsertableObjectFromAnnotatedSchemaRecord,
 	type TreeObjectNode,
 	setField,
 	createUnknownOptionalFieldPolicy,
+	isRecordNodeSchema,
+	type RecordNodeCustomizableSchema,
+	type RecordNodeInsertableData,
+	type RecordNodePojoEmulationSchema,
+	RecordNodeSchema,
+	type TreeRecordNode,
 } from "./node-kinds/index.js";
 export {
 	unhydratedFlexTreeFromInsertable,

package/src/simple-tree/leafNodeSchema.ts

@@ -12,7 +12,12 @@ import {
 	valueSchemaAllows,
 } from "../feature-libraries/index.js";
 
-import { NodeKind, type TreeNodeSchema, type TreeNodeSchemaNonClass } from "./core/index.js";
+import {
+	NodeKind,
+	type NormalizedAnnotatedAllowedTypes,
+	type TreeNodeSchema,
+	type TreeNodeSchemaNonClass,
+} from "./core/index.js";
 import type { NodeSchemaMetadata, TreeLeafValue } from "./schemaTypes.js";
 import type { SimpleLeafNodeSchema } from "./simpleSchema.js";
 import type { JsonCompatibleReadOnlyObject } from "../util/index.js";
@@ -34,6 +39,7 @@ export class LeafNodeSchema<Name extends string, const T extends ValueSchema>
 	public readonly info: T;
 	public readonly implicitlyConstructable = true as const;
 	public readonly childTypes: ReadonlySet<TreeNodeSchema> = new Set();
+	public readonly childAnnotatedAllowedTypes: readonly NormalizedAnnotatedAllowedTypes[] = [];
 
 	public create(data: TreeValue<T> | FlexTreeNode): TreeValue<T> {
 		if (isFlexTreeNode(data)) {

package/src/simple-tree/node-kinds/array/arrayNode.ts

@@ -7,14 +7,14 @@ import { Lazy, oob, fail } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import { EmptyKey } from "../../../core/index.js";
-import {
-	type FlexibleFieldContent,
-	type FlexTreeNode,
-	type FlexTreeSequenceField,
-	isFlexTreeNode,
+import type {
+	FlexibleFieldContent,
+	FlexTreeNode,
+	FlexTreeSequenceField,
 } from "../../../feature-libraries/index.js";
 import {
 	normalizeAllowedTypes,
+	normalizeAnnotatedAllowedTypes,
 	unannotateImplicitAllowedTypes,
 	type ImplicitAllowedTypes,
 	type ImplicitAnnotatedAllowedTypes,
@@ -41,6 +41,8 @@ import {
 	getKernel,
 	type UnhydratedFlexTreeNode,
 	UnhydratedSequenceField,
+	getOrCreateNodeFromInnerUnboxedNode,
+	type NormalizedAnnotatedAllowedTypes,
 } from "../../core/index.js";
 import {
 	type InsertableContent,
@@ -735,9 +737,9 @@ function createArrayNodeProxy(
 			}
 
 			const maybeContent = field.at(maybeIndex);
-			return isFlexTreeNode(maybeContent)
-				? getOrCreateNodeFromInnerNode(maybeContent)
-				: maybeContent;
+			return maybeContent === undefined
+				? undefined
+				: getOrCreateNodeFromInnerUnboxedNode(maybeContent);
 		},
 		set: (target, key, newValue, receiver) => {
 			if (key === "length") {
@@ -796,7 +798,7 @@ function createArrayNodeProxy(
 				// To satisfy 'deepEquals' level scrutiny, the property descriptor for indexed properties must
 				// be a simple value property (as opposed to using getter) and declared writable/enumerable/configurable.
 				return {
-					value: isFlexTreeNode(val) ? getOrCreateNodeFromInnerNode(val) : val,
+					value: val === undefined ? undefined : getOrCreateNodeFromInnerUnboxedNode(val),
 					writable: true, // For MVP, setting indexed properties is reported as allowed here (for deep equals compatibility noted above), but not actually supported.
 					enumerable: true,
 					configurable: true,
@@ -1110,6 +1112,7 @@ export function arraySchema<
 	const unannotatedTypes = unannotateImplicitAllowedTypes(info);
 
 	const lazyChildTypes = new Lazy(() => normalizeAllowedTypes(unannotatedTypes));
+	const lazyAnnotatedTypes = new Lazy(() => [normalizeAnnotatedAllowedTypes(info)]);
 	const lazyAllowedTypesIdentifiers = new Lazy(
 		() => new Set([...lazyChildTypes.value].map((type) => type.identifier)),
 	);
@@ -1192,6 +1195,9 @@ export function arraySchema<
 		public static get childTypes(): ReadonlySet<TreeNodeSchema> {
 			return lazyChildTypes.value;
 		}
+		public static get childAnnotatedAllowedTypes(): readonly NormalizedAnnotatedAllowedTypes[] {
+			return lazyAnnotatedTypes.value;
+		}
 		public static readonly metadata: NodeSchemaMetadata<TCustomMetadata> = metadata ?? {};
 		public static readonly persistedMetadata: JsonCompatibleReadOnlyObject | undefined =
 			persistedMetadata;

package/src/simple-tree/node-kinds/index.ts

@@ -33,7 +33,18 @@ export {
 	isObjectNodeSchema,
 	type ObjectFromSchemaRecord,
 	ObjectNodeSchema,
+	type ObjectNodeSchemaPrivate,
 	objectSchema,
 	setField,
 	type TreeObjectNode,
 } from "./object/index.js";
+
+export {
+	isRecordNodeSchema,
+	type RecordNodeCustomizableSchema,
+	type RecordNodeInsertableData,
+	type RecordNodePojoEmulationSchema,
+	RecordNodeSchema,
+	recordSchema,
+	type TreeRecordNode,
+} from "./record/index.js";

package/src/simple-tree/node-kinds/map/mapNode.ts

@@ -10,11 +10,12 @@ import type {
 	FlexTreeOptionalField,
 	OptionalFieldEditBuilder,
 } from "../../../feature-libraries/index.js";
-import { getTreeNodeForField } from "../../getTreeNodeForField.js";
+import { tryGetTreeNodeForField } from "../../getTreeNodeForField.js";
 import {
 	createFieldSchema,
 	FieldKind,
 	normalizeAllowedTypes,
+	normalizeAnnotatedAllowedTypes,
 	unannotateImplicitAllowedTypes,
 	type ImplicitAllowedTypes,
 	type ImplicitAnnotatedAllowedTypes,
@@ -36,6 +37,7 @@ import {
 	getOrCreateInnerNode,
 	type InternalTreeNode,
 	type UnhydratedFlexTreeNode,
+	type NormalizedAnnotatedAllowedTypes,
 } from "../../core/index.js";
 import {
 	unhydratedFlexTreeFromInsertable,
@@ -177,14 +179,14 @@ abstract class CustomMapNodeBase<const T extends ImplicitAllowedTypes> extends T
 		for (const key of node.keys()) {
 			yield [
 				key,
-				getTreeNodeForField(node.getBoxed(key)) as TreeNodeFromImplicitAllowedTypes<T>,
+				tryGetTreeNodeForField(node.getBoxed(key)) as TreeNodeFromImplicitAllowedTypes<T>,
 			];
 		}
 	}
 	public get(key: string): TreeNodeFromImplicitAllowedTypes<T> {
 		const node = this.innerNode;
 		const field = node.getBoxed(brand(key));
-		return getTreeNodeForField(field) as TreeNodeFromImplicitAllowedTypes<T>;
+		return tryGetTreeNodeForField(field) as TreeNodeFromImplicitAllowedTypes<T>;
 	}
 	public has(key: string): boolean {
 		return this.innerNode.tryGetField(brand(key)) !== undefined;
@@ -221,7 +223,7 @@ abstract class CustomMapNodeBase<const T extends ImplicitAllowedTypes> extends T
 		thisArg?: unknown,
 	): void {
 		for (const field of getOrCreateInnerNode(this)) {
-			const node = getTreeNodeForField(field) as TreeNodeFromImplicitAllowedTypes<T>;
+			const node = tryGetTreeNodeForField(field) as TreeNodeFromImplicitAllowedTypes<T>;
 			callbackFn.call(thisArg, node, field.key, this);
 		}
 	}
@@ -252,6 +254,7 @@ export function mapSchema<
 	const lazyChildTypes = new Lazy(() =>
 		normalizeAllowedTypes(unannotateImplicitAllowedTypes(info)),
 	);
+	const lazyAnnotatedTypes = new Lazy(() => [normalizeAnnotatedAllowedTypes(info)]);
 	const lazyAllowedTypesIdentifiers = new Lazy(
 		() => new Set([...lazyChildTypes.value].map((type) => type.identifier)),
 	);
@@ -300,6 +303,9 @@ export function mapSchema<
 		public static get childTypes(): ReadonlySet<TreeNodeSchema> {
 			return lazyChildTypes.value;
 		}
+		public static get childAnnotatedAllowedTypes(): readonly NormalizedAnnotatedAllowedTypes[] {
+			return lazyAnnotatedTypes.value;
+		}
 		public static readonly metadata: NodeSchemaMetadata<TCustomMetadata> = metadata ?? {};
 		public static readonly persistedMetadata: JsonCompatibleReadOnlyObject | undefined =
 			persistedMetadata;

package/src/simple-tree/node-kinds/object/index.ts

@@ -13,4 +13,8 @@ export {
 	setField,
 	type TreeObjectNode,
 } from "./objectNode.js";
-export { isObjectNodeSchema, ObjectNodeSchema } from "./objectNodeTypes.js";
+export {
+	isObjectNodeSchema,
+	ObjectNodeSchema,
+	type ObjectNodeSchemaPrivate,
+} from "./objectNodeTypes.js";

package/src/simple-tree/node-kinds/object/objectNode.ts

@@ -32,13 +32,15 @@ import {
 	type Context,
 	type UnhydratedFlexTreeNode,
 	getOrCreateInnerNode,
+	type NormalizedAnnotatedAllowedTypes,
 } from "../../core/index.js";
 import { getUnhydratedContext } from "../../createContext.js";
-import { getTreeNodeForField } from "../../getTreeNodeForField.js";
+import { tryGetTreeNodeForField } from "../../getTreeNodeForField.js";
 import {
 	isObjectNodeSchema,
 	type ObjectNodeSchema,
 	type ObjectNodeSchemaInternalData,
+	type ObjectNodeSchemaPrivate,
 } from "./objectNodeTypes.js";
 import { prepareForInsertion } from "../../prepareForInsertion.js";
 import {
@@ -54,7 +56,6 @@ import {
 	type FieldSchemaAlpha,
 	ObjectFieldSchema,
 	type ImplicitAnnotatedFieldSchema,
-	unannotateSchemaRecord,
 	type UnannotateSchemaRecord,
 } from "../../schemaTypes.js";
 import type { SimpleObjectFieldSchema } from "../../simpleSchema.js";
@@ -194,7 +195,9 @@ export type SimpleKeyMap = ReadonlyMap<
 /**
  * Caches the mappings from property keys to stored keys for the provided object field schemas in {@link simpleKeyToFlexKeyCache}.
  */
-function createFlexKeyMapping(fields: Record<string, ImplicitFieldSchema>): SimpleKeyMap {
+function createFlexKeyMapping(
+	fields: Record<string, ImplicitAnnotatedFieldSchema>,
+): SimpleKeyMap {
 	const keyMap: Map<string | symbol, { storedKey: FieldKey; schema: FieldSchema }> = new Map();
 	for (const [propertyKey, fieldSchema] of Object.entries(fields)) {
 		const storedKey = getStoredKey(propertyKey, fieldSchema);
@@ -214,7 +217,7 @@ function createFlexKeyMapping(fields: Record<string, ImplicitFieldSchema>): Simp
  * If not provided `{}` is used for the target.
  */
 function createProxyHandler(
-	schema: ObjectNodeSchema & ObjectNodeSchemaInternalData,
+	schema: ObjectNodeSchemaPrivate,
 	allowAdditionalProperties: boolean,
 ): ProxyHandler<TreeNode> {
 	// To satisfy 'deepEquals' level scrutiny, the target of the proxy must be an object with the same
@@ -234,13 +237,13 @@ function createProxyHandler(
 				debugAssert(() => !flexNode.context.isDisposed() || "FlexTreeNode is disposed");
 				const field = flexNode.tryGetField(fieldInfo.storedKey);
 				if (field !== undefined) {
-					return getTreeNodeForField(field);
+					return tryGetTreeNodeForField(field);
 				}
 
 				return undefined;
 			}
 
-			// POJO mode objects don't have TreeNode's build in members on their targets, so special case them:
+			// POJO mode objects don't have TreeNode's built in members on their targets, so special case them:
 			if (propertyKey === typeSchemaSymbol) {
 				return schema;
 			}
@@ -302,7 +305,7 @@ function createProxyHandler(
 			const field = getOrCreateInnerNode(proxy).tryGetField(fieldInfo.storedKey);
 
 			const p: PropertyDescriptor = {
-				value: field === undefined ? undefined : getTreeNodeForField(field),
+				value: field === undefined ? undefined : tryGetTreeNodeForField(field),
 				writable: true,
 				// Report empty fields as own properties so they shadow inherited properties (even when empty) to match TypeScript typing.
 				// Make empty fields not enumerable so they get skipped when iterating over an object to better align with
@@ -372,14 +375,12 @@ export function objectSchema<
 	// Field set can't be modified after this since derived data is stored in maps.
 	Object.freeze(info);
 
-	const unannotatedInfo = unannotateSchemaRecord(info);
-
 	// Ensure no collisions between final set of property keys, and final set of stored keys (including those
 	// implicitly derived from property keys)
-	assertUniqueKeys(identifier, unannotatedInfo);
+	assertUniqueKeys(identifier, info);
 
 	// Performance optimization: cache property key => stored key and schema.
-	const flexKeyMap: SimpleKeyMap = createFlexKeyMapping(unannotatedInfo);
+	const flexKeyMap: SimpleKeyMap = createFlexKeyMapping(info);
 
 	const identifierFieldKeys: FieldKey[] = [];
 	for (const item of flexKeyMap.values()) {
@@ -391,6 +392,12 @@ export function objectSchema<
 	const lazyChildTypes = new Lazy(
 		() => new Set(Array.from(flexKeyMap.values(), (f) => [...f.schema.allowedTypeSet]).flat()),
 	);
+	const lazyAnnotatedTypes = new Lazy(() =>
+		Array.from(
+			flexKeyMap.values(),
+			({ schema }) => normalizeFieldSchema(schema).annotatedAllowedTypesNormalized,
+		),
+	);
 
 	let handler: ProxyHandler<object>;
 	let customizable: boolean;
@@ -470,7 +477,7 @@ export function objectSchema<
 		protected static override oneTimeSetup<T2>(this: typeof TreeNodeValid<T2>): Context {
 			// One time initialization that required knowing the most derived type (from this.constructor) and thus has to be lazy.
 			customizable = (this as unknown) !== CustomObjectNode;
-			const schema = this as unknown as ObjectNodeSchema & ObjectNodeSchemaInternalData;
+			const schema = this as unknown as ObjectNodeSchemaPrivate;
 			handler = createProxyHandler(schema, customizable);
 			unhydratedContext = getUnhydratedContext(schema);
 
@@ -509,6 +516,9 @@ export function objectSchema<
 		public static get childTypes(): ReadonlySet<TreeNodeSchema> {
 			return lazyChildTypes.value;
 		}
+		public static get childAnnotatedAllowedTypes(): readonly NormalizedAnnotatedAllowedTypes[] {
+			return lazyAnnotatedTypes.value;
+		}
 		public static readonly metadata: NodeSchemaMetadata<TCustomMetadata> = metadata ?? {};
 		public static readonly persistedMetadata: JsonCompatibleReadOnlyObject | undefined =
 			persistedMetadata;
@@ -537,7 +547,7 @@ const targetToProxy: WeakMap<object, TreeNode> = new WeakMap();
  */
 function assertUniqueKeys<
 	const Name extends number | string,
-	const Fields extends RestrictiveStringRecord<ImplicitFieldSchema>,
+	const Fields extends RestrictiveStringRecord<ImplicitAnnotatedFieldSchema>,
 >(schemaName: Name, fields: Fields): void {
 	// Verify that there are no duplicates among the explicitly specified stored keys.
 	const explicitStoredKeys = new Set<string>();

package/src/simple-tree/node-kinds/object/objectNodeTypes.ts

@@ -83,8 +83,11 @@ export const ObjectNodeSchema = {
 	},
 } as const;
 
-export function isObjectNodeSchema(
-	schema: TreeNodeSchema,
-): schema is ObjectNodeSchema & ObjectNodeSchemaInternalData {
+/**
+ * {@link ObjectNodeSchema} with data that is not part of the package-exported API surface.
+ */
+export type ObjectNodeSchemaPrivate = ObjectNodeSchema & ObjectNodeSchemaInternalData;
+
+export function isObjectNodeSchema(schema: TreeNodeSchema): schema is ObjectNodeSchemaPrivate {
 	return schema.kind === NodeKind.Object;
 }

package/src/simple-tree/node-kinds/record/index.ts

@@ -0,0 +1,14 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+export { recordSchema } from "./recordNode.js";
+export {
+	isRecordNodeSchema,
+	type RecordNodeCustomizableSchema,
+	type RecordNodeInsertableData,
+	type RecordNodePojoEmulationSchema,
+	RecordNodeSchema,
+	type TreeRecordNode,
+} from "./recordNodeTypes.js";