@fluidframework/tree 2.23.0 → 2.30.0

package/src/simple-tree/api/conciseTree.ts

@@ -8,7 +8,12 @@ import type { IFluidHandle } from "@fluidframework/core-interfaces";
 import type { ITreeCursor } from "../../core/index.js";
 import type { TreeLeafValue, ImplicitAllowedTypes } from "../schemaTypes.js";
 import type { TreeNodeSchema } from "../core/index.js";
-import { customFromCursor, type EncodeOptions } from "./customTree.js";
+import {
+	customFromCursor,
+	replaceHandles,
+	type EncodeOptions,
+	type HandleConverter,
+} from "./customTree.js";
 import { getUnhydratedContext } from "../createContext.js";
 
 /**
@@ -36,12 +41,12 @@ export type ConciseTree<THandle = IFluidHandle> =
 /**
  * Used to read a node cursor as a ConciseTree.
  */
-export function conciseFromCursor<TCustom>(
+export function conciseFromCursor(
 	reader: ITreeCursor,
 	rootSchema: ImplicitAllowedTypes,
-	options: EncodeOptions<TCustom>,
-): ConciseTree<TCustom> {
-	const config: Required<EncodeOptions<TCustom>> = {
+	options: EncodeOptions,
+): ConciseTree {
+	const config: Required<EncodeOptions> = {
 		useStoredKeys: false,
 		...options,
 	};
@@ -50,10 +55,22 @@ export function conciseFromCursor<TCustom>(
 	return conciseFromCursorInner(reader, config, schemaMap);
 }
 
-function conciseFromCursorInner<TCustom>(
+function conciseFromCursorInner(
 	reader: ITreeCursor,
-	options: Required<EncodeOptions<TCustom>>,
+	options: Required<EncodeOptions>,
 	schema: ReadonlyMap<string, TreeNodeSchema>,
-): ConciseTree<TCustom> {
+): ConciseTree {
 	return customFromCursor(reader, options, schema, conciseFromCursorInner);
 }
+
+/**
+ * Clones tree, replacing any handles.
+ * @remarks A strongly typed version of {@link replaceHandles}.
+ * @alpha
+ */
+export function replaceConciseTreeHandles<T>(
+	tree: ConciseTree,
+	replacer: HandleConverter<T>,
+): ConciseTree<T> {
+	return replaceHandles(tree, replacer) as ConciseTree<T>;
+}

package/src/simple-tree/api/create.ts

@@ -26,7 +26,7 @@ import {
 	defaultSchemaPolicy,
 	FieldKinds,
 	mapTreeFromCursor,
-	type NodeKeyManager,
+	type NodeIdentifierManager,
 } from "../../feature-libraries/index.js";
 import { isFieldInSchema } from "../../feature-libraries/index.js";
 import { toStoredSchema } from "../toStoredSchema.js";
@@ -53,7 +53,7 @@ export function createFromInsertable<
 		? ImplicitFieldSchema
 		: TSchema & ImplicitFieldSchema,
 	data: InsertableField<TSchema>,
-	context?: NodeKeyManager | undefined,
+	context?: NodeIdentifierManager | undefined,
 ): Unhydrated<
 	TSchema extends ImplicitFieldSchema
 		? TreeFieldFromImplicitField<TSchema>
@@ -84,7 +84,7 @@ export function cursorFromInsertable<
 		? ImplicitFieldSchema
 		: TSchema & ImplicitFieldSchema,
 	data: InsertableField<TSchema>,
-	context?: NodeKeyManager | undefined,
+	context?: NodeIdentifierManager | undefined,
 ):
 	| ITreeCursorSynchronous
 	| (TSchema extends FieldSchema<FieldKind.Optional> ? undefined : never) {

package/src/simple-tree/api/customTree.ts

@@ -14,11 +14,12 @@ import {
 	LeafNodeStoredSchema,
 	mapCursorField,
 	ObjectNodeStoredSchema,
+	type FieldKey,
 	type ITreeCursor,
 	type TreeNodeSchemaIdentifier,
 	type TreeNodeStoredSchema,
 } from "../../core/index.js";
-import { fail } from "../../util/index.js";
+import { fail, cloneWithReplacements } from "../../util/index.js";
 import type { TreeLeafValue } from "../schemaTypes.js";
 import { NodeKind, type TreeNodeSchema } from "../core/index.js";
 import {
@@ -35,13 +36,7 @@ import { FieldKinds, valueSchemaAllows } from "../../feature-libraries/index.js"
  * Options for how to encode a tree.
  * @alpha
  */
-export interface EncodeOptions<TCustom> {
-	/**
-	 * How to encode any {@link @fluidframework/core-interfaces#IFluidHandle|IFluidHandles} in the tree.
-	 * @remarks
-	 * See note on {@link ParseOptions.valueConverter}.
-	 */
-	valueConverter(data: IFluidHandle): TCustom;
+export interface EncodeOptions {
 	/**
 	 * If true, interpret the input keys of object nodes as stored keys.
 	 * If false, interpret them as property keys.
@@ -50,15 +45,36 @@ export interface EncodeOptions<TCustom> {
 	readonly useStoredKeys?: boolean;
 }
 
+/**
+ * Options for how to transcode handles.
+ * @remarks
+ * Can be applied using {@link replaceHandles}.
+ * @alpha
+ */
+export type HandleConverter<TCustom> = (data: IFluidHandle) => TCustom;
+
+/**
+ * Options for how to interpret a `ConciseTree<TCustom>` without relying on schema.
+ */
+export interface SchemalessParseOptions {
+	/**
+	 * Converts stored keys into whatever key the tree is using in its encoding.
+	 */
+	keyConverter?: {
+		parse(type: string, inputKey: string): FieldKey;
+		encode(type: string, key: FieldKey): string;
+	};
+}
+
 /**
  * Tree representation with fields as properties and customized handle and child representations.
  */
-export type CustomTree<TChild, THandle> = CustomTreeNode<TChild> | CustomTreeValue<THandle>;
+export type CustomTree<TChild> = CustomTreeNode<TChild> | CustomTreeValue;
 
 /**
  * TreeLeafValue except the handle type is customized.
  */
-export type CustomTreeValue<THandle> = Exclude<TreeLeafValue, IFluidHandle> | THandle;
+export type CustomTreeValue = TreeLeafValue;
 
 /**
  * Tree node representation with fields as properties and customized child representation.
@@ -68,16 +84,16 @@ export type CustomTreeNode<TChild> = TChild[] | { [key: string]: TChild };
 /**
  * Builds an {@link CustomTree} from a cursor in Nodes mode.
  */
-export function customFromCursor<TChild, THandle>(
+export function customFromCursor<TChild>(
 	reader: ITreeCursor,
-	options: Required<EncodeOptions<THandle>>,
+	options: Required<EncodeOptions>,
 	schema: ReadonlyMap<string, TreeNodeSchema>,
 	childHandler: (
 		reader: ITreeCursor,
-		options: Required<EncodeOptions<THandle>>,
+		options: Required<EncodeOptions>,
 		schema: ReadonlyMap<string, TreeNodeSchema>,
 	) => TChild,
-): CustomTree<TChild, THandle> {
+): CustomTree<TChild> {
 	const type = reader.type;
 	const nodeSchema = schema.get(type) ?? fail(0xb2e /* missing schema for type in cursor */);
 
@@ -92,7 +108,7 @@ export function customFromCursor<TChild, THandle>(
 		case handleSchema.identifier:
 			assert(reader.value !== undefined, 0xa52 /* out of schema: missing value */);
 			assert(isFluidHandle(reader.value), 0xa53 /* out of schema: expected FluidHandle */);
-			return options.valueConverter(reader.value);
+			return reader.value;
 		default: {
 			assert(reader.value === undefined, 0xa54 /* out of schema: unexpected value */);
 			if (nodeSchema.kind === NodeKind.Array) {
@@ -136,7 +152,7 @@ export function customFromCursorStored<TChild>(
 		reader: ITreeCursor,
 		schema: ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeStoredSchema>,
 	) => TChild,
-): CustomTree<TChild, IFluidHandle> {
+): CustomTree<TChild> {
 	const type = reader.type;
 	const nodeSchema = schema.get(type) ?? fail(0xb30 /* missing schema for type in cursor */);
 
@@ -186,3 +202,30 @@ export function tryStoredSchemaAsArray(
 		}
 	}
 }
+
+/**
+ * Clones tree, replacing any handles.
+ * @remarks
+ * This can be useful converting data containing handles to JSON compatible formats,
+ * or just asserting that data does not contain handles.
+ *
+ * Reversing this replacement depends on how the replacer encodes handles, and can often be impossible if the replacer
+ * does not have all the necessary context to restore the handles
+ * (e.g. if the handles are something insufficiently descriptive,
+ * if data referenced by the handle got garbage collected,
+ * if the encoded form of the handle can't be differentiated from other data,
+ * or the replacer doesn't have access to the correct Fluid container to to restore them from).
+ *
+ * Code attempting to reverse this replacement may want to use {@link cloneWithReplacements}.
+ * @alpha
+ */
+export function replaceHandles<T>(tree: unknown, replacer: HandleConverter<T>): unknown {
+	return cloneWithReplacements(tree, "", (key, value) => {
+		// eslint-disable-next-line unicorn/prefer-ternary
+		if (isFluidHandle(value)) {
+			return { clone: false, value: replacer(value) };
+		} else {
+			return { clone: true, value };
+		}
+	});
+}

package/src/simple-tree/api/index.ts

@@ -95,6 +95,7 @@ export {
 	applySchemaToParserOptions,
 	cursorFromVerbose,
 	verboseFromCursor,
+	replaceVerboseTreeHandles,
 } from "./verboseTree.js";
 
 export {
@@ -103,9 +104,15 @@ export {
 	type CustomTreeNode,
 	type CustomTreeValue,
 	tryStoredSchemaAsArray,
+	replaceHandles,
+	type HandleConverter,
 } from "./customTree.js";
 
-export { type ConciseTree, conciseFromCursor } from "./conciseTree.js";
+export {
+	type ConciseTree,
+	conciseFromCursor,
+	replaceConciseTreeHandles,
+} from "./conciseTree.js";
 
 export { TreeBeta, type NodeChangedData, type TreeChangeEventsBeta } from "./treeApiBeta.js";
 export { createSimpleTreeIndex, type SimpleTreeIndex } from "./simpleTreeIndex.js";
@@ -133,6 +140,8 @@ export {
 	rollback,
 } from "./transactionTypes.js";
 
+export { generateSchemaFromSimpleSchema } from "./schemaFromSimple.js";
+
 // Exporting the schema (RecursiveObject) to test that recursive types are working correctly.
 // These are `@internal` so they can't be included in the `InternalClassTreeTypes` due to https://github.com/microsoft/rushstack/issues/3639
 export {

package/src/simple-tree/api/schemaFactory.ts

@@ -8,11 +8,10 @@ import { assert, unreachableCase } from "@fluidframework/core-utils/internal";
 // which degrades the API-Extractor report quality since API-Extractor can not tell the inline import is the same as the non-inline one.
 // eslint-disable-next-line unused-imports/no-unused-imports
 import type { IFluidHandle as _dummyImport } from "@fluidframework/core-interfaces";
-import { UsageError } from "@fluidframework/telemetry-utils/internal";
 import { isFluidHandle } from "@fluidframework/runtime-utils/internal";
 
 import type { TreeValue } from "../../core/index.js";
-import type { NodeKeyManager } from "../../feature-libraries/index.js";
+import type { NodeIdentifierManager } from "../../feature-libraries/index.js";
 import {
 	type RestrictiveStringRecord,
 	getOrCreate,
@@ -25,7 +24,6 @@ import type { TreeAlpha } from "../../shared-tree/index.js";
 import {
 	booleanSchema,
 	handleSchema,
-	LeafNodeSchema,
 	nullSchema,
 	numberSchema,
 	stringSchema,
@@ -41,8 +39,8 @@ import {
 	type DefaultProvider,
 	getDefaultProvider,
 	type NodeSchemaOptions,
+	markSchemaMostDerived,
 } from "../schemaTypes.js";
-import { inPrototypeChain } from "../core/index.js";
 import type {
 	NodeKind,
 	WithType,
@@ -76,7 +74,6 @@ import type {
 	Unenforced,
 } from "./typesUnsafe.js";
 import { createFieldSchemaUnsafe } from "./schemaFactoryRecursive.js";
-import { TreeNodeValid } from "../treeNodeValid.js";
 import { isLazy } from "../flexList.js";
 
 /**
@@ -145,7 +142,7 @@ export interface SchemaFactoryObjectOptions<TCustomMetadata = unknown>
 	 * ```
 	 *
 	 * If an application wants to be particularly careful to preserve all data on a node when editing it, it can use
-	 * {@link TreeAlpha.(importVerbose:2)|import}/{@link TreeAlpha.(exportVerbose:2)|export} APIs with persistent keys.
+	 * {@link TreeAlpha.importVerbose|import}/{@link TreeAlpha.exportVerbose|export} APIs with persistent keys.
 	 *
 	 * Note that public API methods which operate on entire nodes (such as `moveTo`, `moveToEnd`, etc. on arrays) do not encounter
 	 * this problem as SharedTree's implementation stores the entire node in its lower layers. It's only when application code
@@ -329,6 +326,17 @@ export const schemaStatics = {
  * On the other hand, two calls to `array(name, allowedTypes)` will always return different {@link TreeNodeSchema} instances
  * and it is an error to use both in the same tree (since their identifiers are not unique).
  *
+ * For "customizable" schema (those which can be subclassed to customize them, see details below) some additional rules must be followed:
+ *
+ * 1. Only a single {@link TreeNodeSchema|schema} can be used from the class hierarchy deriving from the base class produced by this factory.
+ * It is legal to subclass the returned class, and even subclass that class,
+ * but only a single class from that class hierarchy can ever be instantiated or passed to any API as a {@link TreeNodeSchema|schema}.
+ * These base classes can be used with `instanceof`, but not with schema based APIs like `Tree.is`.
+ * 2. If overriding the constructor, the constructor must accept the same argument as the base constructor `super` and forward it to `super` unchanged.
+ * 3. Properties for fields defined in the schema should not be overridden.
+ * 4. Additional static members added to schema should pick relatively unique keys to reduce the risk of colliding with implementation details what are not exposed in the API.
+ * 5. If exporting the schema from a package which uses API-Extractor, export the base class and derived class separately to work around [a known limitation](https://github.com/microsoft/rushstack/issues/4429).
+ *
  * Note:
  * POJO stands for Plain Old JavaScript Object.
  * This means an object that works like a `{}` style object literal.
@@ -901,8 +909,10 @@ export class SchemaFactory<
 	 */
 	public get identifier(): FieldSchema<FieldKind.Identifier, typeof this.string> {
 		const defaultIdentifierProvider: DefaultProvider = getDefaultProvider(
-			(nodeKeyManager: NodeKeyManager) => {
-				return nodeKeyManager.stabilizeNodeKey(nodeKeyManager.generateLocalNodeKey());
+			(nodeKeyManager: NodeIdentifierManager) => {
+				return nodeKeyManager.stabilizeNodeIdentifier(
+					nodeKeyManager.generateLocalNodeIdentifier(),
+				);
 			},
 		);
 		return createFieldSchema(FieldKind.Identifier, this.string, {
@@ -1070,27 +1080,3 @@ export function structuralName<const T extends string>(
 	}
 	return `${collectionName}<${inner}>`;
 }
-
-/**
- * Indicates that a schema is the "most derived" version which is allowed to be used, see {@link MostDerivedData}.
- * Calling helps with error messages about invalid schema usage (using more than one type from single schema factor produced type,
- * and thus calling this for one than one subclass).
- * @remarks
- * Helper for invoking {@link TreeNodeValid.markMostDerived} for any {@link TreeNodeSchema} if it needed.
- */
-export function markSchemaMostDerived(schema: TreeNodeSchema): void {
-	if (schema instanceof LeafNodeSchema) {
-		return;
-	}
-
-	if (!inPrototypeChain(schema, TreeNodeValid)) {
-		// Use JSON.stringify to quote and escape identifier string.
-		throw new UsageError(
-			`Schema for ${JSON.stringify(
-				schema.identifier,
-			)} does not extend a SchemaFactory generated class. This is invalid.`,
-		);
-	}
-
-	(schema as typeof TreeNodeValid & TreeNodeSchema).markMostDerived();
-}

package/src/simple-tree/api/schemaFactoryAlpha.ts

@@ -250,4 +250,16 @@ export class SchemaFactoryAlpha<
 			TCustomMetadata
 		>;
 	}
+
+	/**
+	 * Create a {@link SchemaFactory} with a {@link SchemaFactory.scope|scope} which is a combination of this factory's scope and the provided name.
+	 * @remarks
+	 * The main use-case for this is when creating a collection of related schema (for example using a function that creates multiple schema).
+	 * Creating such related schema using a sub-scope helps ensure they won't collide with other schema in the parent scope.
+	 */
+	public scopedFactory<const T extends TName, TNameInner extends number | string = string>(
+		name: T,
+	): SchemaFactoryAlpha<ScopedSchemaName<TScope, T>, TNameInner> {
+		return new SchemaFactoryAlpha(this.scoped2(name));
+	}
 }

package/src/simple-tree/api/schemaFromSimple.ts

@@ -0,0 +1,68 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { unreachableCase } from "@fluidframework/core-utils/internal";
+import { fail } from "../../util/index.js";
+import { NodeKind, type TreeNodeSchema } from "../core/index.js";
+import { createFieldSchema, type FieldSchema, type AllowedTypes } from "../schemaTypes.js";
+import { SchemaFactory } from "./schemaFactory.js";
+import type { SimpleFieldSchema, SimpleNodeSchema, SimpleTreeSchema } from "./simpleSchema.js";
+
+const factory = new SchemaFactory(undefined);
+
+/**
+ * Create {@link FieldSchema} from a SimpleTreeSchema.
+ * @remarks
+ * Only use this API if hand written schema (produced using {@link SchemaFactory} cannot be provided.
+ *
+ * Using generated schema with schema aware APIs (designed to work with strongly typed schema) like {@link TreeViewConfiguration}
+ * will produce a poor TypeScript typing experience which is subject to change.
+ *
+ * Editing through a view produced using this schema can easily violate invariants other users of the document might expect and must be done with great care.
+ * @internal
+ */
+export function generateSchemaFromSimpleSchema(simple: SimpleTreeSchema): FieldSchema {
+	const context: Context = new Map(
+		[...simple.definitions].map(([id, schema]): [string, () => TreeNodeSchema] => [
+			id,
+			// This relies on the caching in evaluateLazySchema so that it only runs once.
+			() => generateNode(id, schema, context),
+		]),
+	);
+	return generateFieldSchema(simple, context);
+}
+
+type Context = ReadonlyMap<string, () => TreeNodeSchema>;
+
+function generateFieldSchema(simple: SimpleFieldSchema, context: Context): FieldSchema {
+	return createFieldSchema(simple.kind, generateAllowedTypes(simple.allowedTypes, context));
+}
+
+function generateAllowedTypes(allowed: ReadonlySet<string>, context: Context): AllowedTypes {
+	return [...allowed].map((id) => context.get(id) ?? fail(0xb5a /* Missing schema */));
+}
+
+function generateNode(id: string, schema: SimpleNodeSchema, context: Context): TreeNodeSchema {
+	switch (schema.kind) {
+		case NodeKind.Object: {
+			const fields: Record<string, FieldSchema> = {};
+			for (const [key, field] of Object.entries(schema.fields)) {
+				fields[key] = generateFieldSchema(field, context);
+			}
+			return factory.object(id, fields);
+		}
+		case NodeKind.Array:
+			return factory.array(id, generateAllowedTypes(schema.allowedTypes, context));
+		case NodeKind.Map:
+			return factory.map(id, generateAllowedTypes(schema.allowedTypes, context));
+		case NodeKind.Leaf:
+			return (
+				SchemaFactory.leaves.find((leaf) => leaf.identifier === id) ??
+				fail(0xb5b /* Missing schema */)
+			);
+		default:
+			return unreachableCase(schema);
+	}
+}

package/src/simple-tree/api/simpleSchema.ts

@@ -7,6 +7,13 @@ import type { ValueSchema } from "../../core/index.js";
 import type { NodeKind } from "../core/index.js";
 import type { FieldKind, FieldSchemaMetadata, NodeSchemaMetadata } from "../schemaTypes.js";
 
+/*
+ * TODO:
+ * - Make TreeNodeSchema implement these interfaces directly.
+ * - Customize their JSON serialization to use these formats or provide some other serialization scheme.
+ * - Promote these to alpha
+ */
+
 /**
  * Base interface for all {@link SimpleNodeSchema} implementations.
  *
@@ -38,6 +45,9 @@ export interface SimpleObjectNodeSchema extends SimpleNodeSchemaBase<NodeKind.Ob
 	 * Schemas for each of the object's fields, keyed off of schema's keys.
 	 * @remarks
 	 * Depending on how this schema was exported, the string keys may be either the property keys or the stored keys.
+	 * @privateRemarks
+	 * TODO: if these are supposed to be JSON compatible,
+	 * then using a record here makes sense, but if not, this should use a map, and the allowedTypes sets elsewhere should be arrays for JSON compatibility.
 	 */
 	readonly fields: Record<string, SimpleFieldSchema>;
 }

package/src/simple-tree/api/tree.ts

@@ -27,13 +27,13 @@ import {
 	type TreeFieldFromImplicitField,
 	type UnsafeUnknownSchema,
 	FieldKind,
+	markSchemaMostDerived,
 } from "../schemaTypes.js";
 import { NodeKind, type TreeNodeSchema } from "../core/index.js";
 import { toStoredSchema } from "../toStoredSchema.js";
 import { LeafNodeSchema } from "../leafNodeSchema.js";
 import { assert } from "@fluidframework/core-utils/internal";
 import { isObjectNodeSchema, type ObjectNodeSchema } from "../objectNodeTypes.js";
-import { markSchemaMostDerived } from "./schemaFactory.js";
 import { fail, getOrCreate } from "../../util/index.js";
 import type { MakeNominal } from "../../util/index.js";
 import { walkFieldSchema } from "../walkFieldSchema.js";
@@ -264,7 +264,17 @@ export class TreeViewConfiguration<
 	public readonly preventAmbiguity: boolean;
 
 	/**
+	 * Construct a new {@link TreeViewConfiguration}.
+	 *
 	 * @param props - Property bag of configuration options.
+	 *
+	 * @remarks
+	 * Performing this construction deeply validates the provided schema.
+	 * This means that when this constructor is called, all {@link LazyItem} {@link TreeNodeSchema} references will be evaluated (using {@link evaluateLazySchema}).
+	 * This means that the declarations for all transitively reachable {@link TreeNodeSchema} must be available at this time.
+	 *
+	 * For example, a schema reachable from this configuration cannot reference this configuration during its declaration,
+	 * since this would be a cyclic dependency that will cause an error when constructing this configuration.
 	 */
 	public constructor(props: ITreeViewConfiguration<TSchema>) {
 		const config = { ...defaultTreeConfigurationOptions, ...props };
@@ -280,7 +290,7 @@ export class TreeViewConfiguration<
 			// This ensures if multiple schema extending the same schema factory generated class are present (or have been constructed, or get constructed in the future),
 			// an error is reported.
 
-			node: markSchemaMostDerived,
+			node: (schema) => markSchemaMostDerived(schema, true),
 			allowedTypes(types): void {
 				if (config.preventAmbiguity) {
 					checkUnion(types, ambiguityErrors);
@@ -725,12 +735,7 @@ export interface SchemaCompatibilityStatus {
  * Events for {@link TreeBranch}.
  * @sealed @alpha
  */
-export interface TreeBranchEvents {
-	/**
-	 * The stored schema for the document has changed.
-	 */
-	schemaChanged(): void;
-
+export interface TreeBranchEvents extends Omit<TreeViewEvents, "commitApplied"> {
 	/**
 	 * Fired when a change is made to the branch. Includes data about the change that is made which listeners
 	 * can use to filter on changes they care about (e.g. local vs. remote changes).

package/src/simple-tree/api/treeNodeApi.ts

@@ -7,13 +7,14 @@ import { assert, oob } from "@fluidframework/core-utils/internal";
 
 import { EmptyKey, rootFieldKey } from "../../core/index.js";
 import { type TreeStatus, isTreeValue, FieldKinds } from "../../feature-libraries/index.js";
-import { fail, extractFromOpaque, isReadonlyArray } from "../../util/index.js";
+import { fail, extractFromOpaque } from "../../util/index.js";
 import {
 	type TreeLeafValue,
 	type ImplicitFieldSchema,
 	FieldSchema,
 	type ImplicitAllowedTypes,
 	type TreeNodeFromImplicitAllowedTypes,
+	normalizeAllowedTypes,
 } from "../schemaTypes.js";
 import {
 	booleanSchema,
@@ -39,8 +40,6 @@ import {
 	getOrCreateInnerNode,
 } from "../core/index.js";
 import { isObjectNodeSchema } from "../objectNodeTypes.js";
-import { isLazy, type LazyItem } from "../flexList.js";
-import { markSchemaMostDerived } from "./schemaFactory.js";
 
 /**
  * Provides various functions for analyzing {@link TreeNode}s.
@@ -212,19 +211,7 @@ export const treeNodeApi: TreeNodeApi = {
 		if (actualSchema === undefined) {
 			return false;
 		}
-		if (isReadonlyArray<LazyItem<TreeNodeSchema>>(schema)) {
-			for (const singleSchema of schema) {
-				const testSchema = isLazy(singleSchema) ? singleSchema() : singleSchema;
-				markSchemaMostDerived(testSchema);
-				if (testSchema === actualSchema) {
-					return true;
-				}
-			}
-			return false;
-		} else {
-			markSchemaMostDerived(schema);
-			return schema === actualSchema;
-		}
+		return normalizeAllowedTypes(schema).has(actualSchema);
 	},
 	schema(node: TreeNode | TreeLeafValue): TreeNodeSchema {
 		return tryGetSchema(node) ?? fail(0xb37 /* Not a tree node */);
@@ -258,7 +245,7 @@ export const treeNodeApi: TreeNodeApi = {
 				const identifierValue = identifier.value as string;
 
 				const localNodeKey =
-					identifier.context.nodeKeyManager.tryLocalizeNodeKey(identifierValue);
+					identifier.context.nodeKeyManager.tryLocalizeNodeIdentifier(identifierValue);
 				return localNodeKey !== undefined ? extractFromOpaque(localNodeKey) : identifierValue;
 			}
 			default: