@fluidframework/tree 2.2.0 → 2.3.0

package/src/feature-libraries/flex-map-tree/mapTreeNode.ts

@@ -4,6 +4,8 @@
  */
 
 import { assert, oob } from "@fluidframework/core-utils/internal";
+import { UsageError } from "@fluidframework/telemetry-utils/internal";
+
 import {
 	type AnchorNode,
 	EmptyKey,
@@ -11,43 +13,40 @@ import {
 	type FieldKey,
 	type FieldUpPath,
 	type MapTree,
+	type TreeFieldStoredSchema,
 	type TreeNodeSchemaIdentifier,
-	type TreeValue,
+	type TreeStoredSchema,
 	type Value,
 } from "../../core/index.js";
-import { brand, fail, getOrCreate, mapIterable } from "../../util/index.js";
+import { brand, fail, getOrCreate, isReadonlyArray, mapIterable } from "../../util/index.js";
 import {
+	type FlexTreeContext,
 	FlexTreeEntityKind,
 	type FlexTreeField,
-	type FlexTreeLeafNode,
-	type FlexTreeMapNode,
 	type FlexTreeNode,
 	type FlexTreeOptionalField,
 	type FlexTreeRequiredField,
 	type FlexTreeSequenceField,
 	type FlexTreeTypedField,
-	type FlexTreeTypedNode,
-	type FlexTreeTypedNodeUnion,
-	type FlexTreeUnboxField,
-	type FlexTreeUnboxNodeUnion,
+	type FlexTreeUnknownUnboxed,
 	flexTreeMarker,
 	indexForAt,
+	type FlexTreeHydratedContext,
 } from "../flex-tree/index.js";
 import {
 	type FlexAllowedTypes,
 	FlexFieldSchema,
-	type FlexMapNodeSchema,
 	type FlexTreeNodeSchema,
-	type LeafNodeSchema,
+	type FlexTreeSchema,
+	intoStoredSchemaCollection,
 	isLazy,
-	schemaIsLeaf,
-	schemaIsMap,
-	schemaIsObjectNode,
 } from "../typed-schema/index.js";
-import { type FlexImplicitAllowedTypes, normalizeAllowedTypes } from "../schemaBuilderBase.js";
 import type { FlexFieldKind } from "../modular-schema/index.js";
-import { FieldKinds, type SequenceFieldEditBuilder } from "../default-schema/index.js";
-import { UsageError } from "@fluidframework/telemetry-utils/internal";
+import {
+	defaultSchemaPolicy,
+	FieldKinds,
+	type SequenceFieldEditBuilder,
+} from "../default-schema/index.js";
 
 // #region Nodes
 
@@ -70,9 +69,9 @@ export function isMapTreeNode(flexNode: FlexTreeNode): flexNode is MapTreeNode {
 /**
  * Checks if the given {@link FlexTreeField} is a {@link MapTreeSequenceField}.
  */
-export function isMapTreeSequenceField<T extends FlexAllowedTypes>(
-	field: FlexTreeSequenceField<T> | FlexTreeField,
-): field is MapTreeSequenceField<T> {
+export function isMapTreeSequenceField(
+	field: FlexTreeSequenceField | FlexTreeField,
+): field is MapTreeSequenceField {
 	return field instanceof EagerMapTreeSequenceField;
 }
 
@@ -80,8 +79,7 @@ export function isMapTreeSequenceField<T extends FlexAllowedTypes>(
  * An unhydrated {@link FlexTreeSequenceField}, which has additional editing capabilities.
  * @remarks When doing a removal edit, a {@link MapTreeSequenceField}'s `editor` returns ownership of the removed {@link ExclusiveMapTree}s to the caller.
  */
-export interface MapTreeSequenceField<T extends FlexAllowedTypes>
-	extends FlexTreeSequenceField<T> {
+export interface MapTreeSequenceField extends FlexTreeSequenceField {
 	readonly editor: MapTreeSequenceFieldEditBuilder;
 }
 
@@ -109,7 +107,11 @@ interface LocationInField {
  * MapTreeNodes are unconditionally cached -
  * when retrieved via {@link getOrCreateNode}, the same {@link MapTree} object will always produce the same `MapTreeNode` object.
  */
-export class EagerMapTreeNode<TSchema extends FlexTreeNodeSchema> implements MapTreeNode {
+export class EagerMapTreeNode implements MapTreeNode {
+	public get schema(): TreeNodeSchemaIdentifier {
+		return this.flexSchema.name;
+	}
+
 	public readonly [flexTreeMarker] = FlexTreeEntityKind.Node as const;
 
 	/**
@@ -121,7 +123,8 @@ export class EagerMapTreeNode<TSchema extends FlexTreeNodeSchema> implements Map
 	 * Instead, it should always be acquired via {@link getOrCreateNode}.
 	 */
 	public constructor(
-		public readonly schema: TSchema,
+		public readonly context: UnhydratedContext,
+		public readonly flexSchema: FlexTreeNodeSchema,
 		/** The underlying {@link MapTree} that this `MapTreeNode` reads its data from */
 		public readonly mapTree: ExclusiveMapTree,
 		private location = unparentedLocation,
@@ -140,7 +143,7 @@ export class EagerMapTreeNode<TSchema extends FlexTreeNodeSchema> implements Map
 	}
 
 	public get type(): TreeNodeSchemaIdentifier {
-		return this.schema.name;
+		return this.flexSchema.name;
 	}
 
 	/**
@@ -177,37 +180,38 @@ export class EagerMapTreeNode<TSchema extends FlexTreeNodeSchema> implements Map
 		return this.location;
 	}
 
-	public is<TSchemaInner extends FlexTreeNodeSchema>(
-		schema: TSchemaInner,
-	): this is FlexTreeTypedNode<TSchemaInner> {
-		return (schema as unknown) === this.schema;
+	public is(schema: FlexTreeNodeSchema): boolean {
+		return (schema as unknown) === this.flexSchema;
 	}
 
-	public tryGetField(key: FieldKey): EagerMapTreeField<FlexAllowedTypes> | undefined {
+	public tryGetField(key: FieldKey): EagerMapTreeField | undefined {
 		const field = this.mapTree.fields.get(key);
 		// Only return the field if it is not empty, in order to fulfill the contract of `tryGetField`.
 		if (field !== undefined && field.length > 0) {
-			return getOrCreateField(this, key, this.schema.getFieldSchema(key));
+			return getOrCreateField(this, key, this.flexSchema.getFieldSchema(key));
 		}
 	}
 
 	public getBoxed(key: string): FlexTreeField {
 		const fieldKey: FieldKey = brand(key);
-		return getOrCreateField(this, fieldKey, this.schema.getFieldSchema(fieldKey));
+		return getOrCreateField(this, fieldKey, this.flexSchema.getFieldSchema(fieldKey));
 	}
 
 	public boxedIterator(): IterableIterator<FlexTreeField> {
 		return mapIterable(this.mapTree.fields.entries(), ([key]) =>
-			getOrCreateField(this, key, this.schema.getFieldSchema(key)),
+			getOrCreateField(this, key, this.flexSchema.getFieldSchema(key)),
 		);
 	}
 
+	public keys(): IterableIterator<FieldKey> {
+		// TODO: how this should handle missing defaults (and empty keys if they end up being allowed) needs to be determined.
+		return this.mapTree.fields.keys();
+	}
+
 	public get value(): Value {
 		return this.mapTree.value;
 	}
 
-	public context = undefined;
-
 	public get anchorNode(): AnchorNode {
 		// This API is relevant to `LazyTreeNode`s, but not `MapTreeNode`s.
 		// TODO: Refactor the FlexTreeNode interface so that stubbing this out isn't necessary.
@@ -216,11 +220,12 @@ export class EagerMapTreeNode<TSchema extends FlexTreeNodeSchema> implements Map
 
 	private walkTree(): void {
 		for (const [key, mapTrees] of this.mapTree.fields) {
-			const field = getOrCreateField(this, key, this.schema.getFieldSchema(key));
+			const field = getOrCreateField(this, key, this.flexSchema.getFieldSchema(key));
 			for (let index = 0; index < field.length; index++) {
 				const child = getOrCreateChild(
+					this.context,
 					mapTrees[index] ?? oob(),
-					this.schema.getFieldSchema(key).allowedTypes,
+					this.flexSchema.getFieldSchema(key).allowedTypes,
 					{ parent: field, index },
 				);
 				// These next asserts detect the case where `getOrCreateChild` gets a cache hit of a different node than the one we're trying to create
@@ -236,84 +241,31 @@ export class EagerMapTreeNode<TSchema extends FlexTreeNodeSchema> implements Map
 	}
 }
 
+// #endregion Nodes
+
 /**
- * The implementation of a map node created by {@link getOrCreateNode}.
+ * Implementation of `FlexTreeContext`.
+ *
+ * @remarks An editor is required to edit the FlexTree.
  */
-export class EagerMapTreeMapNode<TSchema extends FlexMapNodeSchema>
-	extends EagerMapTreeNode<TSchema>
-	implements FlexTreeMapNode<TSchema>
-{
-	public keys(): IterableIterator<FieldKey> {
-		return this.mapTree.fields.keys();
-	}
-
-	public values(): IterableIterator<FlexTreeUnboxField<TSchema["info"], "notEmpty">> {
-		return mapIterable(
-			this.mapTree.fields.keys(),
-			(key) =>
-				unboxedField(
-					this.tryGetField(key) ?? fail("Unexpected empty map field"),
-					key,
-					this.mapTree,
-					this,
-				) as FlexTreeUnboxField<TSchema["info"], "notEmpty">,
-		);
-	}
-
-	public entries(): IterableIterator<
-		[FieldKey, FlexTreeUnboxField<TSchema["info"], "notEmpty">]
-	> {
-		return mapIterable(this.mapTree.fields.keys(), (key) => [
-			key,
-			unboxedField(
-				this.tryGetField(key) ?? fail("Unexpected empty map field"),
-				key,
-				this.mapTree,
-				this,
-			) as FlexTreeUnboxField<TSchema["info"], "notEmpty">,
-		]);
-	}
-
-	public forEach(
-		callbackFn: (
-			value: FlexTreeUnboxField<TSchema["info"], "notEmpty">,
-			key: FieldKey,
-			map: FlexTreeMapNode<TSchema>,
-		) => void,
-		thisArg?: unknown,
-	): void {
-		const fn = thisArg !== undefined ? callbackFn.bind(thisArg) : callbackFn;
-		for (const [key, value] of this.entries()) {
-			fn(value, key, this);
-		}
-	}
+export class UnhydratedContext implements FlexTreeContext {
+	public readonly schema: TreeStoredSchema;
 
-	public override getBoxed(key: string): FlexTreeTypedField<TSchema["info"]> {
-		return super.getBoxed(key) as FlexTreeTypedField<TSchema["info"]>;
-	}
-
-	public [Symbol.iterator](): IterableIterator<
-		[FieldKey, FlexTreeUnboxField<TSchema["info"], "notEmpty">]
-	> {
-		return this.entries();
+	/**
+	 * @param flexSchema - Schema to use when working with the tree.
+	 */
+	public constructor(public readonly flexSchema: FlexTreeSchema) {
+		this.schema = {
+			rootFieldSchema: flexSchema.rootFieldSchema.stored,
+			...intoStoredSchemaCollection(flexSchema),
+		};
 	}
 
-	public override boxedIterator(): IterableIterator<FlexTreeTypedField<TSchema["info"]>> {
-		return super.boxedIterator() as IterableIterator<FlexTreeTypedField<TSchema["info"]>>;
+	public isHydrated(): this is FlexTreeHydratedContext {
+		return false;
 	}
 }
 
-class EagerMapTreeLeafNode<TSchema extends LeafNodeSchema>
-	extends EagerMapTreeNode<TSchema>
-	implements FlexTreeLeafNode<TSchema>
-{
-	public override get value(): TreeValue<TSchema["info"]> {
-		return super.value as TreeValue<TSchema["info"]>;
-	}
-}
-
-// #endregion Nodes
-
 // #region Fields
 
 /**
@@ -325,6 +277,13 @@ interface MapTreeField extends FlexTreeField {
 	readonly mapTrees: readonly MapTree[];
 }
 
+const emptyContext = new UnhydratedContext({
+	adapters: {},
+	nodeSchema: new Map(),
+	policy: defaultSchemaPolicy,
+	rootFieldSchema: FlexFieldSchema.empty,
+});
+
 /**
  * A special singleton that is the implicit {@link LocationInField} of all un-parented {@link EagerMapTreeNode}s.
  * @remarks This exists because {@link EagerMapTreeNode.parentField} must return a field.
@@ -339,8 +298,8 @@ const unparentedLocation: LocationInField = {
 		length: 0,
 		key: EmptyKey,
 		parent: undefined,
-		is<TSchema extends FlexFieldSchema>(schema: TSchema) {
-			return schema === (FlexFieldSchema.empty as FlexFieldSchema);
+		is<TKind2 extends FlexFieldKind>(kind: TKind2) {
+			return this.schema.kind === kind.identifier;
 		},
 		boxedIterator(): IterableIterator<FlexTreeNode> {
 			return [].values();
@@ -348,20 +307,28 @@ const unparentedLocation: LocationInField = {
 		boxedAt(index: number): FlexTreeNode | undefined {
 			return undefined;
 		},
-		schema: FlexFieldSchema.empty,
-		context: undefined,
+		schema: FlexFieldSchema.empty.stored,
+		context: emptyContext,
 		mapTrees: [],
+		getFieldPath() {
+			fail("unsupported");
+		},
 	},
 	index: -1,
 };
 
-class EagerMapTreeField<T extends FlexAllowedTypes> implements MapTreeField {
+class EagerMapTreeField implements MapTreeField {
 	public [flexTreeMarker] = FlexTreeEntityKind.Field as const;
 
+	public get schema(): TreeFieldStoredSchema {
+		return this.flexSchema.stored;
+	}
+
 	public constructor(
-		public readonly schema: FlexFieldSchema<FlexFieldKind, T>,
+		public readonly context: UnhydratedContext,
+		public readonly flexSchema: FlexFieldSchema,
 		public readonly key: FieldKey,
-		public readonly parent: EagerMapTreeNode<FlexTreeNodeSchema>,
+		public readonly parent: EagerMapTreeNode,
 	) {
 		const fieldKeyCache = getFieldKeyCache(parent);
 		assert(!fieldKeyCache.has(key), 0x990 /* A field already exists for the given MapTrees */);
@@ -388,40 +355,36 @@ class EagerMapTreeField<T extends FlexAllowedTypes> implements MapTreeField {
 		return this.mapTrees.length;
 	}
 
-	public is<TSchemaInner extends FlexFieldSchema>(
-		schema: TSchemaInner,
-	): this is FlexTreeTypedField<TSchemaInner> {
-		return this.schema.equals(schema);
+	public is<TKind2 extends FlexFieldKind>(kind: TKind2): this is FlexTreeTypedField<TKind2> {
+		return this.schema.kind === kind.identifier;
 	}
 
-	public boxedIterator(): IterableIterator<FlexTreeTypedNodeUnion<T>> {
+	public boxedIterator(): IterableIterator<FlexTreeNode> {
 		return this.mapTrees
 			.map(
 				(m, index) =>
-					getOrCreateChild(m, this.schema.allowedTypes, {
+					getOrCreateChild(this.context, m, this.flexSchema.allowedTypes, {
 						parent: this,
 						index,
-					}) as FlexTreeNode as FlexTreeTypedNodeUnion<T>,
+					}) as FlexTreeNode,
 			)
 			.values();
 	}
 
-	public boxedAt(index: number): FlexTreeTypedNodeUnion<T> | undefined {
+	public boxedAt(index: number): FlexTreeNode | undefined {
 		const i = indexForAt(index, this.length);
 		if (i === undefined) {
 			return undefined;
 		}
 		const m = this.mapTrees[i];
 		if (m !== undefined) {
-			return getOrCreateChild(m, this.schema.allowedTypes, {
+			return getOrCreateChild(this.context, m, this.flexSchema.allowedTypes, {
 				parent: this,
 				index: i,
-			}) as FlexTreeNode as FlexTreeTypedNodeUnion<T>;
+			}) as FlexTreeNode;
 		}
 	}
 
-	public context: undefined;
-
 	/**
 	 * Mutate this field.
 	 * @param edit - A function which receives the current `MapTree`s that comprise the contents of the field so that it may be mutated.
@@ -439,12 +402,13 @@ class EagerMapTreeField<T extends FlexAllowedTypes> implements MapTreeField {
 			this.parent.mapTree.fields.delete(this.key);
 		}
 	}
+
+	public getFieldPath(): FieldUpPath {
+		throw unsupportedUsageError("Editing an array");
+	}
 }
 
-class EagerMapTreeOptionalField<T extends FlexAllowedTypes>
-	extends EagerMapTreeField<T>
-	implements FlexTreeOptionalField<T>
-{
+class EagerMapTreeOptionalField extends EagerMapTreeField implements FlexTreeOptionalField {
 	public readonly editor = {
 		set: (newContent: ExclusiveMapTree | undefined): void => {
 			// If the new content is a MapTreeNode, it needs to have its parent pointer updated
@@ -467,10 +431,10 @@ class EagerMapTreeOptionalField<T extends FlexAllowedTypes>
 		},
 	};
 
-	public get content(): FlexTreeUnboxNodeUnion<T> | undefined {
+	public get content(): FlexTreeUnknownUnboxed | undefined {
 		const value = this.mapTrees[0];
 		if (value !== undefined) {
-			return unboxedUnion(this.schema, value, {
+			return unboxed(this.flexSchema, value, {
 				parent: this,
 				index: 0,
 			});
@@ -480,19 +444,16 @@ class EagerMapTreeOptionalField<T extends FlexAllowedTypes>
 	}
 }
 
-class EagerMapTreeRequiredField<T extends FlexAllowedTypes>
-	extends EagerMapTreeOptionalField<T>
-	implements FlexTreeRequiredField<T>
+class EagerMapTreeRequiredField
+	extends EagerMapTreeOptionalField
+	implements FlexTreeRequiredField
 {
-	public override get content(): FlexTreeUnboxNodeUnion<T> {
+	public override get content(): FlexTreeUnknownUnboxed {
 		return super.content ?? fail("Expected EagerMapTree required field to have a value");
 	}
 }
 
-class EagerMapTreeSequenceField<T extends FlexAllowedTypes>
-	extends EagerMapTreeField<T>
-	implements FlexTreeSequenceField<T>
-{
+class EagerMapTreeSequenceField extends EagerMapTreeField implements FlexTreeSequenceField {
 	public readonly editor: MapTreeSequenceFieldEditBuilder = {
 		insert: (index, newContent): void => {
 			for (let i = 0; i < newContent.length; i++) {
@@ -524,44 +485,32 @@ class EagerMapTreeSequenceField<T extends FlexAllowedTypes>
 		},
 	};
 
-	public at(index: number): FlexTreeUnboxNodeUnion<T> | undefined {
+	public at(index: number): FlexTreeUnknownUnboxed | undefined {
 		const i = indexForAt(index, this.length);
 		if (i === undefined) {
 			return undefined;
 		}
-		return unboxedUnion(this.schema, this.mapTrees[i] ?? oob(), { parent: this, index: i });
+		return unboxed(this.flexSchema, this.mapTrees[i] ?? oob(), { parent: this, index: i });
 	}
-	public map<U>(callbackfn: (value: FlexTreeUnboxNodeUnion<T>, index: number) => U): U[] {
+	public map<U>(callbackfn: (value: FlexTreeUnknownUnboxed, index: number) => U): U[] {
 		return Array.from(this, callbackfn);
 	}
-	public mapBoxed<U>(callbackfn: (value: FlexTreeTypedNodeUnion<T>, index: number) => U): U[] {
-		return Array.from(this.boxedIterator(), callbackfn);
-	}
 
-	public *[Symbol.iterator](): IterableIterator<FlexTreeUnboxNodeUnion<T>> {
+	public *[Symbol.iterator](): IterableIterator<FlexTreeUnknownUnboxed> {
 		for (const [i, mapTree] of this.mapTrees.entries()) {
-			yield unboxedUnion(this.schema, mapTree, { parent: this, index: i });
+			yield unboxed(this.flexSchema, mapTree, { parent: this, index: i });
 		}
 	}
-
-	public getFieldPath(): FieldUpPath {
-		throw unsupportedUsageError("Editing an array");
-	}
 }
 
 // #endregion Fields
 
 // #region Caching and unboxing utilities
 
-const nodeCache = new WeakMap<MapTree, EagerMapTreeNode<FlexTreeNodeSchema>>();
+const nodeCache = new WeakMap<MapTree, EagerMapTreeNode>();
 /** Node Parent -\> Field Key -\> Field */
-const fieldCache = new WeakMap<
-	MapTreeNode,
-	Map<FieldKey, EagerMapTreeField<FlexAllowedTypes>>
->();
-function getFieldKeyCache(
-	parent: MapTreeNode,
-): WeakMap<FieldKey, EagerMapTreeField<FlexAllowedTypes>> {
+const fieldCache = new WeakMap<MapTreeNode, Map<FieldKey, EagerMapTreeField>>();
+function getFieldKeyCache(parent: MapTreeNode): WeakMap<FieldKey, EagerMapTreeField> {
 	return getOrCreate(fieldCache, parent, () => new Map());
 }
 
@@ -580,59 +529,52 @@ export function tryGetMapTreeNode(mapTree: MapTree): MapTreeNode | undefined {
  * @remarks It must conform to the `nodeSchema`.
  */
 export function getOrCreateMapTreeNode(
+	context: UnhydratedContext,
 	nodeSchema: FlexTreeNodeSchema,
 	mapTree: ExclusiveMapTree,
-): EagerMapTreeNode<FlexTreeNodeSchema> {
-	return nodeCache.get(mapTree) ?? createNode(nodeSchema, mapTree, undefined);
+): EagerMapTreeNode {
+	return nodeCache.get(mapTree) ?? createNode(context, nodeSchema, mapTree, undefined);
 }
 
 /** Helper for creating a `MapTreeNode` given the parent field (e.g. when "walking down") */
 function getOrCreateChild(
+	context: UnhydratedContext,
 	mapTree: ExclusiveMapTree,
-	implicitAllowedTypes: FlexImplicitAllowedTypes,
+	allowedTypes: FlexAllowedTypes,
 	parent: LocationInField | undefined,
-): EagerMapTreeNode<FlexTreeNodeSchema> {
+): EagerMapTreeNode {
 	const cached = nodeCache.get(mapTree);
 	if (cached !== undefined) {
 		return cached;
 	}
 
-	const allowedTypes = normalizeAllowedTypes(implicitAllowedTypes);
+	assert(isReadonlyArray(allowedTypes), 0xa25 /* invalid types */);
 	const nodeSchema =
 		allowedTypes
 			.map((t) => (isLazy(t) ? t() : t))
 			.find((t): t is FlexTreeNodeSchema => {
-				assert(t !== "Any", 0x993 /* 'Any' type is not supported */);
 				return t.name === mapTree.type;
 			}) ?? fail("Unsupported node schema");
 
-	return createNode(nodeSchema, mapTree, parent);
+	return createNode(context, nodeSchema, mapTree, parent);
 }
 
 /** Always constructs a new node, therefore may not be called twice for the same `MapTree`. */
-function createNode<TSchema extends FlexTreeNodeSchema>(
-	nodeSchema: TSchema,
+function createNode(
+	context: UnhydratedContext,
+	nodeSchema: FlexTreeNodeSchema,
 	mapTree: ExclusiveMapTree,
 	parentField: LocationInField | undefined,
-): EagerMapTreeNode<TSchema> {
-	if (schemaIsLeaf(nodeSchema)) {
-		return new EagerMapTreeLeafNode(nodeSchema, mapTree, parentField);
-	}
-	if (schemaIsMap(nodeSchema)) {
-		return new EagerMapTreeMapNode(nodeSchema, mapTree, parentField);
-	}
-	if (schemaIsObjectNode(nodeSchema)) {
-		return new EagerMapTreeNode(nodeSchema, mapTree, parentField);
-	}
-	assert(false, 0x994 /* Unrecognized node kind */);
+): EagerMapTreeNode {
+	return new EagerMapTreeNode(context, nodeSchema, mapTree, parentField);
 }
 
 /** Creates a field with the given attributes, or returns a cached field if there is one */
 function getOrCreateField(
-	parent: EagerMapTreeNode<FlexTreeNodeSchema>,
+	parent: EagerMapTreeNode,
 	key: FieldKey,
 	schema: FlexFieldSchema,
-): EagerMapTreeField<FlexFieldSchema["allowedTypes"]> {
+): EagerMapTreeField {
 	const cached = getFieldKeyCache(parent).get(key);
 	if (cached !== undefined) {
 		return cached;
@@ -642,70 +584,32 @@ function getOrCreateField(
 		schema.kind.identifier === FieldKinds.required.identifier ||
 		schema.kind.identifier === FieldKinds.identifier.identifier
 	) {
-		return new EagerMapTreeRequiredField(schema, key, parent);
+		return new EagerMapTreeRequiredField(parent.context, schema, key, parent);
 	}
 
 	if (schema.kind.identifier === FieldKinds.optional.identifier) {
-		return new EagerMapTreeOptionalField(schema, key, parent);
+		return new EagerMapTreeOptionalField(parent.context, schema, key, parent);
 	}
 
 	if (schema.kind.identifier === FieldKinds.sequence.identifier) {
-		return new EagerMapTreeSequenceField(schema, key, parent);
+		return new EagerMapTreeSequenceField(parent.context, schema, key, parent);
 	}
 
-	return new EagerMapTreeField(schema, key, parent);
+	return new EagerMapTreeField(parent.context, schema, key, parent);
 }
 
-/** Unboxes non-polymorphic leaf nodes to their values, if applicable */
-function unboxedUnion<TTypes extends FlexAllowedTypes>(
-	schema: FlexFieldSchema<FlexFieldKind, TTypes>,
+/** Unboxes leaf nodes to their values */
+function unboxed(
+	schema: FlexFieldSchema,
 	mapTree: ExclusiveMapTree,
 	parent: LocationInField,
-): FlexTreeUnboxNodeUnion<TTypes> {
-	const type = schema.monomorphicChildType;
-	if (type !== undefined) {
-		if (schemaIsLeaf(type)) {
-			return mapTree.value as FlexTreeUnboxNodeUnion<TTypes>;
-		}
-		return getOrCreateChild(mapTree, type, parent) as FlexTreeUnboxNodeUnion<TTypes>;
+): FlexTreeUnknownUnboxed {
+	const value = mapTree.value;
+	if (value !== undefined) {
+		return value;
 	}
 
-	return getOrCreateChild(
-		mapTree,
-		schema.allowedTypes,
-		parent,
-	) as FlexTreeUnboxNodeUnion<TTypes>;
-}
-
-/** Unboxes non-polymorphic required and optional fields holding leaf nodes to their values, if applicable */
-function unboxedField<TFieldSchema extends FlexFieldSchema>(
-	field: EagerMapTreeField<FlexAllowedTypes>,
-	key: FieldKey,
-	mapTree: ExclusiveMapTree,
-	parentNode: EagerMapTreeNode<FlexTreeNodeSchema>,
-): FlexTreeUnboxField<TFieldSchema> {
-	const fieldSchema = field.schema;
-	const mapTrees =
-		mapTree.fields.get(key) ?? fail("Key does not exist in unhydrated map tree");
-
-	if (fieldSchema.kind === FieldKinds.required) {
-		return unboxedUnion(fieldSchema, mapTrees[0] ?? oob(), {
-			parent: field,
-			index: 0,
-		}) as FlexTreeUnboxField<TFieldSchema>;
-	}
-	if (fieldSchema.kind === FieldKinds.optional) {
-		return (
-			mapTrees.length > 0
-				? unboxedUnion(fieldSchema, mapTrees[0] ?? oob(), {
-						parent: field,
-						index: 0,
-					})
-				: undefined
-		) as FlexTreeUnboxField<TFieldSchema>;
-	}
-
-	return getOrCreateField(parentNode, key, fieldSchema) as FlexTreeUnboxField<TFieldSchema>;
+	return getOrCreateChild(parent.parent.context, mapTree, schema.allowedTypes, parent);
 }
 
 // #endregion Caching and unboxing utilities