fix(tree): walk allowed types

.changeset/six-steaks-clean.md

@@ -0,0 +1,9 @@
+---
+"fluid-framework": minor
+"@fluidframework/tree": minor
+"__section": tree
+---
+Rename and change type of `annotatedAllowedTypeSet` on `FieldSchemaAlpha` to more closely align with `allowedTypesSet`
+
+This changes the `annotatedAllowedTypeSet` property on [`FieldSchemaAlpha`](https://fluidframework.com/docs/api/fluid-framework/fieldschemaalpha-class).
+It is now called `annotatedAllowedTypesNormalized` and stores evaluated schemas along with their annotations in a list of objects rather than as a mapping from the schemas to their annotations. This makes the API easier to use and better aligns with the current public APIs.

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

@@ -32,9 +32,9 @@ export interface AllowedTypesMetadata {
 export function allowUnused<T>(t?: T): void;
 
 // @alpha
-export interface AnnotatedAllowedType<T extends TreeNodeSchema = TreeNodeSchema> {
+export interface AnnotatedAllowedType<T = LazyItem<TreeNodeSchema>> {
     readonly metadata: AllowedTypeMetadata;
-    readonly type: LazyItem<T>;
+    readonly type: T;
 }
 
 // @alpha
@@ -214,7 +214,7 @@ export class FieldSchemaAlpha<Kind extends FieldKind = FieldKind, Types extends
     readonly allowedTypesMetadata: AllowedTypesMetadata;
     // (undocumented)
     readonly annotatedAllowedTypes: ImplicitAnnotatedAllowedTypes;
-    get annotatedAllowedTypeSet(): ReadonlyMap<TreeNodeSchema, AllowedTypeMetadata>;
+    get annotatedAllowedTypesNormalized(): NormalizedAnnotatedAllowedTypes;
     get persistedMetadata(): JsonCompatibleReadOnlyObject | undefined;
 }
 
@@ -618,6 +618,12 @@ export interface NodeSchemaOptionsAlpha<out TCustomMetadata = unknown> extends N
 // @alpha
 export const noopValidator: JsonValidator;
 
+// @alpha
+export interface NormalizedAnnotatedAllowedTypes {
+    readonly metadata: AllowedTypesMetadata;
+    readonly types: readonly AnnotatedAllowedType<TreeNodeSchema>[];
+}
+
 // @public @system
 export type ObjectFromSchemaRecord<T extends RestrictiveStringRecord<ImplicitFieldSchema>> = RestrictiveStringRecord<ImplicitFieldSchema> extends T ? {} : {
     -readonly [Property in keyof T]: Property extends string ? TreeFieldFromImplicitField<T[Property]> : unknown;

packages/dds/tree/src/index.ts

@@ -158,6 +158,7 @@ export {
 	type ImplicitAnnotatedFieldSchema,
 	type AnnotatedAllowedType,
 	type AnnotatedAllowedTypes,
+	type NormalizedAnnotatedAllowedTypes,
 	type AllowedTypeMetadata,
 	type AllowedTypesMetadata,
 	type InsertableObjectFromAnnotatedSchemaRecord,

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

@@ -22,7 +22,6 @@ import {
 	TreeStatus,
 } from "../feature-libraries/index.js";
 import {
-	type FieldSchema,
 	type ImplicitFieldSchema,
 	type SchemaCompatibilityStatus,
 	type TreeView,
@@ -53,6 +52,7 @@ import {
 	areImplicitFieldSchemaEqual,
 	createUnknownOptionalFieldPolicy,
 	prepareForInsertionContextless,
+	type FieldSchema,
 } from "../simple-tree/index.js";
 import {
 	type Breakable,
@@ -353,7 +353,10 @@ export class SchematizingSimpleTreeView<
 			);
 			this.checkout.forest.anchors.slots.set(
 				SimpleContextSlot,
-				new HydratedContext(this.rootFieldSchema.allowedTypeSet, view.context),
+				new HydratedContext(
+					normalizeFieldSchema(this.rootFieldSchema).annotatedAllowedTypesNormalized,
+					view.context,
+				),
 			);
 
 			const unregister = this.checkout.storedSchema.events.on("afterSchemaChange", () => {

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

@@ -9,7 +9,9 @@ import { UsageError } from "@fluidframework/telemetry-utils/internal";
 import {
 	type FieldSchemaAlpha,
 	type ImplicitFieldSchema,
+	evaluateLazySchema,
 	FieldKind,
+	isAnnotatedAllowedType,
 	markSchemaMostDerived,
 	normalizeFieldSchema,
 } from "../schemaTypes.js";
@@ -202,8 +204,12 @@ export class TreeViewConfiguration<
 				debugAssert(() => !definitions.has(schema.identifier));
 				definitions.set(schema.identifier, schema as SimpleNodeSchema & TreeNodeSchema);
 			},
-			allowedTypes(types): void {
-				checkUnion(types, config.preventAmbiguity, ambiguityErrors);
+			allowedTypes({ types }): void {
+				checkUnion(
+					types.map((t) => evaluateLazySchema(isAnnotatedAllowedType(t) ? t.type : t)),
+					config.preventAmbiguity,
+					ambiguityErrors,
+				);
 			},
 		});
 

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

@@ -10,7 +10,7 @@ import type {
 } from "../../feature-libraries/index.js";
 import { brand } from "../../util/index.js";
 
-import type { TreeNodeSchema } from "./treeNodeSchema.js";
+import type { NormalizedAnnotatedAllowedTypes, TreeNodeSchema } from "./treeNodeSchema.js";
 import { walkAllowedTypes } from "./walkSchema.js";
 
 /**
@@ -52,7 +52,7 @@ export class Context {
 	 * Since this walks the schema, it must not be invoked during schema declaration or schema forward references could fail to be resolved.
 	 */
 	public constructor(
-		rootSchema: Iterable<TreeNodeSchema>,
+		rootSchema: NormalizedAnnotatedAllowedTypes,
 		public readonly flexContext: FlexTreeContext,
 	) {
 		const schema: Map<TreeNodeSchemaIdentifier, TreeNodeSchema> = new Map();
@@ -71,7 +71,7 @@ export class Context {
  */
 export class HydratedContext extends Context {
 	public constructor(
-		rootSchema: Iterable<TreeNodeSchema>,
+		rootSchema: NormalizedAnnotatedAllowedTypes,
 		public override readonly flexContext: FlexTreeHydratedContext,
 	) {
 		super(rootSchema, flexContext);

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

@@ -31,6 +31,9 @@ export {
 	type TreeNodeSchemaNonClass,
 	type TreeNodeSchemaCore,
 	type TreeNodeSchemaBoth,
+	type AnnotatedAllowedType,
+	type NormalizedAnnotatedAllowedTypes,
+	isAnnotatedAllowedTypes,
 } from "./treeNodeSchema.js";
 export { walkAllowedTypes, type SchemaVisitor } from "./walkSchema.js";
 export { Context, HydratedContext, SimpleContextSlot } from "./context.js";

packages/dds/tree/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,
+		"All implementations of TreeNodeSchemaCore must also implement TreeNodeSchemaCorePrivate",
+	);
+	assert(
+		Array.isArray((schema as TreeNodeSchemaCorePrivate).childAnnotatedAllowedTypes),
+		"All implementations of TreeNodeSchemaCore must also implement TreeNodeSchemaCorePrivate",
+	);
+	return schema as TreeNodeSchemaCorePrivate;
+}
+
 /**
  * Kind of tree node.
  * @remarks

packages/dds/tree/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;
 }

packages/dds/tree/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);
 	});
 }

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

@@ -8,6 +8,8 @@ export {
 	typeSchemaSymbol,
 	type WithType,
 	type TreeNodeSchema,
+	type AnnotatedAllowedType,
+	type NormalizedAnnotatedAllowedTypes,
 	NodeKind,
 	type TreeNodeSchemaClass,
 	type TreeNodeSchemaNonClass,
@@ -138,8 +140,6 @@ export type {
 export {
 	type NodeFromSchema,
 	isTreeNodeSchemaClass,
-	type AnnotatedAllowedType,
-	type AnnotatedAllowedTypes,
 	type ImplicitFieldSchema,
 	type ImplicitAnnotatedFieldSchema,
 	type TreeFieldFromImplicitField,
@@ -158,6 +158,7 @@ export {
 	type AllowedTypes,
 	type AllowedTypeMetadata,
 	type AllowedTypesMetadata,
+	type AnnotatedAllowedTypes,
 	FieldKind,
 	FieldSchema,
 	type FieldSchemaAlpha,

packages/dds/tree/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)) {