@fluidframework/tree 2.72.0 → 2.73.0

package/src/simple-tree/index.ts

@@ -44,7 +44,7 @@ export {
 	type UnannotateAllowedTypesList,
 	type AllowedTypeMetadata,
 	type AnnotatedAllowedTypes,
-	type SchemaUpgrade,
+	SchemaUpgrade,
 	type LazyItem,
 	type FlexList,
 	type FlexListToUnion,
@@ -59,6 +59,8 @@ export {
 	type AllowedTypesFullFromMixed,
 	AnnotatedAllowedTypesInternal,
 	type NumberKeys,
+	ExpectStored,
+	createSchemaUpgrade,
 } from "./core/index.js";
 export { walkFieldSchema } from "./walkFieldSchema.js";
 export type { UnsafeUnknownSchema, Insertable } from "./unsafeUnknownSchema.js";
@@ -283,4 +285,7 @@ export type { LeafSchema } from "./leafNodeSchema.js";
 export {
 	encodeSimpleSchema,
 	decodeSimpleSchema,
+	exportCompatibilitySchemaSnapshot,
+	importCompatibilitySchemaSnapshot,
+	checkCompatibility,
 } from "./api/index.js";

package/src/simple-tree/leafNodeSchema.ts

@@ -7,7 +7,7 @@ import { assert } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 import { isFluidHandle } from "@fluidframework/runtime-utils/internal";
 
-import { LeafNodeStoredSchema, type TreeValue, ValueSchema } from "../core/index.js";
+import { type TreeValue, ValueSchema } from "../core/index.js";
 import {
 	type FlexTreeNode,
 	isFlexTreeNode,
@@ -64,7 +64,6 @@ export class LeafNodeSchema<Name extends string, const T extends ValueSchema>
 				): FlexContent => leafToFlexContent(data, this, allowedTypes),
 			})),
 		childAllowedTypes: [],
-		toStored: () => new LeafNodeStoredSchema(this.leafKind),
 	};
 	#initializedData: TreeNodeSchemaInitializedData | undefined;
 

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

@@ -45,7 +45,6 @@ import {
 	createTreeNodeSchemaPrivateData,
 	type FlexContent,
 	type TreeNodeSchemaPrivateData,
-	convertAllowedTypes,
 	withBufferedTreeEvents,
 	AnnotatedAllowedTypesInternal,
 } from "../../core/index.js";
@@ -66,9 +65,13 @@ import type {
 	ArrayNodePojoEmulationSchema,
 	ArrayNodeSchema,
 } from "./arrayNodeTypes.js";
-import { brand, type JsonCompatibleReadOnlyObject } from "../../../util/index.js";
+import {
+	brand,
+	validateIndex,
+	validateIndexRange,
+	type JsonCompatibleReadOnlyObject,
+} from "../../../util/index.js";
 import { nullSchema } from "../../leafNodeSchema.js";
