@fluidframework/tree 2.3.1 → 2.4.0-297027

package/src/simple-tree/arrayNode.ts

@@ -3,31 +3,23 @@
  * Licensed under the MIT License.
  */
 
-import { oob } from "@fluidframework/core-utils/internal";
+import { Lazy, oob } from "@fluidframework/core-utils/internal";
 import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import { EmptyKey, type ExclusiveMapTree } from "../core/index.js";
 import {
-	type FlexTreeNodeSchema,
 	type FlexTreeNode,
 	type FlexTreeSequenceField,
-	type MapTreeNode,
-	getOrCreateMapTreeNode,
 	getSchemaAndPolicy,
 	isFlexTreeNode,
-	isMapTreeSequenceField,
-	UnhydratedContext,
 } from "../feature-libraries/index.js";
-import { getOrCreateNodeFromFlexTreeNode, prepareContentForHydration } from "./proxies.js";
+import { prepareContentForHydration } from "./proxies.js";
 import { getOrCreateInnerNode } from "./proxyBinding.js";
-// This import seems to trigger a false positive `Type import "TreeNodeFromImplicitAllowedTypes" is used by decorator metadata` lint error.
-// Other ways to import (ex: import the module with the items as type imports) give different more real errors, and auto fix to this format,
-// so there does not seem to be a clean workaround to make the linter happy.
-// eslint-disable-next-line @typescript-eslint/consistent-type-imports
-import type {
-	ImplicitAllowedTypes,
-	InsertableTreeNodeFromImplicitAllowedTypes,
-	TreeNodeFromImplicitAllowedTypes,
+import {
+	normalizeAllowedTypes,
+	type ImplicitAllowedTypes,
+	type InsertableTreeNodeFromImplicitAllowedTypes,
+	type TreeNodeFromImplicitAllowedTypes,
 } from "./schemaTypes.js";
 import {
 	type WithType,
@@ -36,15 +28,22 @@ import {
 	NodeKind,
 	type TreeNode,
 	type InternalTreeNode,
-	type TreeNodeSchemaClass,
 	type TreeNodeSchema,
 	typeSchemaSymbol,
+	type Context,
+	getOrCreateNodeFromInnerNode,
+	type TreeNodeSchemaBoth,
+	getSimpleNodeSchemaFromInnerNode,
 } from "./core/index.js";
 import { type InsertableContent, mapTreeFromNodeData } from "./toMapTree.js";
 import { fail } from "../util/index.js";
-import { getFlexSchema, toFlexSchema } from "./toFlexSchema.js";
-import { getKernel } from "./core/index.js";
+import {
+	getKernel,
+	UnhydratedFlexTreeNode,
+	UnhydratedTreeSequenceField,
+} from "./core/index.js";
 import { TreeNodeValid, type MostDerivedData } from "./treeNodeValid.js";
+import { getUnhydratedContext } from "./createContext.js";
 
 /**
  * A generic array type, used to defined types like {@link (TreeArrayNode:interface)}.
@@ -132,23 +131,89 @@ export interface TreeArrayNodeBase<out T, in TNew, in TMoveFrom>
 
 	/**
 	 * Moves the specified item to the desired location in the array.
-	 * @param index - The index to move the item to.
-	 * This is based on the state of the array before moving the source item.
+	 *
+	 * WARNING - This API is easily misused.
+	 * Please read the documentation for the `destinationGap` parameter carefully.
+	 *
+	 * @param destinationGap - The location *between* existing items that the moved item should be moved to.
+	 *
+	 * WARNING - `destinationGap` describes a location between existing items *prior to applying the move operation*.
+	 *
+	 * For example, if the array contains items `[A, B, C]` before the move, the `destinationGap` must be one of the following:
+	 *
+	 * - `0` (between the start of the array and `A`'s original position)
+	 * - `1` (between `A`'s original position and `B`'s original position)
+	 * - `2` (between `B`'s original position and `C`'s original position)
+	 * - `3` (between `C`'s original position and the end of the array)
+	 *
+	 * So moving `A` between `B` and `C` would require `destinationGap` to be `2`.
+	 *
+	 * This interpretation of `destinationGap` makes it easy to specify the desired destination relative to a sibling item that is not being moved,
+	 * or relative to the start or end of the array:
+	 *
+	 * - Move to the start of the array: `array.moveToIndex(0, ...)` (see also `moveToStart`)
+	 * - Move to before some item X: `array.moveToIndex(indexOfX, ...)`
+	 * - Move to after some item X: `array.moveToIndex(indexOfX + 1`, ...)
+	 * - Move to the end of the array: `array.moveToIndex(array.length, ...)` (see also `moveToEnd`)
+	 *
+	 * This interpretation of `destinationGap` does however make it less obvious how to move an item relative to its current position:
+	 *
+	 * - Move item B before its predecessor: `array.moveToIndex(indexOfB - 1, ...)`
+	 * - Move item B after its successor: `array.moveToIndex(indexOfB + 2, ...)`
+	 *
+	 * Notice the asymmetry between `-1` and `+2` in the above examples.
+	 * In such scenarios, it can often be easier to approach such edits by swapping adjacent items:
+	 * If items A and B are adjacent, such that A precedes B,
+	 * then they can be swapped with `array.moveToIndex(indexOfA, indexOfB)`.
+	 *
 	 * @param sourceIndex - The index of the item to move.
 	 * @throws Throws if any of the input indices are not in the range [0, `array.length`).
 	 */
-	moveToIndex(index: number, sourceIndex: number): void;
+	moveToIndex(destinationGap: number, sourceIndex: number): void;
 
 	/**
 	 * Moves the specified item to the desired location in the array.
-	 * @param index - The index to move the item to in the range [0, `array.length`].
-	 * This is based on the state of the array before moving the source item.
+	 *
+	 * WARNING - This API is easily misused.
+	 * Please read the documentation for the `destinationGap` parameter carefully.
+	 *
+	 * @param destinationGap - The location *between* existing items that the moved item should be moved to.
+	 *
+	 * WARNING - `destinationGap` describes a location between existing items *prior to applying the move operation*.
+	 *
+	 * For example, if the array contains items `[A, B, C]` before the move, the `destinationGap` must be one of the following:
+	 *
+	 * - `0` (between the start of the array and `A`'s original position)
+	 * - `1` (between `A`'s original position and `B`'s original position)
+	 * - `2` (between `B`'s original position and `C`'s original position)
+	 * - `3` (between `C`'s original position and the end of the array)
+	 *
+	 * So moving `A` between `B` and `C` would require `destinationGap` to be `2`.
+	 *
+	 * This interpretation of `destinationGap` makes it easy to specify the desired destination relative to a sibling item that is not being moved,
+	 * or relative to the start or end of the array:
+	 *
+	 * - Move to the start of the array: `array.moveToIndex(0, ...)` (see also `moveToStart`)
+	 * - Move to before some item X: `array.moveToIndex(indexOfX, ...)`
+	 * - Move to after some item X: `array.moveToIndex(indexOfX + 1`, ...)
+	 * - Move to the end of the array: `array.moveToIndex(array.length, ...)` (see also `moveToEnd`)
+	 *
+	 * This interpretation of `destinationGap` does however make it less obvious how to move an item relative to its current position:
+	 *
+	 * - Move item B before its predecessor: `array.moveToIndex(indexOfB - 1, ...)`
+	 * - Move item B after its successor: `array.moveToIndex(indexOfB + 2, ...)`
+	 *
+	 * Notice the asymmetry between `-1` and `+2` in the above examples.
+	 * In such scenarios, it can often be easier to approach such edits by swapping adjacent items:
+	 * If items A and B are adjacent, such that A precedes B,
+	 * then they can be swapped with `array.moveToIndex(indexOfA, indexOfB)`.
+	 *
 	 * @param sourceIndex - The index of the item to move.
 	 * @param source - The source array to move the item out of.
 	 * @throws Throws if any of the source index is not in the range [0, `array.length`),
 	 * or if the index is not in the range [0, `array.length`].
 	 */
-	moveToIndex(index: number, sourceIndex: number, source: TMoveFrom): void;
+	moveToIndex(destinationGap: number, sourceIndex: number, source: TMoveFrom): void;
 
 	/**
 	 * Moves the specified items to the start of the array.
@@ -192,18 +257,85 @@ export interface TreeArrayNodeBase<out T, in TNew, in TMoveFrom>
 
 	/**
 	 * Moves the specified items to the desired location within the array.
-	 * @param index - The index to move the items to.
-	 * This is based on the state of the array before moving the source items.
+	 *
+	 * WARNING - This API is easily misused.
+	 * Please read the documentation for the `destinationGap` parameter carefully.
+	 *
+	 * @param destinationGap - The location *between* existing items that the moved item should be moved to.
+	 *
+	 * WARNING - `destinationGap` describes a location between existing items *prior to applying the move operation*.
+	 *
+	 * For example, if the array contains items `[A, B, C]` before the move, the `destinationGap` must be one of the following:
+	 *
+	 * - `0` (between the start of the array and `A`'s original position)
+	 * - `1` (between `A`'s original position and `B`'s original position)
+	 * - `2` (between `B`'s original position and `C`'s original position)
+	 * - `3` (between `C`'s original position and the end of the array)
+	 *
+	 * So moving `A` between `B` and `C` would require `destinationGap` to be `2`.
+	 *
+	 * This interpretation of `destinationGap` makes it easy to specify the desired destination relative to a sibling item that is not being moved,
+	 * or relative to the start or end of the array:
+	 *
+	 * - Move to the start of the array: `array.moveToIndex(0, ...)` (see also `moveToStart`)
+	 * - Move to before some item X: `array.moveToIndex(indexOfX, ...)`
+	 * - Move to after some item X: `array.moveToIndex(indexOfX + 1`, ...)
+	 * - Move to the end of the array: `array.moveToIndex(array.length, ...)` (see also `moveToEnd`)
+	 *
+	 * This interpretation of `destinationGap` does however make it less obvious how to move an item relative to its current position:
+	 *
+	 * - Move item B before its predecessor: `array.moveToIndex(indexOfB - 1, ...)`
+	 * - Move item B after its successor: `array.moveToIndex(indexOfB + 2, ...)`
+	 *
+	 * Notice the asymmetry between `-1` and `+2` in the above examples.
+	 * In such scenarios, it can often be easier to approach such edits by swapping adjacent items:
+	 * If items A and B are adjacent, such that A precedes B,
+	 * then they can be swapped with `array.moveToIndex(indexOfA, indexOfB)`.
+	 *
 	 * @param sourceStart - The starting index of the range to move (inclusive).
 	 * @param sourceEnd - The ending index of the range to move (exclusive)
 	 * @throws Throws if any of the input indices are not in the range [0, `array.length`) or if `sourceStart` is greater than `sourceEnd`.
 	 * if any of the input indices are not in the range [0, `array.length`], or if `sourceStart` is greater than `sourceEnd`.
 	 */
-	moveRangeToIndex(index: number, sourceStart: number, sourceEnd: number): void;
+	moveRangeToIndex(destinationGap: number, sourceStart: number, sourceEnd: number): void;
 
 	/**
 	 * Moves the specified items to the desired location within the array.
-	 * @param index - The index to move the items to.
+	 *
+	 * WARNING - This API is easily misused.
+	 * Please read the documentation for the `destinationGap` parameter carefully.
+	 *
+	 * @param destinationGap - The location *between* existing items that the moved item should be moved to.
+	 *
+	 * WARNING - `destinationGap` describes a location between existing items *prior to applying the move operation*.
+	 *
+	 * For example, if the array contains items `[A, B, C]` before the move, the `destinationGap` must be one of the following:
+	 *
+	 * - `0` (between the start of the array and `A`'s original position)
+	 * - `1` (between `A`'s original position and `B`'s original position)
+	 * - `2` (between `B`'s original position and `C`'s original position)
+	 * - `3` (between `C`'s original position and the end of the array)
+	 *
+	 * So moving `A` between `B` and `C` would require `destinationGap` to be `2`.
+	 *
+	 * This interpretation of `destinationGap` makes it easy to specify the desired destination relative to a sibling item that is not being moved,
+	 * or relative to the start or end of the array:
+	 *
+	 * - Move to the start of the array: `array.moveToIndex(0, ...)` (see also `moveToStart`)
+	 * - Move to before some item X: `array.moveToIndex(indexOfX, ...)`
+	 * - Move to after some item X: `array.moveToIndex(indexOfX + 1`, ...)
+	 * - Move to the end of the array: `array.moveToIndex(array.length, ...)` (see also `moveToEnd`)
+	 *
+	 * This interpretation of `destinationGap` does however make it less obvious how to move an item relative to its current position:
+	 *
+	 * - Move item B before its predecessor: `array.moveToIndex(indexOfB - 1, ...)`
+	 * - Move item B after its successor: `array.moveToIndex(indexOfB + 2, ...)`
+	 *
+	 * Notice the asymmetry between `-1` and `+2` in the above examples.
+	 * In such scenarios, it can often be easier to approach such edits by swapping adjacent items:
+	 * If items A and B are adjacent, such that A precedes B,
+	 * then they can be swapped with `array.moveToIndex(indexOfA, indexOfB)`.
+	 *
 	 * @param sourceStart - The starting index of the range to move (inclusive).
 	 * @param sourceEnd - The ending index of the range to move (exclusive)
 	 * @param source - The source array to move items out of.
@@ -211,7 +343,7 @@ export interface TreeArrayNodeBase<out T, in TNew, in TMoveFrom>
 	 * if any of the input indices are not in the range [0, `array.length`], or if `sourceStart` is greater than `sourceEnd`.
 	 */
 	moveRangeToIndex(
-		index: number,
+		destinationGap: number,
 		sourceStart: number,
 		sourceEnd: number,
 		source: TMoveFrom,
@@ -559,7 +691,7 @@ function createArrayNodeProxy(
 
 			const maybeContent = field.at(maybeIndex);
 			return isFlexTreeNode(maybeContent)
-				? getOrCreateNodeFromFlexTreeNode(maybeContent)
+				? getOrCreateNodeFromInnerNode(maybeContent)
 				: maybeContent;
 		},
 		set: (target, key, newValue, receiver) => {
@@ -619,7 +751,7 @@ function createArrayNodeProxy(
 				// To satisfy 'deepEquals' level scrutiny, the property descriptor for indexed properties must
 				// be a simple value property (as opposed to using getter) and declared writable/enumerable/configurable.
 				return {
-					value: isFlexTreeNode(val) ? getOrCreateNodeFromFlexTreeNode(val) : val,
+					value: isFlexTreeNode(val) ? getOrCreateNodeFromInnerNode(val) : val,
 					writable: true, // For MVP, setting indexed properties is reported as allowed here (for deep equals compatibility noted above), but not actually supported.
 					enumerable: true,
 					configurable: true,
@@ -665,6 +797,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 	public static readonly kind = NodeKind.Array;
 
 	protected abstract get simpleSchema(): T;
+	protected abstract get allowedTypes(): ReadonlySet<TreeNodeSchema>;
 
 	public constructor(
 		input: Iterable<InsertableTreeNodeFromImplicitAllowedTypes<T>> | InternalTreeNode,
@@ -736,7 +869,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 			return val;
 		}
 
-		return getOrCreateNodeFromFlexTreeNode(val) as TreeNodeFromImplicitAllowedTypes<T>;
+		return getOrCreateNodeFromInnerNode(val) as TreeNodeFromImplicitAllowedTypes<T>;
 	}
 	public insertAt(index: number, ...value: Insertable<T>): void {
 		const field = getSequenceField(this);
@@ -781,16 +914,16 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		this.moveRangeToIndex(this.length, sourceIndex, sourceIndex + 1, source);
 	}
 	public moveToIndex(
-		destinationIndex: number,
+		destinationGap: number,
 		sourceIndex: number,
 		source?: TreeArrayNode,
 	): void {
 		const sourceArray = source ?? this;
 		const sourceField = getSequenceField(sourceArray);
 		const destinationField = getSequenceField(this);
-		validateIndex(destinationIndex, destinationField, "moveToIndex", true);
+		validateIndex(destinationGap, destinationField, "moveToIndex", true);
 		validateIndex(sourceIndex, sourceField, "moveToIndex");
-		this.moveRangeToIndex(destinationIndex, sourceIndex, sourceIndex + 1, source);
+		this.moveRangeToIndex(destinationGap, sourceIndex, sourceIndex + 1, source);
 	}
 	public moveRangeToStart(
 		sourceStart: number,
@@ -815,22 +948,24 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		this.moveRangeToIndex(this.length, sourceStart, sourceEnd, source);
 	}
 	public moveRangeToIndex(
-		destinationIndex: number,
+		destinationGap: number,
 		sourceStart: number,
 		sourceEnd: number,
 		source?: TreeArrayNode,
 	): void {
 		const destinationField = getSequenceField(this);
+		const destinationSchema = this.allowedTypes;
 		const sourceField = source !== undefined ? getSequenceField(source) : destinationField;
 
-		validateIndex(destinationIndex, destinationField, "moveRangeToIndex", true);
+		validateIndex(destinationGap, destinationField, "moveRangeToIndex", true);
 		validateIndexRange(sourceStart, sourceEnd, source ?? destinationField, "moveRangeToIndex");
 
 		// TODO: determine support for move across different sequence types
 		if (sourceField !== destinationField) {
 			for (let i = sourceStart; i < sourceEnd; i++) {
 				const sourceNode = sourceField.boxedAt(i) ?? oob();
-				if (!destinationField.schema.types.has(sourceNode.schema)) {
+				const sourceSchema = getSimpleNodeSchemaFromInnerNode(sourceNode);
+				if (!destinationSchema.has(sourceSchema)) {
 					throw new UsageError("Type in source sequence is not allowed in destination.");
 				}
 			}
@@ -838,20 +973,26 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 
 		const movedCount = sourceEnd - sourceStart;
 		if (!destinationField.context.isHydrated()) {
-			if (!isMapTreeSequenceField(sourceField)) {
+			if (!(sourceField instanceof UnhydratedTreeSequenceField)) {
 				throw new UsageError(
 					"Cannot move elements from an unhydrated array to a hydrated array.",
 				);
 			}
 
-			if (sourceField !== destinationField || destinationIndex < sourceStart) {
+			if (sourceField.context.isHydrated()) {
+				throw new UsageError(
+					"Cannot move elements from an unhydrated array to a hydrated array.",
+				);
+			}
+
+			if (sourceField !== destinationField || destinationGap < sourceStart) {
 				destinationField.editor.insert(
-					destinationIndex,
+					destinationGap,
 					sourceField.editor.remove(sourceStart, movedCount),
 				);
-			} else if (destinationIndex > sourceStart + movedCount) {
+			} else if (destinationGap > sourceStart + movedCount) {
 				destinationField.editor.insert(
-					destinationIndex - movedCount,
+					destinationGap - movedCount,
 					sourceField.editor.remove(sourceStart, movedCount),
 				);
 			}
@@ -870,7 +1011,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 				sourceStart,
 				movedCount,
 				destinationField.getFieldPath(),
-				destinationIndex,
+				destinationGap,
 			);
 		}
 	}
@@ -910,7 +1051,7 @@ export function arraySchema<
 	implicitlyConstructable: ImplicitlyConstructable,
 	customizable: boolean,
 ) {
-	type Output = TreeNodeSchemaClass<
+	type Output = TreeNodeSchemaBoth<
 		TName,
 		NodeKind.Array,
 		TreeArrayNode<T> & WithType<TName, NodeKind.Array>,
@@ -918,8 +1059,10 @@ export function arraySchema<
 		ImplicitlyConstructable,
 		T
 	>;
-	let flexSchema: FlexTreeNodeSchema;
-	let unhydratedContext: UnhydratedContext;
+
+	const lazyChildTypes = new Lazy(() => normalizeAllowedTypes(info));
+
+	let unhydratedContext: Context;
 
 	// This class returns a proxy from its constructor to handle numeric indexing.
 	// Alternatively it could extend a normal class which gets tons of numeric properties added.
@@ -948,20 +1091,18 @@ export function arraySchema<
 			this: typeof TreeNodeValid<T2>,
 			instance: TreeNodeValid<T2>,
 			input: T2,
-		): MapTreeNode {
-			return getOrCreateMapTreeNode(
+		): UnhydratedFlexTreeNode {
+			return UnhydratedFlexTreeNode.getOrCreate(
 				unhydratedContext,
-				flexSchema,
 				mapTreeFromNodeData(input as object, this as unknown as ImplicitAllowedTypes),
 			);
 		}
 
 		protected static override constructorCached: MostDerivedData | undefined = undefined;
 
-		protected static override oneTimeSetup<T2>(this: typeof TreeNodeValid<T2>): void {
+		protected static override oneTimeSetup<T2>(this: typeof TreeNodeValid<T2>): Context {
 			const schema = this as unknown as TreeNodeSchema;
-			flexSchema = getFlexSchema(schema);
-			unhydratedContext = new UnhydratedContext(toFlexSchema(schema));
+			unhydratedContext = getUnhydratedContext(schema);
 
 			// First run, do extra validation.
 			// TODO: provide a way for TreeConfiguration to trigger this same validation to ensure it gets run early.
@@ -987,12 +1128,17 @@ export function arraySchema<
 					prototype = Reflect.getPrototypeOf(prototype) as object;
 				}
 			}
+
+			return unhydratedContext;
 		}
 
 		public static readonly identifier = identifier;
 		public static readonly info = info;
 		public static readonly implicitlyConstructable: ImplicitlyConstructable =
 			implicitlyConstructable;
+		public static get childTypes(): ReadonlySet<TreeNodeSchema> {
+			return lazyChildTypes.value;
+		}
 
 		// eslint-disable-next-line import/no-deprecated
 		public get [typeNameSymbol](): TName {
@@ -1005,6 +1151,9 @@ export function arraySchema<
 		protected get simpleSchema(): T {
 			return info;
 		}
+		protected get allowedTypes(): ReadonlySet<TreeNodeSchema> {
+			return lazyChildTypes.value;
+		}
 	}
 
 	const output: Output = Schema;

package/src/simple-tree/core/context.ts

@@ -0,0 +1,78 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { anchorSlot, type TreeNodeSchemaIdentifier } from "../../core/index.js";
+import type {
+	FlexTreeContext,
+	FlexTreeHydratedContext,
+} from "../../feature-libraries/index.js";
+import { brand } from "../../util/index.js";
+import type { TreeNodeSchema } from "./treeNodeSchema.js";
+import { walkAllowedTypes } from "./walkSchema.js";
+
+/**
+ * Creating multiple simple tree contexts for the same branch, and thus with the same underlying AnchorSet does not work due to how TreeNode caching works.
+ * This slot is used to detect if one already exists and error if creating a second.
+ * @remarks
+ * See also {@link ContextSlot} in which the flex-tree context is stored.
+ */
+export const SimpleContextSlot = anchorSlot<HydratedContext>();
+
+/**
+ * Additional information about a collection of {@link TreeNode}s.
+ * @remarks
+ * Each TreeNode provides a way to navigate to a context for it with additional information.
+ *
+ * For unhydrated nodes, this information is rather limited since the node doesn't know what tree it might get inserted into (if any),
+ * and thus is limited to information derived about that particular unhydrated tree and its schema.
+ *
+ * Hydrated nodes have more contextual information, and thus can provide a single {@link HydratedContext} for all nodes in the document which can have additional information.
+ *
+ * @privateRemarks
+ * This design is the same as {@link FlexTreeContext} with its base type and {@link FlexTreeHydratedContext} extending it.
+ */
+export class Context {
+	/**
+	 * All schema which could transitively be used under the associated node.
+	 * @remarks
+	 * While generally {@link TreeNodeSchema} are referenced as objects and thus do not need to be looked up by identifier,
+	 * there are a few cases (mainly constructing new TreeNodes from existing tree data) where such a lookup is useful.
+	 * Having this map in the context addresses this use-case.
+	 * @privateRemarks
+	 * This design mirrors how {@link FlexTreeSchema} are accessed off the {@link FlexTreeContext}, making the migration away from them simpler.
+	 */
+	public readonly schema: ReadonlyMap<TreeNodeSchemaIdentifier, TreeNodeSchema>;
+
+	/**
+	 * Builds the context.
+	 * @remarks
+	 * 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>,
+		public readonly flexContext: FlexTreeContext,
+	) {
+		const schema: Map<TreeNodeSchemaIdentifier, TreeNodeSchema> = new Map();
+		walkAllowedTypes(rootSchema, {
+			node(nodeSchema) {
+				schema.set(brand(nodeSchema.identifier), nodeSchema);
+			},
+		});
+		this.schema = schema;
+	}
+}
+
+/**
+ * Extends {@link Context} with additional information from the view and/or tree which becomes known when a {@link TreeNode} is hydrated,
+ * associating it with the containing {@link TreeView}.
+ */
+export class HydratedContext extends Context {
+	public constructor(
+		rootSchema: Iterable<TreeNodeSchema>,
+		public override readonly flexContext: FlexTreeHydratedContext,
+	) {
+		super(rootSchema, flexContext);
+	}
+}

package/src/simple-tree/core/getOrCreateNode.ts

@@ -0,0 +1,38 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import type { TreeValue } from "../../core/index.js";
+import type { FlexTreeNode } from "../../feature-libraries/index.js";
+import { fail } from "../../util/index.js";
+import { type InnerNode, mapTreeNodeToProxy, proxySlot } from "./treeNodeKernel.js";
+import { getSimpleNodeSchemaFromInnerNode } from "./schemaCaching.js";
+import type { TreeNode, InternalTreeNode } from "./types.js";
+import { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
+
+/**
+ * Returns the TreeNode or TreeValue for the provided {@link InnerNode}.
+ * This will allocate a new one if needed, and otherwise return one from cache.
+ * @remarks
+ * This supports both hydrated and unhydrated nodes.
+ */
+export function getOrCreateNodeFromInnerNode(flexNode: InnerNode): TreeNode | TreeValue {
+	const cached =
+		flexNode instanceof UnhydratedFlexTreeNode
+			? mapTreeNodeToProxy.get(flexNode)
+			: flexNode.anchorNode.slots.get(proxySlot);
+
+	if (cached !== undefined) {
+		return cached;
+	}
+
+	const classSchema = getSimpleNodeSchemaFromInnerNode(flexNode) ?? fail("Missing schema");
+	const node = flexNode as unknown as InternalTreeNode;
+	// eslint-disable-next-line unicorn/prefer-ternary
+	if (typeof classSchema === "function") {
+		return new classSchema(node) as TreeNode;
+	} else {
+		return (classSchema as { create(data: FlexTreeNode): TreeValue }).create(flexNode);
+	}
+}

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

@@ -9,8 +9,8 @@ export {
 	getKernel,
 	tryGetTreeNodeSchema,
 	type InnerNode,
-	tryGetCachedTreeNode,
 	tryDisposeTreeNode,
+	tryGetTreeNodeFromMapNode,
 } from "./treeNodeKernel.js";
 export { type WithType, typeNameSymbol, typeSchemaSymbol } from "./withType.js";
 export {
@@ -27,11 +27,15 @@ export {
 	type TreeNodeSchemaClass,
 	type TreeNodeSchemaNonClass,
 	type TreeNodeSchemaCore,
+	type TreeNodeSchemaBoth,
 } from "./treeNodeSchema.js";
+export { getSimpleNodeSchemaFromInnerNode } from "./schemaCaching.js";
+export { walkAllowedTypes, type SchemaVisitor } from "./walkSchema.js";
+export { Context, HydratedContext, SimpleContextSlot } from "./context.js";
+export { getOrCreateNodeFromInnerNode } from "./getOrCreateNode.js";
 export {
-	getSimpleNodeSchema,
-	setFlexSchemaFromClassSchema,
-	tryGetSimpleNodeSchema,
-	cachedFlexSchemaFromClassSchema,
-	getSimpleNodeSchemaFromNode,
-} from "./schemaCaching.js";
+	UnhydratedFlexTreeNode,
+	UnhydratedTreeSequenceField,
+	tryUnhydratedFlexTreeNode,
+	UnhydratedContext,
+} from "./unhydratedFlexTree.js";