Add ensureSchema

Author: noencke

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

@@ -965,6 +965,9 @@ export interface SchemaStaticsAlpha {
     readonly typesRecursive: <const T extends readonly Unenforced<AnnotatedAllowedType | LazyItem<TreeNodeSchema>>[]>(t: T, metadata?: AllowedTypesMetadata) => AllowedTypesFullFromMixedUnsafe<T>;
 }
 
+// @alpha
+export const schemaSymbol: unique symbol;
+
 // @alpha @sealed
 export class SchemaUpgrade {
     // (undocumented)
@@ -1348,6 +1351,7 @@ export interface TreeAlpha {
     child(node: TreeNode, key: string | number): TreeNode | TreeLeafValue | undefined;
     children(node: TreeNode): Iterable<[propertyKey: string | number, child: TreeNode | TreeLeafValue]>;
     create<const TSchema extends ImplicitFieldSchema | UnsafeUnknownSchema>(schema: UnsafeUnknownSchema extends TSchema ? ImplicitFieldSchema : TSchema & ImplicitFieldSchema, data: InsertableField<TSchema>): Unhydrated<TSchema extends ImplicitFieldSchema ? TreeFieldFromImplicitField<TSchema> : TreeNode | TreeLeafValue | undefined>;
+    ensureSchema<TSchema extends TreeNodeSchema, TContent extends InsertableField<TSchema>>(schema: TSchema, content: TContent): TContent;
     exportCompressed(tree: TreeNode | TreeLeafValue, options: {
         idCompressor?: IIdCompressor;
     } & Pick<CodecWriteOptions, "minVersionForCollab">): JsonCompatible<IFluidHandle>;

packages/dds/tree/src/index.ts

@@ -123,6 +123,7 @@ export {
 	type WithType,
 	type NodeChangedData,
 	type SchemaUpgrade,
+	schemaSymbol,
 	// Types not really intended for public use, but used in links.
 	// Can not be moved to internalTypes since doing so causes app code to throw errors like:
 	// Error: src/simple-tree/objectNode.ts:72:1 - (ae-unresolved-link) The @link reference could not be resolved: The package "@fluidframework/tree" does not have an export "TreeNodeApi"

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

@@ -58,6 +58,8 @@ import {
 	importConcise,
 	exportConcise,
 	borrowCursorFromTreeNodeOrValue,
+	schemaSymbol,
+	type TreeNodeSchema,
 } from "../simple-tree/index.js";
 import { brand, extractFromOpaque, type JsonCompatible } from "../util/index.js";
 import {
@@ -519,6 +521,31 @@ export interface TreeAlpha {
 		onInvalidation: () => void,
 		trackDuring: () => TResult,
 	): ObservationResults<TResult>;
+
+	/**
+	 * Ensures that the provided content will be interpreted as the given schema when inserting into the tree.
+	 * @returns `content`, for convenience.
+	 * @remarks
+	 * If applicable, this will tag the given content with a {@link schemaSymbol | special property} that indicates its intended schema.
+	 * The `content` will be interpreted as the given `schema` when later inserted into the tree.
+	 * This is particularly useful when the content's schema cannot be inferred from its structure alone because it is compatible with multiple schemas.
+	 * @example
+	 * ```typescript
+	 * const sf = new SchemaFactory("example");
+	 * class Dog extends sf.object("Dog", { name: sf.string() }) {}
+	 * class Cat extends sf.object("Cat", { name: sf.string() }) {}
+	 * class Root extends sf.object("Root", { pet: [Dog, Cat] }) {}
+	 * // ...
+	 * const pet = { name: "Max" };
+	 * view.root.pet = pet; // Error: ambiguous schema - is it a Dog or a Cat?
+	 * TreeAlpha.ensureSchema(Dog, pet); // Tags `pet` as a Dog.
+	 * view.root.pet = pet; // No error - it's a Dog.
+	 * ```
+	 */
+	ensureSchema<TSchema extends TreeNodeSchema, TContent extends InsertableField<TSchema>>(
+		schema: TSchema,
+		content: TContent,
+	): TContent;
 }
 
 /**
@@ -1003,6 +1030,21 @@ export const TreeAlpha: TreeAlpha = {
 		}
 		return result;
 	},
+
+	ensureSchema<TSchema extends TreeNodeSchema, TNode extends InsertableField<TSchema>>(
+		schema: TSchema,
+		node: TNode,
+	): TNode {
+		if (typeof node === "object" && node !== null) {
+			Reflect.defineProperty(node, schemaSymbol, {
+				configurable: false,
+				enumerable: false,
+				writable: true,
+				value: schema.identifier,
+			});
+		}
+		return node;
+	},
 };
 
 /**

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

@@ -16,7 +16,12 @@ export {
 	SimpleContextSlot,
 	withBufferedTreeEvents,
 } from "./treeNodeKernel.js";
-export { type WithType, typeNameSymbol, typeSchemaSymbol } from "./withType.js";
+export {
+	type WithType,
+	typeNameSymbol,
+	typeSchemaSymbol,
+	schemaSymbol,
+} from "./withType.js";
 export {
 	type Unhydrated,
 	type InternalTreeNode,

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

@@ -5,6 +5,9 @@
 
 import type { TreeNode } from "./treeNode.js";
 import type { NodeKind, TreeNodeSchemaClass } from "./treeNodeSchema.js";
+// Used by doc links:
+// eslint-disable-next-line @typescript-eslint/no-unused-vars, unused-imports/no-unused-imports
+import type { TreeAlpha } from "../../shared-tree/index.js";
 
 /**
  * The type of a {@link TreeNode}.
@@ -39,6 +42,27 @@ export const typeNameSymbol: unique symbol = Symbol("TreeNode Type");
  */
 export const typeSchemaSymbol: unique symbol = Symbol("TreeNode Schema");
 
+/**
+ * The intended type of insertable content that is to become a {@link TreeNode}.
+ * @remarks **Note:** Whenever possible, use the type-safe {@link (TreeAlpha:interface).ensureSchema} function rather than manually employing this symbol.
+ *
+ * If a property with this symbol key is present on an object that is inserted into the tree,
+ * the tree will use the schema identifier specified by the value of this property when creating the node.
+ * This is particularly useful for specifying the intended schema of untyped content when it would otherwise be ambiguous.
+ * @example
+ * ```typescript
+ * const sf = new SchemaFactory("example");
+ * class Dog extends sf.object("Dog", { name: sf.string() }) {}
+ * class Cat extends sf.object("Cat", { name: sf.string() }) {}
+ * class Root extends sf.object("Root", { pet: [Dog, Cat] }) {}
+ * // ...
+ * view.root.pet = { name: "Max" }; // Error: ambiguous schema - is it a Dog or a Cat?
+ * view.root.pet = { name: "Max", [schemaSymbol]: "example.Dog" }; // No error - it's a Dog.
+ * ```
+ * @alpha
+ */
+export const schemaSymbol: unique symbol = Symbol("SharedTree Schema");
+
 /**
  * Adds a type symbol to a type for stronger typing.
  *

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

@@ -6,6 +6,7 @@
 export {
 	typeNameSymbol,
 	typeSchemaSymbol,
+	schemaSymbol,
 	type WithType,
 	type TreeNodeSchema,
 	type AnnotatedAllowedType,

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

@@ -17,6 +17,7 @@ import {
 	isTreeNode,
 	type TreeNode,
 	type TreeNodeSchema,
+	schemaSymbol,
 	type Unhydrated,
 	UnhydratedFlexTreeNode,
 } from "./core/index.js";
@@ -125,16 +126,24 @@ For class-based schema, this can be done by replacing an expression like "{foo:
 
 /**
  * Returns all types for which the data is schema-compatible.
+ * @remarks This will respect the {@link schemaSymbol} property on data to disambiguate types - if present, only that type will be returned.
  */
 export function getPossibleTypes(
 	allowedTypes: ReadonlySet<TreeNodeSchema>,
 	data: FactoryContent,
 ): TreeNodeSchema[] {
 	assert(data !== undefined, 0x889 /* undefined cannot be used as FactoryContent. */);
+	const type =
+		typeof data === "object" && data !== null
+			? (data as Partial<{ [schemaSymbol]: string }>)[schemaSymbol]
+			: undefined;
 
 	let best = CompatibilityLevel.None;
 	const possibleTypes: TreeNodeSchema[] = [];
 	for (const schema of allowedTypes) {
+		if (type !== undefined && schema.identifier === type) {
+			return [schema];
+		}
 		const handler = getTreeNodeSchemaPrivateData(schema).idempotentInitialize();
 		const level = handler.shallowCompatibilityTest(data);
 		if (level > best) {

packages/dds/tree/src/test/simple-tree/api/treeNodeApi.spec.ts

@@ -3544,6 +3544,138 @@ describe("treeNodeApi", () => {
 			}
 		});
 	});
+
+	describe("ensureType", () => {
+		const sf = new SchemaFactory("test");
+		class Son extends sf.object("Son", { value: sf.number }) {}
+		class Daughter extends sf.object("Daughter", { value: sf.number }) {}
+		class Parent extends sf.object("Parent", {
+			child: sf.optional([Son, Daughter]),
+		}) {}
+		it("returns the same value that was passed in", () => {
+			const child = { value: 3 };
+			const son = TreeAlpha.ensureSchema(Son, child);
+			assert.equal(son, child);
+		});
+
+		it("allows leaf types", () => {
+			const nullValue = null;
+			assert.equal(TreeAlpha.ensureSchema(schema.null, nullValue), nullValue);
+			const booleanValue = true;
+			assert.equal(TreeAlpha.ensureSchema(schema.boolean, booleanValue), booleanValue);
+			const numberValue = 3;
+			assert.equal(TreeAlpha.ensureSchema(schema.number, numberValue), numberValue);
+			const stringValue = "hello";
+			assert.equal(TreeAlpha.ensureSchema(schema.string, stringValue), stringValue);
+		});
+
+		it("tags an object that is otherwise ambiguous", () => {
+			const child = { value: 3 };
+			// `child` could be either a Son or a Daughter, so we can't disambiguate.
+			assert.throws(
+				() => {
+					hydrate(Parent, { child });
+				},
+				validateUsageError(/compatible with more than one type/),
+			);
+			// If we explicitly tag it as a Daughter, it is thereafter interpreted as such.
+			const daughter = TreeAlpha.ensureSchema(Daughter, child);
+			const parent = hydrate(Parent, { child: daughter });
+			assert(Tree.is(parent.child, Daughter));
+		});
+
+		it("tags an array that is otherwise ambiguous", () => {
+			class Sons extends sf.array("Sons", Son) {}
+			class Daughters extends sf.array("Daughters", Daughter) {}
+			class ArrayParent extends sf.object("Parent", {
+				children: sf.optional([Sons, Daughters]),
+			}) {}
+			const children = [{ value: 3 }, { value: 4 }];
+			assert.throws(
+				() => {
+					hydrate(ArrayParent, { children });
+				},
+				validateUsageError(/compatible with more than one type/),
+			);
+			const daughters = TreeAlpha.ensureSchema(Daughters, children);
+			const parent = hydrate(ArrayParent, { children: daughters });
+			assert(Tree.is(parent.children, Daughters));
+		});
+
+		it("tags a map that is otherwise ambiguous", () => {
+			class SonMap extends sf.map("SonMap", Son) {}
+			class DaughterMap extends sf.map("DaughterMap", Daughter) {}
+			class MapParent extends sf.object("Parent", {
+				children: sf.optional([SonMap, DaughterMap]),
+			}) {}
+
+			const children = {
+				a: { value: 3 },
+				b: { value: 4 },
+			};
+			assert.throws(
+				() => {
+					hydrate(MapParent, { children });
+				},
+				validateUsageError(/compatible with more than one type/),
+			);
+			const daughterMap = TreeAlpha.ensureSchema(DaughterMap, children);
+			const parent = hydrate(MapParent, { children: daughterMap });
+			assert(Tree.is(parent.children, DaughterMap));
+		});
+
+		it("can re-tag an object that has already been tagged", () => {
+			const child = { value: 3 };
+			const daughter = TreeAlpha.ensureSchema(Daughter, child);
+			const son = TreeAlpha.ensureSchema(Son, daughter);
+			const parent = hydrate(Parent, { child: son });
+			assert(Tree.is(parent.child, Son));
+		});
+
+		it("can be used to disambiguate deep trees", () => {
+			class Father extends sf.object("Father", {
+				child: sf.optional([Son, Daughter]),
+			}) {}
+			class Mother extends sf.object("Mother", {
+				child: sf.optional([Son, Daughter]),
+			}) {}
+			class GrandParent extends sf.object("GrandParent", {
+				parent: sf.optional([Father, Mother]),
+			}) {}
+			// Ambiguous parent and child
+			assert.throws(() => {
+				hydrate(GrandParent, {
+					parent: {
+						child: { value: 3 },
+					},
+				});
+			});
+			// Tagged parent, but ambiguous child
+			assert.throws(() => {
+				hydrate(GrandParent, {
+					parent: {
+						child: TreeAlpha.ensureSchema(Son, { value: 3 }),
+					},
+				});
+			});
+			// Ambiguous parent, but tagged child
+			assert.throws(() => {
+				hydrate(GrandParent, {
+					parent: TreeAlpha.ensureSchema(Father, {
+						child: { value: 3 },
+					}),
+				});
+			});
+			// Both parent and child tagged
+			const grandParent = hydrate(GrandParent, {
+				parent: TreeAlpha.ensureSchema(Father, {
+					child: TreeAlpha.ensureSchema(Son, { value: 3 }),
+				}),
+			});
+			assert.ok(Tree.is(grandParent.parent, Father));
+			assert.ok(Tree.is(grandParent.parent?.child, Son));
+		});
+	});
 });
 
 function checkoutWithInitialTree(

packages/dds/tree/src/test/simple-tree/unhydratedFlexTreeFromInsertable.spec.ts

@@ -43,7 +43,7 @@ import {
 } from "../../feature-libraries/index.js";
 import { validateUsageError } from "../utils.js";
 // eslint-disable-next-line import/no-internal-modules
-import { UnhydratedFlexTreeNode } from "../../simple-tree/core/index.js";
+import { schemaSymbol, UnhydratedFlexTreeNode } from "../../simple-tree/core/index.js";
 // eslint-disable-next-line import/no-internal-modules
 import { getUnhydratedContext } from "../../simple-tree/createContext.js";
 // eslint-disable-next-line import/no-internal-modules
@@ -1459,6 +1459,29 @@ describe("unhydratedFlexTreeFromInsertable", () => {
 				[Optional],
 			);
 		});
+
+		it("with type symbol", () => {
+			const f = new SchemaFactory("test");
+			class A extends f.object("A", {
+				value: f.string,
+			}) {}
+			class B extends f.object("B", {
+				value: f.string,
+			}) {}
+			// Type symbol specified when there is only one valid option
+			const content = { [schemaSymbol]: "test.A", value: "hello" };
+			assert.deepEqual(getPossibleTypes(new Set([A]), content), [A]);
+
+			// Type symbol specified when there are multiple valid options
+			const ambiguousContent = { [schemaSymbol]: "test.B", value: "hello" };
+			// Only B, even though A is also valid based on properties
+			assert.deepEqual(getPossibleTypes(new Set([A, B]), ambiguousContent), [B]);
+
+			// Type symbol specified that does not match any options
+			const invalidContent = { [schemaSymbol]: "test.C", value: "hello" };
+			// Should fall back to A, as if no type symbol was provided
+			assert.deepEqual(getPossibleTypes(new Set([A]), invalidContent), [A]);
+		});
 	});
 
 	describe("defaults", () => {