-import { arrayNodeStoredSchema } from "../../toStoredSchema.js";
 import type { SimpleAllowedTypeAttributes } from "../../simpleSchema.js";
 
 /**
@@ -951,7 +954,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 	}
 	public insertAt(index: number, ...value: Insertable<T>): void {
 		const field = getSequenceField(this);
-		validateIndex(index, field, "insertAt", true);
+		validateIndex(index, field, "TreeArrayNode.insertAt", true);
 		const content = this.#mapTreesFromFieldData(value);
 		field.editor.insert(index, content);
 	}
@@ -963,32 +966,30 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 	}
 	public removeAt(index: number): void {
 		const field = getSequenceField(this);
-		validateIndex(index, field, "removeAt");
+		validateIndex(index, field, "TreeArrayNode.removeAt");
 		field.editor.remove(index, 1);
 	}
 	public removeRange(start?: number, end?: number): void {
 		const field = getSequenceField(this);
 		const { length } = field;
 		const removeStart = start ?? 0;
+		validateIndex(removeStart, field, "TreeArrayNode.removeRange", true);
+
 		const removeEnd = Math.min(length, end ?? length);
-		validatePositiveIndex(removeStart);
-		validatePositiveIndex(removeEnd);
-		if (removeEnd < removeStart) {
-			// This catches both the case where start is > array.length and when start is > end.
-			throw new UsageError('Too large of "start" value passed to TreeArrayNode.removeRange.');
-		}
+		validateIndexRange(removeStart, removeEnd, field, "TreeArrayNode.removeRange");
+
 		field.editor.remove(removeStart, removeEnd - removeStart);
 	}
 	public moveToStart(sourceIndex: number, source?: ReadonlyArrayNode): void {
 		const sourceArray = source ?? this;
 		const sourceField = getSequenceField(sourceArray);
-		validateIndex(sourceIndex, sourceField, "moveToStart");
+		validateIndex(sourceIndex, sourceField, "TreeArrayNode.moveToStart");
 		this.moveRangeToIndex(0, sourceIndex, sourceIndex + 1, source);
 	}
 	public moveToEnd(sourceIndex: number, source?: ReadonlyArrayNode): void {
 		const sourceArray = source ?? this;
 		const sourceField = getSequenceField(sourceArray);
-		validateIndex(sourceIndex, sourceField, "moveToEnd");
+		validateIndex(sourceIndex, sourceField, "TreeArrayNode.moveToEnd");
 		this.moveRangeToIndex(this.length, sourceIndex, sourceIndex + 1, source);
 	}
 	public moveToIndex(
@@ -999,8 +1000,8 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		const sourceArray = source ?? this;
 		const sourceField = getSequenceField(sourceArray);
 		const destinationField = getSequenceField(this);
-		validateIndex(destinationGap, destinationField, "moveToIndex", true);
-		validateIndex(sourceIndex, sourceField, "moveToIndex");
+		validateIndex(destinationGap, destinationField, "TreeArrayNode.moveToIndex", true);
+		validateIndex(sourceIndex, sourceField, "TreeArrayNode.moveToIndex");
 		this.moveRangeToIndex(destinationGap, sourceIndex, sourceIndex + 1, source);
 	}
 	public moveRangeToStart(
@@ -1012,7 +1013,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 			sourceStart,
 			sourceEnd,
 			source ?? getSequenceField(this),
-			"moveRangeToStart",
+			"TreeArrayNode.moveRangeToStart",
 		);
 		this.moveRangeToIndex(0, sourceStart, sourceEnd, source);
 	}
@@ -1025,7 +1026,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 			sourceStart,
 			sourceEnd,
 			source ?? getSequenceField(this),
-			"moveRangeToEnd",
+			"TreeArrayNode.moveRangeToEnd",
 		);
 		this.moveRangeToIndex(this.length, sourceStart, sourceEnd, source);
 	}
@@ -1044,8 +1045,13 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		).getFieldSchema(EmptyKey).types;
 		const sourceField = source !== undefined ? getSequenceField(source) : destinationField;
 
-		validateIndex(destinationGap, destinationField, "moveRangeToIndex", true);
-		validateIndexRange(sourceStart, sourceEnd, source ?? destinationField, "moveRangeToIndex");
+		validateIndex(destinationGap, destinationField, "TreeArrayNode.moveRangeToIndex", true);
+		validateIndexRange(
+			sourceStart,
+			sourceEnd,
+			source ?? destinationField,
+			"TreeArrayNode.moveRangeToIndex",
+		);
 
 		// TODO: determine support for move across different sequence types
 		if (sourceField !== destinationField) {
@@ -1282,12 +1288,7 @@ export function arraySchema<
 		}
 
 		public static get [privateDataSymbol](): TreeNodeSchemaPrivateData {
-			return (privateData ??= createTreeNodeSchemaPrivateData(
-				this,
-				[normalizedTypes],
-				(storedOptions) =>
-					arrayNodeStoredSchema(convertAllowedTypes(info, storedOptions), persistedMetadata),
-			));
+			return (privateData ??= createTreeNodeSchemaPrivateData(this, [normalizedTypes]));
 		}
 	}
 
@@ -1295,56 +1296,6 @@ export function arraySchema<
 	return output;
 }
 
-function validateSafeInteger(index: number): void {
-	if (!Number.isSafeInteger(index)) {
-		throw new UsageError(`Expected a safe integer, got ${index}.`);
-	}
-}
-
-function validatePositiveIndex(index: number): void {
-	validateSafeInteger(index);
-	if (index < 0) {
-		throw new UsageError(`Expected non-negative index, got ${index}.`);
-	}
-}
-
-function validateIndex(
-	index: number,
-	array: { readonly length: number },
-	methodName: string,
-	allowOnePastEnd: boolean = false,
-): void {
-	validatePositiveIndex(index);
-	if (allowOnePastEnd) {
-		if (index > array.length) {
-			throw new UsageError(
-				`Index value passed to TreeArrayNode.${methodName} is out of bounds.`,
-			);
-		}
-	} else {
-		if (index >= array.length) {
-			throw new UsageError(
-				`Index value passed to TreeArrayNode.${methodName} is out of bounds.`,
-			);
-		}
-	}
-}
-
-function validateIndexRange(
-	startIndex: number,
-	endIndex: number,
-	array: { readonly length: number },
-	methodName: string,
-): void {
-	validateIndex(startIndex, array, methodName, true);
-	validateIndex(endIndex, array, methodName, true);
-	if (startIndex > endIndex || array.length < endIndex) {
-		throw new UsageError(
-			`Index value passed to TreeArrayNode.${methodName} is out of bounds.`,
-		);
-	}
-}
-
 /**
  * Transforms data for a child of an array.
  * @param child - The tree data to be transformed.

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

@@ -7,7 +7,6 @@ import { assert, Lazy } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import {
-	FieldKinds,
 	isTreeValue,
 	type FlexibleNodeContent,
 	type FlexTreeNode,
@@ -42,7 +41,6 @@ import {
 	createTreeNodeSchemaPrivateData,
 	type FlexContent,
 	type TreeNodeSchemaPrivateData,
-	convertAllowedTypes,
 	AnnotatedAllowedTypesInternal,
 } from "../../core/index.js";
 import {
@@ -343,19 +341,7 @@ export function mapSchema<
 		}
 
 		public static get [privateDataSymbol](): TreeNodeSchemaPrivateData {
-			return (privateData ??= createTreeNodeSchemaPrivateData(
-				this,
-				[normalizedTypes],
-				(storedOptions) =>
-					new MapNodeStoredSchema(
-						{
-							kind: FieldKinds.optional.identifier,
-							types: convertAllowedTypes(info, storedOptions),
-							persistedMetadata,
-						},
-						persistedMetadata,
-					),
-			));
+			return (privateData ??= createTreeNodeSchemaPrivateData(this, [normalizedTypes]));
 		}
 	}
 	const schemaErased: MapNodeCustomizableSchema<

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

@@ -86,7 +86,7 @@ import {
 	type FactoryContentObject,
 	type InsertableContent,
 } from "../../unhydratedFlexTreeFromInsertable.js";
-import { convertField, convertFieldKind } from "../../toStoredSchema.js";
+import { convertFieldKind } from "../../toStoredSchema.js";
 import type { ObjectSchemaOptionsAlpha } from "../../api/index.js";
 
 /**
@@ -585,20 +585,6 @@ export function objectSchema<
 			return (privateData ??= createTreeNodeSchemaPrivateData(
 				this,
 				Array.from(CustomObjectNode.fields.values(), (schema) => schema.allowedTypesFull),
-				(storedOptions) => {
-					const fields: Map<FieldKey, TreeFieldStoredSchema> = new Map();
-					for (const fieldSchema of flexKeyMap.values()) {
-						assert(
-							fieldSchema.schema instanceof FieldSchemaAlpha,
-							0xc19 /* Expected FieldSchemaAlpha */,
-						);
-						fields.set(
-							brand(fieldSchema.storedKey),
-							convertField(fieldSchema.schema, storedOptions),
-						);
-					}
-					return new ObjectNodeStoredSchema(fields, nodeOptions.persistedMetadata);
-				},
 			));
 		}
 	}
