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,15 @@ export interface NodeSchemaOptionsAlpha<out TCustomMetadata = unknown> extends N
 // @alpha
 export const noopValidator: JsonValidator;
 
+// @alpha
+export function normalizeAllowedTypes(types: ImplicitAllowedTypes): ReadonlySet<TreeNodeSchema>;
+
+// @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,14 @@ 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) =>
+						isAnnotatedAllowedType(t) ? evaluateLazySchema(t.type) : evaluateLazySchema(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,14 @@
  * Licensed under the MIT License.
  */
 
-import type { TreeLeafValue } from "../schemaTypes.js";
+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 +69,47 @@ 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 (
+		typeof allowedTypes === "object" && "metadata" in allowedTypes && "types" in allowedTypes
+	);
+}
+
 /**
  * Schema which is not a class.
  * @remarks
@@ -231,8 +279,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 +369,71 @@ 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 direct child of a node with this schema could have, grouped by field.
+	 * If this node does not have fields, it will contain a single array with all its allowed types.
+	 *
+	 * Equivalently, this is also all schema directly referenced when defining this schema's allowed child types,
+	 * which is also the same as the set of schema referenced directly by the `Info` type parameter and the `info` property.
+	 * This property is simply re-exposing that information in an easier to traverse format consistent across all node kinds.
+	 * @remarks
+	 * Some kinds of nodes may have additional restrictions on children:
+	 * this set simply enumerates all directly referenced schema, and can be use to walk over all referenced schema types.
+	 *
+	 * 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 set.
+	 *
+	 * @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 {
+	if (
+		"childAnnotatedAllowedTypes" in schema &&
+		Array.isArray(schema.childAnnotatedAllowedTypes) &&
+		(schema.childAnnotatedAllowedTypes.length === 0 ||
+			isAnnotatedAllowedTypes(
+				schema.childAnnotatedAllowedTypes[0] as ImplicitAnnotatedAllowedTypes,
+			))
+	) {
+		return schema as TreeNodeSchemaCorePrivate;
+	}
+	throw new Error(
+		"All implementations of TreeNodeSchemaCore must also implement TreeNodeSchemaCorePrivate",
+	);
+}
+
 /**
  * Kind of tree node.
  * @remarks

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

@@ -3,57 +3,73 @@
  * Licensed under the MIT License.
  */
 
-import type { TreeNodeSchema } from "./treeNodeSchema.js";
+import { fail } from "@fluidframework/core-utils/internal";
+import type { AllowedTypeMetadata } from "../schemaTypes.js";
+import {
+	asTreeNodeSchemaCorePrivate,
+	type NormalizedAnnotatedAllowedTypes,
+	type TreeNodeSchema,
+} from "./treeNodeSchema.js";
 
 /**
  * Traverses all {@link TreeNodeSchema} schema reachable from `schema`, applying the visitor pattern.
  */
 export function walkNodeSchema(
 	schema: TreeNodeSchema,
+	annotations: AllowedTypeMetadata,
 	visitor: SchemaVisitor,
 	visitedSet: Set<TreeNodeSchema>,
 ): void {
 	if (visitedSet.has(schema)) {
 		return;
 	}
+
 	visitedSet.add(schema);
 
-	walkAllowedTypes(schema.childTypes, visitor, visitedSet);
+	const annotatedAllowedTypes =
+		asTreeNodeSchemaCorePrivate(schema).childAnnotatedAllowedTypes ??
+		fail("TreeNodeSchemas must implement TreeNodeSchemaCorePrivate");
+
+	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,
 	// working from the inner types out to the types that use them will probably go better than the reverse.
 	// This does not however ensure all types referenced by a type are visited before it, since in recursive cases thats impossible.
-	visitor.node?.(schema);
+	visitor.node?.(schema, annotations);
 }
 
 /**
  * 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 annotatedAllowedType of annotatedAllowedTypes.types) {
+		const { type, metadata } = annotatedAllowedType;
+		walkNodeSchema(type, metadata, visitor, visitedSet);
 	}
-	visitor.allowedTypes?.(allowedTypes);
+	visitor.allowedTypes?.(annotatedAllowedTypes);
 }
 
 /**
  * Callbacks for use in {@link walkFieldSchema} / {@link walkAllowedTypes} / {@link walkNodeSchema}.
  */
 export interface SchemaVisitor {
 	/**
-	 * Called once for each node schema.
+	 * Called for each node schema. This may be called multiple times for the same node schema e.g. if the same schema
+	 * is allowed on different fields.
 	 */
-	node?: (schema: TreeNodeSchema) => void;
+	node?: (schema: TreeNodeSchema, annotations: AllowedTypeMetadata) => void;
 	/**
 	 * Called once for each set of allowed types.
 	 * Includes implicit allowed types (when a single type was used instead of an array).
 	 *
 	 * 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);
 	});
 }