@fluidframework/tree 2.4.0-297027 → 2.4.0-299374

package/src/simple-tree/{typesUnsafe.ts → api/typesUnsafe.ts}

@@ -3,27 +3,27 @@
  * Licensed under the MIT License.
  */
 
-import type { FlexListToUnion, Unenforced } from "../feature-libraries/index.js";
-import type { RestrictiveStringRecord, _InlineTrick } from "../util/index.js";
+import type { RestrictiveStringRecord, _InlineTrick } from "../../util/index.js";
 
 import type {
-	AllowedTypes,
 	ApplyKind,
 	FieldKind,
 	FieldSchema,
 	ImplicitAllowedTypes,
 	ImplicitFieldSchema,
-	NodeFromSchema,
-	TreeNodeFromImplicitAllowedTypes,
-} from "./schemaTypes.js";
+	TreeLeafValue,
+} from "../schemaTypes.js";
 import type {
 	NodeKind,
-	TreeNodeSchema,
 	WithType,
 	TreeNode,
 	Unhydrated,
-} from "./core/index.js";
-import type { TreeArrayNodeBase, TreeArrayNode } from "./arrayNode.js";
+	InternalTreeNode,
+	TreeNodeSchema,
+	TreeNodeSchemaCore,
+} from "../core/index.js";
+import type { TreeArrayNodeBase } from "../arrayNode.js";
+import type { FlexListToUnion, LazyItem } from "../flexList.js";
 
 /*
  * TODO:
@@ -37,6 +37,16 @@ import type { TreeArrayNodeBase, TreeArrayNode } from "./arrayNode.js";
  */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
+/**
+ * A placeholder to use in {@link https://www.typescriptlang.org/docs/handbook/2/generics.html#generic-constraints | extends constraints} when using the real type breaks compilation of some recursive types due to {@link https://github.com/microsoft/TypeScript/issues/55758 | a design limitation of TypeScript}.
+ *
+ * These extends constraints only serve as documentation:
+ * to avoid breaking compilation, this type has to not actually enforce anything, and thus is just `unknown`.
+ * Therefore the type safety is the responsibility of the user of the API.
+ * @public
+ */
+export type Unenforced<_DesiredExtendsConstraint> = unknown;
+
 /**
  * {@link Unenforced} version of `ObjectFromSchemaRecord`.
  * @remarks
@@ -49,6 +59,65 @@ export type ObjectFromSchemaRecordUnsafe<
 	-readonly [Property in keyof T]: TreeFieldFromImplicitFieldUnsafe<T[Property]>;
 };
 
+/**
+ * {@link Unenforced} version of {@link TreeNodeSchema}.
+ * @remarks
+ * Do note use this type directly: its only needed in the implementation of generic logic which define recursive schema, not when using recursive schema.
+ * @system @public
+ */
+export type TreeNodeSchemaUnsafe<
+	Name extends string = string,
+	Kind extends NodeKind = NodeKind,
+	TNode extends Unenforced<TreeNode | TreeLeafValue> = unknown,
+	TBuild = never,
+	ImplicitlyConstructable extends boolean = boolean,
+	Info = unknown,
+> =
+	| TreeNodeSchemaClassUnsafe<Name, Kind, TNode, TBuild, ImplicitlyConstructable, Info>
+	| TreeNodeSchemaNonClassUnsafe<Name, Kind, TNode, TBuild, ImplicitlyConstructable, Info>;
+
+/**
+ * {@link Unenforced} version of {@link TreeNodeSchemaClass}.
+ * @remarks
+ * Do note use this type directly: its only needed in the implementation of generic logic which define recursive schema, not when using recursive schema.
+ * @system @public
+ */
+export interface TreeNodeSchemaClassUnsafe<
+	out Name extends string,
+	out Kind extends NodeKind,
+	// TODO: maybe this can be more specific (exclude leaves)
+	out TNode extends Unenforced<TreeNode | TreeLeafValue>,
+	in TInsertable = never,
+	out ImplicitlyConstructable extends boolean = boolean,
+	out Info = unknown,
+> extends TreeNodeSchemaCore<Name, Kind, ImplicitlyConstructable, Info> {
+	/**
+	 * Constructs an {@link Unhydrated} node with this schema.
+	 * @remarks
+	 * This constructor is also used internally to construct hydrated nodes with a different parameter type.
+	 * Therefore, overriding this constructor with different argument types is not type-safe and is not supported.
+	 * @sealed
+	 */
+	new (data: TInsertable | InternalTreeNode): Unhydrated<TNode>;
+}
+
+/**
+ * {@link Unenforced} version of {@link TreeNodeSchemaNonClass}.
+ * @remarks
+ * Do note use this type directly: its only needed in the implementation of generic logic which define recursive schema, not when using recursive schema.
+ * @system @public
+ */
+export interface TreeNodeSchemaNonClassUnsafe<
+	out Name extends string,
+	out Kind extends NodeKind,
+	out TNode extends Unenforced<TreeNode | TreeLeafValue>,
+	in TInsertable = never,
+	out ImplicitlyConstructable extends boolean = boolean,
+	out Info = unknown,
+> extends TreeNodeSchemaCore<Name, Kind, ImplicitlyConstructable, Info> {
+	create(data: TInsertable): TNode;
+}
+
 /**
  * {@link Unenforced} version of {@link TreeObjectNode}.
  * @remarks
@@ -73,6 +142,14 @@ export type TreeFieldFromImplicitFieldUnsafe<TSchema extends Unenforced<Implicit
 			? TreeNodeFromImplicitAllowedTypesUnsafe<TSchema>
 			: unknown;
 
+/**
+ * {@link Unenforced} version of {@link AllowedTypes}.
+ * @remarks
+ * Do note use this type directly: its only needed in the implementation of generic logic which define recursive schema, not when using recursive schema.
+ * @system @public
+ */
+export type AllowedTypesUnsafe = readonly LazyItem<TreeNodeSchemaUnsafe>[];
+
 /**
  * {@link Unenforced} version of {@link TreeNodeFromImplicitAllowedTypes}.
  * @remarks
@@ -81,13 +158,11 @@ export type TreeFieldFromImplicitFieldUnsafe<TSchema extends Unenforced<Implicit
  */
 export type TreeNodeFromImplicitAllowedTypesUnsafe<
 	TSchema extends Unenforced<ImplicitAllowedTypes>,
-> = TSchema extends ImplicitAllowedTypes
-	? TreeNodeFromImplicitAllowedTypes<TSchema>
-	: TSchema extends TreeNodeSchema
-		? NodeFromSchema<TSchema>
-		: TSchema extends AllowedTypes
-			? NodeFromSchema<FlexListToUnion<TSchema>>
-			: unknown;
+> = TSchema extends TreeNodeSchemaUnsafe
+	? NodeFromSchemaUnsafe<TSchema>
+	: TSchema extends AllowedTypesUnsafe
+		? NodeFromSchemaUnsafe<FlexListToUnion<TSchema>>
+		: unknown;
 
 /**
  * {@link Unenforced} version of {@link InsertableTreeNodeFromImplicitAllowedTypes}.
@@ -97,7 +172,7 @@ export type TreeNodeFromImplicitAllowedTypesUnsafe<
  */
 export type InsertableTreeNodeFromImplicitAllowedTypesUnsafe<
 	TSchema extends Unenforced<ImplicitAllowedTypes>,
-> = TSchema extends AllowedTypes
+> = TSchema extends AllowedTypesUnsafe
 	? InsertableTypedNodeUnsafe<FlexListToUnion<TSchema>>
 	: InsertableTypedNodeUnsafe<TSchema>;
 
@@ -119,7 +194,7 @@ export type InsertableTypedNodeUnsafe<T extends Unenforced<TreeNodeSchema>> = [
  * @system @public
  */
 export type NodeFromSchemaUnsafe<T extends Unenforced<TreeNodeSchema>> =
-	T extends TreeNodeSchema<string, NodeKind, infer TNode> ? TNode : never;
+	T extends TreeNodeSchemaUnsafe<string, NodeKind, infer TNode> ? TNode : never;
 
 /**
  * {@link Unenforced} version of {@link InsertableTreeNodeFromImplicitAllowedTypes}.
@@ -128,7 +203,7 @@ export type NodeFromSchemaUnsafe<T extends Unenforced<TreeNodeSchema>> =
  * @system @public
  */
 export type NodeBuilderDataUnsafe<T extends Unenforced<TreeNodeSchema>> =
-	T extends TreeNodeSchema<string, NodeKind, unknown, infer TBuild> ? TBuild : never;
+	T extends TreeNodeSchemaUnsafe<string, NodeKind, unknown, infer TBuild> ? TBuild : never;
 
 /**
  * {@link Unenforced} version of {@link (TreeArrayNode:interface)}.
@@ -139,8 +214,7 @@ export type NodeBuilderDataUnsafe<T extends Unenforced<TreeNodeSchema>> =
 export interface TreeArrayNodeUnsafe<TAllowedTypes extends Unenforced<ImplicitAllowedTypes>>
 	extends TreeArrayNodeBase<
 		TreeNodeFromImplicitAllowedTypesUnsafe<TAllowedTypes>,
-		InsertableTreeNodeFromImplicitAllowedTypesUnsafe<TAllowedTypes>,
-		TreeArrayNode
+		InsertableTreeNodeFromImplicitAllowedTypesUnsafe<TAllowedTypes>
 	> {}
 
 /**

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

@@ -10,10 +10,7 @@ import { assert } from "@fluidframework/core-utils/internal";
 import {
 	aboveRootPlaceholder,
 	EmptyKey,
-	forEachField,
-	inCursorField,
 	keyAsDetachedField,
-	mapCursorField,
 	type FieldKey,
 	type ITreeCursor,
 	type ITreeCursorSynchronous,
@@ -40,16 +37,22 @@ import {
 	stringSchema,
 } from "../leafNodeSchema.js";
 import { isObjectNodeSchema } from "../objectNodeTypes.js";
-import { walkFieldSchema } from "../walkFieldSchema.js";
+import {
+	customFromCursorInner,
+	type CustomTreeNode,
+	type CustomTreeValue,
+	type EncodeOptions,
+} from "./customTree.js";
 import { getUnhydratedContext } from "../createContext.js";
 
 /**
- * Verbose encoding of a {@link TreeNode} or {@link TreeValue}.
+ * Verbose encoding of a {@link TreeNode} or {@link TreeLeafValue}.
  * @remarks
  * This is verbose meaning that every {@link TreeNode} is a {@link VerboseTreeNode}.
  * Any IFluidHandle values have been replaced by `THandle`.
  * @privateRemarks
- * This can store all possible simple trees, but it can not store all possible trees representable by our internal representations like FlexTree and JsonableTree.
+ * This can store all possible simple trees,
+ * but it can not store all possible trees representable by our internal representations like FlexTree and JsonableTree.
  */
 export type VerboseTree<THandle = IFluidHandle> =
 	| VerboseTreeNode<THandle>
@@ -77,7 +80,8 @@ export type VerboseTree<THandle = IFluidHandle> =
  * This format allows for all simple-tree compatible trees to be represented.
  *
  * Unlike `JsonableTree`, leaf nodes are not boxed into node objects, and instead have their schema inferred from the value.
- * Additionally, sequence fields can only occur on a node that has a single sequence field (with the empty key) replicating the behavior of simple-tree ArrayNodes.
+ * Additionally, sequence fields can only occur on a node that has a single sequence field (with the empty key)
+ * replicating the behavior of simple-tree ArrayNodes.
  */
 export interface VerboseTreeNode<THandle = IFluidHandle> {
 	/**
@@ -144,24 +148,6 @@ export interface SchemalessParseOptions<TCustom> {
 	};
 }
 
-/**
- * Options for how to interpret a `VerboseTree<TCustom>` without relying on schema.
- */
-export interface EncodeOptions<TCustom> {
-	/**
-	 * Fixup custom input formats.
-	 * @remarks
-	 * See note on {@link ParseOptions.valueConverter}.
-	 */
-	valueConverter(data: IFluidHandle): TCustom;
-	/**
-	 * If true, interpret the input keys of object nodes as stored keys.
-	 * If false, interpret them as property keys.
-	 * @defaultValue false.
-	 */
-	readonly useStoredKeys?: boolean;
-}
-
 /**
  * Use info from `schema` to convert `options` to {@link SchemalessParseOptions}.
  */
@@ -345,12 +331,7 @@ export function verboseFromCursor<TCustom>(
 		...options,
 	};
 
-	const schemaMap = new Map<string, TreeNodeSchema>();
-	walkFieldSchema(rootSchema, {
-		node(schema) {
-			schemaMap.set(schema.identifier, schema);
-		},
-	});
+	const schemaMap = getUnhydratedContext(rootSchema).schema;
 
 	return verboseFromCursorInner(reader, config, schemaMap);
 }
@@ -360,50 +341,14 @@ function verboseFromCursorInner<TCustom>(
 	options: Required<EncodeOptions<TCustom>>,
 	schema: ReadonlyMap<string, TreeNodeSchema>,
 ): VerboseTree<TCustom> {
-	const type = reader.type;
-	const nodeSchema = schema.get(type) ?? fail("missing schema for type in cursor");
-
-	switch (type) {
-		case numberSchema.identifier:
-		case booleanSchema.identifier:
-		case nullSchema.identifier:
-		case stringSchema.identifier:
-			assert(reader.value !== undefined, 0xa14 /* out of schema: missing value */);
-			assert(!isFluidHandle(reader.value), 0xa15 /* out of schema: unexpected FluidHandle */);
-			return reader.value;
-		case handleSchema.identifier:
-			assert(reader.value !== undefined, 0xa16 /* out of schema: missing value */);
-			assert(isFluidHandle(reader.value), 0xa17 /* out of schema: unexpected FluidHandle */);
-			return options.valueConverter(reader.value);
-		default: {
-			assert(reader.value === undefined, 0xa18 /* out of schema: unexpected value */);
-			if (nodeSchema.kind === NodeKind.Array) {
-				const fields = inCursorField(reader, EmptyKey, () =>
-					mapCursorField(reader, () => verboseFromCursorInner(reader, options, schema)),
-				);
-				return { type, fields };
-			} else {
-				const fields: Record<string, VerboseTree<TCustom>> = {};
-				forEachField(reader, () => {
-					const children = mapCursorField(reader, () =>
-						verboseFromCursorInner(reader, options, schema),
-					);
-					if (children.length === 1) {
-						const storedKey = reader.getFieldKey();
-						const key =
-							isObjectNodeSchema(nodeSchema) && !options.useStoredKeys
-								? nodeSchema.storedKeyToPropertyKey.get(storedKey) ??
-									fail("missing property key")
-								: storedKey;
-						// Length is checked above.
-						// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-						fields[key] = children[0]!;
-					} else {
-						assert(children.length === 0, 0xa19 /* invalid children number */);
-					}
-				});
-				return { type, fields };
-			}
-		}
+	const fields = customFromCursorInner(reader, options, schema, verboseFromCursorInner);
+	const nodeSchema = schema.get(reader.type) ?? fail("missing schema for type in cursor");
+	if (nodeSchema.kind === NodeKind.Leaf) {
+		return fields as CustomTreeValue<TCustom>;
 	}
+
+	return {
+		type: reader.type,
+		fields: fields as CustomTreeNode<TCustom>,
+	};
 }

package/src/simple-tree/arrayNode.ts

@@ -14,11 +14,11 @@ import {
 	isFlexTreeNode,
 } from "../feature-libraries/index.js";
 import { prepareContentForHydration } from "./proxies.js";
-import { getOrCreateInnerNode } from "./proxyBinding.js";
 import {
 	normalizeAllowedTypes,
 	type ImplicitAllowedTypes,
 	type InsertableTreeNodeFromImplicitAllowedTypes,
+	type TreeLeafValue,
 	type TreeNodeFromImplicitAllowedTypes,
 } from "./schemaTypes.js";
 import {
@@ -34,6 +34,8 @@ import {
 	getOrCreateNodeFromInnerNode,
 	type TreeNodeSchemaBoth,
 	getSimpleNodeSchemaFromInnerNode,
+	getOrCreateInnerNode,
+	type TreeNodeSchemaClass,
 } from "./core/index.js";
 import { type InsertableContent, mapTreeFromNodeData } from "./toMapTree.js";
 import { fail } from "../util/index.js";
@@ -45,6 +47,19 @@ import {
 import { TreeNodeValid, type MostDerivedData } from "./treeNodeValid.js";
 import { getUnhydratedContext } from "./createContext.js";
 
+/**
+ * A covariant base type for {@link (TreeArrayNode:interface)}.
+ *
+ * This provides the readonly subset of TreeArrayNode functionality, and is used as the source interface for moves since that needs to be covariant.
+ * @privateRemarks
+ * Ideally this would just include `TreeNode, WithType<string, NodeKind.Array>` in the extends list but https://github.com/microsoft/TypeScript/issues/16936 prevents that from compiling.
+ * As a workaround around for this TypeScript limitation, the conflicting type intersection is wrapped in `Awaited` (which has no effect on the type in this case) which allows it to compile.
+ * @system @sealed @public
+ */
+export interface ReadonlyArrayNode<out T = TreeNode | TreeLeafValue>
+	extends ReadonlyArray<T>,
+		Awaited<TreeNode & WithType<string, NodeKind.Array>> {}
+
 /**
  * A generic array type, used to defined types like {@link (TreeArrayNode:interface)}.
  *
@@ -53,9 +68,8 @@ import { getUnhydratedContext } from "./createContext.js";
  *
  * @system @sealed @public
  */
-export interface TreeArrayNodeBase<out T, in TNew, in TMoveFrom>
-	extends ReadonlyArray<T>,
-		TreeNode {
+export interface TreeArrayNodeBase<out T, in TNew, in TMoveFrom = ReadonlyArrayNode>
+	extends ReadonlyArrayNode<T> {
 	/**
 	 * Inserts new item(s) at a specified location.
 	 * @param index - The index at which to insert `value`.
@@ -366,8 +380,7 @@ export interface TreeArrayNode<
 	TAllowedTypes extends ImplicitAllowedTypes = ImplicitAllowedTypes,
 > extends TreeArrayNodeBase<
 		TreeNodeFromImplicitAllowedTypes<TAllowedTypes>,
-		InsertableTreeNodeFromImplicitAllowedTypes<TAllowedTypes>,
-		TreeArrayNode
+		InsertableTreeNodeFromImplicitAllowedTypes<TAllowedTypes>
 	> {}
 
 /**
@@ -417,9 +430,7 @@ export class IterableTreeArrayContent<T> implements Iterable<T> {
 /**
  * Given a array node proxy, returns its underlying LazySequence field.
  */
-function getSequenceField<TSimpleType extends ImplicitAllowedTypes>(
-	arrayNode: TreeArrayNode<TSimpleType>,
-): FlexTreeSequenceField {
+function getSequenceField(arrayNode: ReadonlyArrayNode): FlexTreeSequenceField {
 	return getOrCreateInnerNode(arrayNode).getBoxed(EmptyKey) as FlexTreeSequenceField;
 }
 
@@ -799,6 +810,11 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 	protected abstract get simpleSchema(): T;
 	protected abstract get allowedTypes(): ReadonlySet<TreeNodeSchema>;
 
+	public abstract override get [typeSchemaSymbol](): TreeNodeSchemaClass<
+		string,
+		NodeKind.Array
+	>;
+
 	public constructor(
 		input: Iterable<InsertableTreeNodeFromImplicitAllowedTypes<T>> | InternalTreeNode,
 	) {
@@ -901,13 +917,13 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		}
 		field.editor.remove(removeStart, removeEnd - removeStart);
 	}
-	public moveToStart(sourceIndex: number, source?: TreeArrayNode): void {
+	public moveToStart(sourceIndex: number, source?: ReadonlyArrayNode): void {
 		const sourceArray = source ?? this;
 		const sourceField = getSequenceField(sourceArray);
 		validateIndex(sourceIndex, sourceField, "moveToStart");
 		this.moveRangeToIndex(0, sourceIndex, sourceIndex + 1, source);
 	}
-	public moveToEnd(sourceIndex: number, source?: TreeArrayNode): void {
+	public moveToEnd(sourceIndex: number, source?: ReadonlyArrayNode): void {
 		const sourceArray = source ?? this;
 		const sourceField = getSequenceField(sourceArray);
 		validateIndex(sourceIndex, sourceField, "moveToEnd");
@@ -916,7 +932,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 	public moveToIndex(
 		destinationGap: number,
 		sourceIndex: number,
-		source?: TreeArrayNode,
+		source?: ReadonlyArrayNode,
 	): void {
 		const sourceArray = source ?? this;
 		const sourceField = getSequenceField(sourceArray);
@@ -928,7 +944,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 	public moveRangeToStart(
 		sourceStart: number,
 		sourceEnd: number,
-		source?: TreeArrayNode,
+		source?: ReadonlyArrayNode,
 	): void {
 		validateIndexRange(
 			sourceStart,
@@ -938,7 +954,11 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		);
 		this.moveRangeToIndex(0, sourceStart, sourceEnd, source);
 	}
-	public moveRangeToEnd(sourceStart: number, sourceEnd: number, source?: TreeArrayNode): void {
+	public moveRangeToEnd(
+		sourceStart: number,
+		sourceEnd: number,
+		source?: ReadonlyArrayNode,
+	): void {
 		validateIndexRange(
 			sourceStart,
 			sourceEnd,
@@ -951,7 +971,7 @@ abstract class CustomArrayNodeBase<const T extends ImplicitAllowedTypes>
 		destinationGap: number,
 		sourceStart: number,
 		sourceEnd: number,
-		source?: TreeArrayNode,
+		source?: ReadonlyArrayNode,
 	): void {
 		const destinationField = getSequenceField(this);
 		const destinationSchema = this.allowedTypes;

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

@@ -11,6 +11,7 @@ export {
 	type InnerNode,
 	tryDisposeTreeNode,
 	tryGetTreeNodeFromMapNode,
+	getOrCreateInnerNode,
 } from "./treeNodeKernel.js";
 export { type WithType, typeNameSymbol, typeSchemaSymbol } from "./withType.js";
 export {

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

@@ -430,3 +430,16 @@ export function getTreeNodeSchemaFromHydratedFlexNode(flexNode: FlexTreeNode): T
 
 	return context.schema.get(flexNode.schema) ?? fail("Missing schema");
 }
+
+/**
+ * Retrieves the flex node associated with the given target via {@link setInnerNode}.
+ * @remarks
+ * For {@link Unhydrated} nodes, this returns the MapTreeNode.
+ *
+ * For hydrated nodes it returns a FlexTreeNode backed by the forest.
+ * Note that for "marinated" nodes, this FlexTreeNode exists and returns it: it does not return the MapTreeNode which is the current InnerNode.
+ */
+export function getOrCreateInnerNode(treeNode: TreeNode, allowFreed = false): InnerNode {
+	const kernel = getKernel(treeNode);
+	return kernel.getOrCreateInnerNode(allowFreed);
+}

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

@@ -3,7 +3,8 @@
  * Licensed under the MIT License.
  */
 
-import type { InternalTreeNode, Unhydrated } from "./types.js";
+import type { TreeLeafValue } from "../schemaTypes.js";
+import type { InternalTreeNode, TreeNode, Unhydrated } from "./types.js";
 
 /**
  * Schema for a tree node.
@@ -19,7 +20,7 @@ import type { InternalTreeNode, Unhydrated } from "./types.js";
 export type TreeNodeSchema<
 	Name extends string = string,
 	Kind extends NodeKind = NodeKind,
-	TNode = unknown,
+	TNode extends TreeNode | TreeLeafValue = TreeNode | TreeLeafValue,
 	TBuild = never,
 	ImplicitlyConstructable extends boolean = boolean,
 	Info = unknown,
@@ -38,7 +39,7 @@ export type TreeNodeSchema<
 export interface TreeNodeSchemaNonClass<
 	out Name extends string = string,
 	out Kind extends NodeKind = NodeKind,
-	out TNode = unknown,
+	out TNode extends TreeNode | TreeLeafValue = TreeNode | TreeLeafValue,
 	in TInsertable = never,
 	out ImplicitlyConstructable extends boolean = boolean,
 	out Info = unknown,
@@ -94,7 +95,8 @@ export interface TreeNodeSchemaNonClass<
 export interface TreeNodeSchemaClass<
 	out Name extends string = string,
 	out Kind extends NodeKind = NodeKind,
-	out TNode = unknown,
+	// TODO: maybe this can be more specific (exclude leaves)
+	out TNode extends TreeNode | TreeLeafValue = TreeNode | TreeLeafValue,
 	in TInsertable = never,
 	out ImplicitlyConstructable extends boolean = boolean,
 	out Info = unknown,
@@ -115,7 +117,7 @@ export interface TreeNodeSchemaClass<
 export type TreeNodeSchemaBoth<
 	Name extends string = string,
 	Kind extends NodeKind = NodeKind,
-	TNode = unknown,
+	TNode extends TreeNode = TreeNode,
 	TInsertable = never,
 	ImplicitlyConstructable extends boolean = boolean,
 	Info = unknown,

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

@@ -442,7 +442,12 @@ class EagerMapTreeRequiredField
 	implements FlexTreeRequiredField
 {
 	public override get content(): FlexTreeUnknownUnboxed {
-		return super.content ?? fail("Expected EagerMapTree required field to have a value");
+		// This cannot use ?? since null is a legal value here.
+		assert(
+			super.content !== undefined,
+			"Expected EagerMapTree required field to have a value",
+		);
+		return super.content;
 	}
 }
 

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

@@ -3,7 +3,9 @@
  * Licensed under the MIT License.
  */
 
+import type { TreeLeafValue } from "../schemaTypes.js";
 import type { NodeKind, TreeNodeSchemaClass } from "./treeNodeSchema.js";
+import type { TreeNode } from "./types.js";
 
 /**
  * The type of a {@link TreeNode}.
@@ -87,5 +89,12 @@ export interface WithType<
 	/**
 	 * Type symbol, marking a type in a way to increase type safety via strong type checking.
 	 */
-	get [typeSchemaSymbol](): TreeNodeSchemaClass<TName, TKind, unknown, never, boolean, TInfo>;
+	get [typeSchemaSymbol](): TreeNodeSchemaClass<
+		TName,
+		TKind,
+		TreeNode | TreeLeafValue,
+		never,
+		boolean,
+		TInfo
+	>;
 }