Rename maybeDecompressedStringAsNumber and propagate mayContainCompressedIds to non-leaf nodes (#26980)

## Description

Improves clarity around how compressed ids are handled in UniformChunks,
and tracks if a chunk might use them (deeply) in the shape.

As a memory optimization, avoid storing idCompressor when unneeded. For
small leaf chunks (as is common currently) this should help memory use a
little as there are a very large number of such chunks.

Author: CraigMacomber

packages/dds/tree/src/feature-libraries/chunked-forest/chunkTree.ts

@@ -424,8 +424,8 @@ export interface ChunkCompressor {
 	/**
 	 * If the idCompressor is provided, {@link UniformChunk}s with identifiers will be encoded for its in-memory representation.
 	 * @remarks
-	 * This compression applies to {@link UniformChunk}s when {@link TreeShape.maybeDecompressedStringAsNumber} is set.
-	 * If the `policy` does not use UniformChunks or does not set `maybeDecompressedStringAsNumber`, then no compression will be applied even when providing `idCompressor`.
+	 * This compression applies to {@link UniformChunk}s when {@link TreeShape.mayContainCompressedIds} is set.
+	 * If the `policy` does not use UniformChunks or does not set `mayContainCompressedIds`, then no compression will be applied even when providing `idCompressor`.
 	 */
 	readonly idCompressor: IIdCompressor | undefined;
 }
@@ -583,6 +583,7 @@ export function insertValues(
 	// Slow path: walk shape and cursor together, inserting values.
 	if (shape.hasValue) {
 		if (
+			shape.mayContainCompressedIds &&
 			typeof cursor.value === "string" &&
 			idCompressor !== undefined &&
 			isStableNodeIdentifier(cursor.value)

packages/dds/tree/src/feature-libraries/chunked-forest/uniformChunk.ts

@@ -5,7 +5,6 @@
 
 import { assert, compareArrays, oob, fail } from "@fluidframework/core-utils/internal";
 import type { SessionSpaceCompressedId, IIdCompressor } from "@fluidframework/id-compressor";
-import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import {
 	CursorLocationType,
@@ -54,7 +53,7 @@ export class UniformChunk extends ReferenceCountedBase implements TreeChunk {
 		idCompressor?: IIdCompressor,
 	) {
 		super();
-		this.idCompressor = idCompressor;
+		this.idCompressor = shape.treeShape.mayContainCompressedIds ? idCompressor : undefined;
 		assert(
 			shape.treeShape.valuesPerTopLevelNode * shape.topLevelLength === values.length,
 			0x4c3 /* invalid number of values for shape */,
@@ -100,30 +99,42 @@ export class TreeShape {
 	// TODO: this is only needed at chunk roots. Optimize it base on that.
 	public readonly positions: readonly NodePositionInfo[];
 
+	/**
+	 * Whether chunks using this shape (including any descendant leaf within it) may contain values compressed by the {@link UniformChunk.idCompressor}.
+	 *
+	 * @remarks
+	 * For string leaf nodes, this can be explicitly set to `true` to indicate that the value may be a compressed id
+	 * stored as a number that needs to be decompressed back to a string.
+	 * For non-leaf nodes, this is automatically derived from whether any child shapes have it set.
+	 */
+	public readonly mayContainCompressedIds: boolean;
+
 	/**
 	 *
 	 * @param type - {@link TreeNodeSchemaIdentifier} used to compare shapes.
 	 * @param hasValue - whether or not the TreeShape has a value.
 	 * @param fieldsArray - an array of {@link FieldShape} values, which contains a TreeShape for each FieldKey.
 	 *
-	 * @param maybeDecompressedStringAsNumber - used to check whether or not the value could have been compressed by the idCompressor.
-	 * This flag can only be set on string leaf nodes, and will throw a usage error otherwise.
-	 * If set to true, an additional check can be made (example: getting the value of {@link Cursor}) to return the original uncompressed value.
+	 * @param maybeCompressedIdLeaf - whether the value may have been compressed by the {@link UniformChunk.idCompressor}.
+	 * Can only be explicitly set to `true` on string leaf nodes; otherwise this constructor asserts.
+	 * For non-leaf nodes, {@link TreeShape.mayContainCompressedIds} is automatically derived from child shapes.
 	 */
 	public constructor(
 		public readonly type: TreeNodeSchemaIdentifier,
 		public readonly hasValue: boolean,
 		public readonly fieldsArray: readonly FieldShape[],
-		public readonly maybeDecompressedStringAsNumber: boolean = false,
+		maybeCompressedIdLeaf: boolean = false,
 	) {
-		if (
-			maybeDecompressedStringAsNumber &&
-			!(hasValue && type === "com.fluidframework.leaf.string")
-		) {
-			throw new UsageError(
-				"maybeDecompressedStringAsNumber flag can only be set to true for string leaf node.",
+		assert(hasValue === false || fieldsArray.length === 0, "only non-leaf can have fields");
+		if (maybeCompressedIdLeaf) {
+			assert(
+				hasValue && type === "com.fluidframework.leaf.string",
+				"only strings can opt into maybeCompressedIdLeaf",
 			);
 		}
+		// For non-leaf nodes, derive from whether any child shapes contain compressed ids.
+		this.mayContainCompressedIds =
+			maybeCompressedIdLeaf || fieldsArray.some(([, shape]) => shape.mayContainCompressedIds);
 		const fields: Map<FieldKey, OffsetShape> = new Map();
 		let numberOfValues = hasValue ? 1 : 0;
 		const infos: NodePositionInfo[] = [
@@ -146,7 +157,7 @@ export class TreeShape {
 	}
 
 	public equals(other: TreeShape): boolean {
-		// TODO: either dedup instances and/or store a collision resistant hash for fast compare.
+		// TODO: either dedupe instances and/or store a collision resistant hash for fast compare.
 
 		if (
 			!compareArrays(
@@ -157,7 +168,11 @@ export class TreeShape {
 		) {
 			return false;
 		}
-		return this.type === other.type && this.hasValue === other.hasValue;
+		return (
+			this.type === other.type &&
+			this.hasValue === other.hasValue &&
+			this.mayContainCompressedIds === other.mayContainCompressedIds
+		);
 	}
 
 	public withTopLevelLength(topLevelLength: number): ChunkShape {
@@ -551,9 +566,8 @@ class Cursor extends SynchronousCursor implements ChunkedCursor {
 		const info = this.nodeInfo(CursorLocationType.Nodes);
 		if (info.shape.hasValue) {
 			const value = this.chunk.values[info.valueOffset];
-			// If the maybeDecompressedStringAsNumber flag is set to true, we check if the value is a number.
-			// This flag can only ever be set on string leaf nodes, so if the value is a number, we can assume it is a compressible, known stable id.
-			if (info.shape.maybeDecompressedStringAsNumber && typeof value === "number") {
+			// If mayContainCompressedIds is set, check if the value is a number (i.e. a compressed ID that needs decompression).
+			if (info.shape.mayContainCompressedIds && typeof value === "number") {
 				const idCompressor = this.chunk.idCompressor;
 				assert(
 					idCompressor !== undefined,

packages/dds/tree/src/test/feature-libraries/chunked-forest/chunkTree.spec.ts

@@ -104,6 +104,52 @@ describe("chunkTree", () => {
 				assert.deepEqual(values, chunk.values);
 			});
 		}
+
+		it("does not compress string values when shape does not have mayContainCompressedIds", () => {
+			// Simulate a string leaf whose shape does not opt into compressed ids
+			// (e.g. a plain string field in a string | number union).
+			// Even if the string value is a valid stable id and an idCompressor is provided,
+			// insertValues must leave it as a string.
+			const stableId = testIdCompressor.decompress(testIdCompressor.generateCompressedId());
+			const stringShapeNoCompress = new TreeShape(
+				brand(stringSchema.identifier),
+				true,
+				[],
+				false,
+			);
+
+			const values: Value[] = [];
+			const cursor = cursorForJsonableTreeNode({
+				type: brand(stringSchema.identifier),
+				value: stableId,
+			});
+			insertValues(cursor, stringShapeNoCompress, values, testIdCompressor);
+			// The value must remain the original string, not a compressed numeric id.
+			assert.equal(values.length, 1);
+			assert.equal(typeof values[0], "string");
+			assert.equal(values[0], stableId);
+		});
+
+		it("compresses string values when shape has mayContainCompressedIds", () => {
+			const compressedId = testIdCompressor.generateCompressedId();
+			const stableId = testIdCompressor.decompress(compressedId);
+			const stringShapeCompress = new TreeShape(
+				brand(stringSchema.identifier),
+				true,
+				[],
+				true,
+			);
+
+			const values: Value[] = [];
+			const cursor = cursorForJsonableTreeNode({
+				type: brand(stringSchema.identifier),
+				value: stableId,
+			});
+			insertValues(cursor, stringShapeCompress, values, testIdCompressor);
+			// The value must be compressed to the numeric id.
+			assert.equal(values.length, 1);
+			assert.equal(values[0], compressedId);
+		});
 	});
 
 	describe("uniformChunkFromCursor", () => {

packages/dds/tree/src/test/feature-libraries/chunked-forest/uniformChunk.spec.ts

@@ -6,7 +6,7 @@
 import { strict as assert } from "node:assert";
 
 import { BenchmarkType, benchmark } from "@fluid-tools/benchmark";
-import { validateUsageError } from "@fluidframework/test-runtime-utils/internal";
+import { validateAssertionError } from "@fluidframework/test-runtime-utils/internal";
 
 import { EmptyKey, type ITreeCursorSynchronous } from "../../../core/index.js";
 import {
@@ -66,16 +66,62 @@ describe("uniformChunk", () => {
 				validateShape(tree.dataFactory().shape);
 			});
 		}
-		it("shape with maybeCompressedStringAsNumber flag set to true fails if it is not a string leaf node.", () => {
+		it("shape with mayContainCompressedIds flag set to true fails if it is not a string leaf node.", () => {
 			const validShapeWithFlag = new TreeShape(brand(stringSchema.identifier), true, [], true);
-			// Test that a non string leaf node shape with maybeCompressedStringAsNumber set to true fails.
+			// Test that a non string leaf node shape with mayContainCompressedIds set to true fails.
 			assert.throws(
 				() => new TreeShape(brand(numberSchema.identifier), true, [], true),
-				validateUsageError(
-					/maybeDecompressedStringAsNumber flag can only be set to true for string leaf node./,
-				),
+				validateAssertionError("only strings can opt into maybeCompressedIdLeaf"),
 			);
 		});
+
+		it("equals distinguishes shapes differing only by mayContainCompressedIds", () => {
+			const withIds = new TreeShape(brand(stringSchema.identifier), true, [], true);
+			const withoutIds = new TreeShape(brand(stringSchema.identifier), true, [], false);
+			assert.equal(withIds.equals(withoutIds), false);
+			assert.equal(withoutIds.equals(withIds), false);
+			// Self-equality still holds
+			assert.equal(withIds.equals(withIds), true);
+			assert.equal(withoutIds.equals(withoutIds), true);
+		});
+
+		it("mayContainCompressedIds propagates from child shapes to parent shapes", () => {
+			const leafWithIds = new TreeShape(brand(stringSchema.identifier), true, [], true);
+			const leafWithoutIds = new TreeShape(brand(stringSchema.identifier), true, [], false);
+
+			// Parent with a child that has `mayContainCompressedIds` should also have it.
+			const parentWithCompressedChild = new TreeShape(
+				brand(JsonAsTree.JsonObject.identifier),
+				false,
+				[[xField, leafWithIds, 1]],
+			);
+			assert.equal(parentWithCompressedChild.mayContainCompressedIds, true);
+
+			// Parent with no children that have `mayContainCompressedIds` should not have it.
+			const parentWithoutCompressedChild = new TreeShape(
+				brand(JsonAsTree.JsonObject.identifier),
+				false,
+				[[xField, leafWithoutIds, 1]],
+			);
+			assert.equal(parentWithoutCompressedChild.mayContainCompressedIds, false);
+
+			// Propagation through multiple levels: grandparent should inherit from grandchild.
+			const grandparent = new TreeShape(brand(JsonAsTree.Array.identifier), false, [
+				[EmptyKey, parentWithCompressedChild, 1],
+			]);
+			assert.equal(grandparent.mayContainCompressedIds, true);
+
+			// Mixed children: one with and one without compressed IDs.
+			const parentWithMixedChildren = new TreeShape(
+				brand(JsonAsTree.JsonObject.identifier),
+				false,
+				[
+					[xField, leafWithoutIds, 1],
+					[yField, leafWithIds, 1],
+				],
+			);
+			assert.equal(parentWithMixedChildren.mayContainCompressedIds, true);
+		});
 	});
 
 	testSpecializedFieldCursor<TreeChunk, ITreeCursorSynchronous>({