@@ -650,9 +636,8 @@ function assertUniqueKeys<
 }
 
 /**
- * {@link TreeNodeSchemaInitializedData.toFlexContent} for Map nodes.
+ * {@link TreeNodeSchemaInitializedData.toFlexContent} for object nodes.
  *
- * Transforms data under an Object schema.
  * @param data - The tree data to be transformed. Must be a Record-like object.
  * @param schema - The schema to comply with.
  */

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

@@ -31,7 +31,6 @@ import {
 	type FlexContent,
 	CompatibilityLevel,
 	type TreeNodeSchemaPrivateData,
-	convertAllowedTypes,
 	AnnotatedAllowedTypesInternal,
 } from "../../core/index.js";
 import { getTreeNodeSchemaInitializedData } from "../../createContext.js";
@@ -50,7 +49,6 @@ import type {
 	TreeRecordNode,
 } from "./recordNodeTypes.js";
 import {
-	FieldKinds,
 	isTreeValue,
 	type FlexTreeNode,
 	type FlexTreeOptionalField,
@@ -388,19 +386,7 @@ export function recordSchema<
 		}
 
 		public static get [privateDataSymbol](): TreeNodeSchemaPrivateData {
-			return (privateData ??= createTreeNodeSchemaPrivateData(
-				this,
-				[normalizedTypes],
-				(storedOptions) =>
-					new MapNodeStoredSchema(
-						{
-							kind: FieldKinds.optional.identifier,
-							types: convertAllowedTypes(info, storedOptions),
-							persistedMetadata,
-						},
-						persistedMetadata,
-					),
-			));
+			return (privateData ??= createTreeNodeSchemaPrivateData(this, [normalizedTypes]));
 		}
 	}
 

package/src/simple-tree/simpleSchema.ts

@@ -5,7 +5,7 @@
 
 import type { ValueSchema } from "../core/index.js";
 import type { JsonCompatibleReadOnlyObject } from "../util/index.js";
-import type { NodeKind, SimpleNodeSchemaBase } from "./core/index.js";
+import type { NodeKind, SchemaUpgrade, SimpleNodeSchemaBase } from "./core/index.js";
 import type { FieldKind, FieldSchemaMetadata } from "./fieldSchema.js";
 
 /*
@@ -39,6 +39,15 @@ export interface SimpleNodeSchemaBaseAlpha<
 	readonly persistedMetadata: JsonCompatibleReadOnlyObject | undefined;
 }
 
+/**
+ * {@link AllowedTypes} for a location in the tree, expressed for use in the Simple Schema layer of abstraction.
+ *
+ * @remarks
+ * Refers to the types by identifier.
+ * A {@link SimpleTreeSchema} is needed to resolve these identifiers to their schema {@link SimpleTreeSchema.definitions}.
+ */
+export type SimpleAllowedTypes = ReadonlyMap<string, SimpleAllowedTypeAttributes>;
+
 /**
  * A {@link SimpleNodeSchema} for an object node.
  *
@@ -179,13 +188,17 @@ export type SimpleNodeSchema =
  */
 export interface SimpleAllowedTypeAttributes {
 	/**
-	 * True if this schema is included as a {@link SchemaStaticsBeta.staged | staged} schema upgrade,
+	 * {@link SchemaUpgrade} if this schema is included as a {@link SchemaStaticsBeta.staged | staged} schema upgrade,
 	 * allowing the view schema be compatible with stored schema with (post upgrade) or without it (pre-upgrade).
 	 * New documents and schema upgrades will omit any staged schema.
 	 *
 	 * Undefined if derived from a stored schema.
+	 *
+	 * @privateRemarks
+	 * The false and undefined cases here are a bit odd.
+	 * This API should be reevaluated before stabilizing.
 	 */
-	readonly isStaged: boolean | undefined;
+	readonly isStaged: false | SchemaUpgrade | undefined;
 }
 
 /**

package/src/simple-tree/toStoredSchema.ts

@@ -3,7 +3,11 @@
  * Licensed under the MIT License.
  */
 
-import { unreachableCase, fail } from "@fluidframework/core-utils/internal";
+import {
+	unreachableCase,
+	fail,
+	transformMapValues,
+} from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import {
@@ -17,7 +21,6 @@ import {
 	type TreeNodeSchemaIdentifier,
 	type TreeNodeStoredSchema,
 	type TreeStoredSchema,
-	type TreeTypeSet,
 } from "../core/index.js";
 import { FieldKinds, type FlexFieldKind } from "../feature-libraries/index.js";
 import { brand, getOrCreate, type JsonCompatibleReadOnlyObject } from "../util/index.js";
@@ -25,23 +28,23 @@ import { brand, getOrCreate, type JsonCompatibleReadOnlyObject } from "../util/i
 import {
 	allowedTypeFilter,
 	convertAllowedTypes,
-	getTreeNodeSchemaPrivateData,
-	isClassBasedSchema,
+	ExpectStored,
 	NodeKind,
 	type SimpleNodeSchemaBase,
+	type StoredFromViewSchemaGenerationOptions,
 	type StoredSchemaGenerationOptions,
 } from "./core/index.js";
-import {
-	FieldKind,
-	FieldSchemaAlpha,
-	normalizeFieldSchema,
-	type ImplicitFieldSchema,
-} from "./fieldSchema.js";
-import type { SimpleFieldSchema, SimpleNodeSchema } from "./simpleSchema.js";
+import { FieldKind, normalizeFieldSchema, type ImplicitFieldSchema } from "./fieldSchema.js";
+import type {
+	SimpleAllowedTypes,
+	SimpleFieldSchema,
+	SimpleNodeSchema,
+	SimpleTreeSchema,
+} from "./simpleSchema.js";
 import { walkFieldSchema } from "./walkFieldSchema.js";
 
 const viewToStoredCache = new WeakMap<
-	StoredSchemaGenerationOptions,
+	StoredFromViewSchemaGenerationOptions,
 	WeakMap<ImplicitFieldSchema, TreeStoredSchema>
 >();
 
@@ -85,6 +88,9 @@ export const toUnhydratedSchema = permissiveStoredSchemaGenerationOptions;
  * TODO:#38722 When runtime schema upgrades are implemented, this will need to be updated to check if
  * a staged allowed type has been upgraded and if so, include it in the conversion.
  *
+ * Even if this took in a SimpleTreeSchema,
+ * it would still need to walk the schema to avoid including schema that become unreachable due to filtered out staged schema.
+ *
  * @throws
  * Throws a `UsageError` if multiple schemas are encountered with the same identifier.
  */
@@ -114,7 +120,11 @@ export function toStoredSchema(
 					),
 				);
 			},
-			allowedTypeFilter: (allowedType) => allowedTypeFilter(allowedType, options),
+			allowedTypeFilter: (allowedType) =>
+				allowedTypeFilter(
+					{ isStaged: allowedType.metadata.stagedSchemaUpgrade ?? false },
+					options,
+				),
 		});
 
 		const result: TreeStoredSchema = {
@@ -126,24 +136,35 @@ export function toStoredSchema(
 }
 
 /**
- * Normalizes an {@link ImplicitFieldSchema} into a {@link TreeFieldSchema}.
+ * Convert a {@link SimpleTreeSchema} for a stored schema into a {@link TreeStoredSchema}.
+ * @remarks
+ * This only supports simple schemas that are already logically stored schemas.
+ * @privateRemarks
+ * To correctly support view schema here, this would need to filter out unreferenced schema after excluding staged schema.
+ * @see {@link ExpectStored}.
+ */
+export function simpleStoredSchemaToStoredSchema(
+	treeSchema: SimpleTreeSchema,
+): TreeStoredSchema {
+	const result: TreeStoredSchema = {
+		nodeSchema: transformMapValues(treeSchema.definitions, (schema) =>
+			getStoredSchema(schema, ExpectStored),
+		) as Map<TreeNodeSchemaIdentifier, TreeNodeStoredSchema>,
+		rootFieldSchema: convertField(treeSchema.root, ExpectStored),
+	};
+	return result;
+}
+
+/**
+ * Convert a {@link SimpleFieldSchema} into a {@link TreeFieldStoredSchema}.
  */
 export function convertField(
-	schema: SimpleFieldSchema | FieldSchemaAlpha,
+	schema: SimpleFieldSchema,
 	options: StoredSchemaGenerationOptions,
 ): TreeFieldStoredSchema {
 	const kind: FieldKindIdentifier =
 		convertFieldKind.get(schema.kind)?.identifier ?? fail(0xae3 /* Invalid field kind */);
-	let types: TreeTypeSet;
-	// eslint-disable-next-line unicorn/prefer-ternary
-	if (schema instanceof FieldSchemaAlpha) {
-		types = convertAllowedTypes(schema.allowedTypes, options);
-	} else {
-		const allowedTypesIdentifiers: ReadonlySet<string> = new Set(
-			schema.simpleAllowedTypes.keys(),
-		);
-		types = allowedTypesIdentifiers as TreeTypeSet;
-	}
+	const types = convertAllowedTypes(schema.simpleAllowedTypes, options);
 	return { kind, types, persistedMetadata: schema.persistedMetadata };
 }
 
@@ -168,9 +189,6 @@ export function getStoredSchema(
 	schema: SimpleNodeSchema,
 	options: StoredSchemaGenerationOptions,
 ): TreeNodeStoredSchema {
-	if (isClassBasedSchema(schema)) {
-		return getTreeNodeSchemaPrivateData(schema).toStored(options);
-	}
 	const kind = schema.kind;
 	switch (kind) {
 		case NodeKind.Leaf: {
@@ -178,10 +196,7 @@ export function getStoredSchema(
 		}
 		case NodeKind.Map:
 		case NodeKind.Record: {
-			const allowedTypesIdentifiers: ReadonlySet<string> = new Set(
-				schema.simpleAllowedTypes.keys(),
-			);
-			const types = allowedTypesIdentifiers as TreeTypeSet;
+			const types = convertAllowedTypes(schema.simpleAllowedTypes, options);
 			return new MapNodeStoredSchema(
 				{
 					kind: FieldKinds.optional.identifier,
@@ -193,11 +208,11 @@ export function getStoredSchema(
 			);
 		}
 		case NodeKind.Array: {
-			const allowedTypesIdentifiers: ReadonlySet<string> = new Set(
-				schema.simpleAllowedTypes.keys(),
+			return arrayNodeStoredSchema(
+				schema.simpleAllowedTypes,
+				options,
+				schema.persistedMetadata,
 			);
-			const types = allowedTypesIdentifiers as TreeTypeSet;
-			return arrayNodeStoredSchema(types, schema.persistedMetadata);
 		}
 		case NodeKind.Object: {
 			const fields: Map<FieldKey, TreeFieldStoredSchema> = new Map();
@@ -213,12 +228,13 @@ export function getStoredSchema(
 }
 
 export function arrayNodeStoredSchema(
-	types: TreeTypeSet,
+	schema: SimpleAllowedTypes,
+	options: StoredSchemaGenerationOptions,
 	persistedMetadata: JsonCompatibleReadOnlyObject | undefined,
 ): ObjectNodeStoredSchema {
 	const field = {
 		kind: FieldKinds.sequence.identifier,
-		types,
+		types: convertAllowedTypes(schema, options),
 		persistedMetadata,
 	};
 	const fields = new Map([[EmptyKey, field]]);

package/src/simple-tree/unhydratedFlexTreeFromInsertable.ts

@@ -5,9 +5,9 @@
 
 import type { IFluidHandle } from "@fluidframework/core-interfaces";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
-import { assert } from "@fluidframework/core-utils/internal";
+import { assert, fail } from "@fluidframework/core-utils/internal";
 
-import { hasSingle } from "../util/index.js";
+import { filterIterable, hasSingle, oneFromIterable } from "../util/index.js";
 
 import { normalizeFieldSchema, FieldKind, type ImplicitFieldSchema } from "./fieldSchema.js";
 import {
@@ -96,7 +96,13 @@ export function unhydratedFlexTreeFromInsertableNode(
 	const handler = getTreeNodeSchemaPrivateData(schema).idempotentInitialize();
 	const result = handler.toFlexContent(data, allowedTypes);
 
-	return new UnhydratedFlexTreeNode(...result, getUnhydratedContext(schema));
+	// Might not match schema due to fallbacks, see TODO on toFlexContent
+	// TODO: fix TODO in `toFlexContent`, and remove this.
+	const finalSchema =
+		oneFromIterable(filterIterable(allowedTypes, (s) => s.identifier === result[0].type)) ??
+		fail(0xc9d /* missing schema */);
+
+	return new UnhydratedFlexTreeNode(...result, getUnhydratedContext(finalSchema));
 }
 
 function getType(
@@ -133,21 +139,20 @@ export function getPossibleTypes(
 	data: FactoryContent,
 ): TreeNodeSchema[] {
 	assert(data !== undefined, 0x889 /* undefined cannot be used as FactoryContent. */);
-	const type =
-		typeof data === "object" && data !== null
-			? (data as Partial<{ [contentSchemaSymbol]: string }>)[contentSchemaSymbol]
-			: undefined;
 
-	let best = CompatibilityLevel.None;
+	let toCheck: Iterable<TreeNodeSchema>;
+	if (typeof data === "object" && data !== null && contentSchemaSymbol in data) {
+		// If the data has an explicit brand via contentSchemaSymbol, only check that type.
+		const type = data[contentSchemaSymbol];
+		toCheck = filterIterable(allowedTypes, (schema) => schema.identifier === type);
+	} else {
+		toCheck = allowedTypes;
+	}
+
+	// Start at the lowest level of compat we would ever accept: this discards types which are less compatible.
+	let best = CompatibilityLevel.Low;
 	const possibleTypes: TreeNodeSchema[] = [];
-	for (const schema of allowedTypes) {
-		if (type !== undefined) {
-			if (schema.identifier === type) {
-				return [schema];
-			} else {
-				continue;
-			}
-		}
+	for (const schema of toCheck) {
 		const handler = getTreeNodeSchemaPrivateData(schema).idempotentInitialize();
 		const level = handler.shallowCompatibilityTest(data);
 		if (level > best) {
@@ -158,7 +163,7 @@ export function getPossibleTypes(
 			possibleTypes.push(schema);
 		}
 	}
-	return best === CompatibilityLevel.None ? [] : possibleTypes;
+	return possibleTypes;
 }
 
 /**