(tree) Add incremental summary optimzation opt-in to AnnotatedAllowedTypes (#25408)

## Description
With this PR, a set of allowed types in a schema can be opted in to
incremental summary optimization. These allowed types will be optimized
during summary such that if they don't change across summaries, they
will not be encoded and their content will not be included in the
summary that is uploaded to the service. The usage pattern is described
below.

## Usage pattern
- Incremental summarization only works for forest type
`ForestTypeOptimized` when the compression strategy
is `TreeCompressionStrategyExtended.CompressedIncremental`. In addition,
`SharedTreeOptionsInternal.shouldEncodeFieldIncrementally` must be
passed when creating the tree. This callback function will be called for
each allowed types in the schema to determine if it should be
incrementally summarized. These configs can be passed in when creating a
tree via the `configuredSharedTree` API.
- A set of allowed types can be opted in to incremental summary
optimization by adding an `incrementalSummaryHint`symbol as true to the
`AllowedTypesMetadata.custom` property.
- The helper function `shouldIncrementallySummarizeAllowedTypes` (added
in this PR) can be used to implement the
`SharedTreeOptionsInternal.shouldEncodeFieldIncrementally` callback
function. It takes in the schema that contains the allowed types and the
node identifier and field key of the target allowed types. The last two
parameters are the same as the ones in the callback function. It will
find the `AllowedTypesMetadata.custom` for the passed in allowed types
and determine if the `incrementalSummaryHint` symbol is set to true.


[AB#41866](https://dev.azure.com/fluidframework/235294da-091d-4c29-84fc-cdfc3d90890b/_workitems/edit/41866)

Author: agarwal-navin

packages/dds/tree/src/core/forest/forest.ts

@@ -84,15 +84,14 @@ export interface IForestSubscription {
 	clone(schema: TreeStoredSchemaSubscription, anchors: AnchorSet): IEditableForest;
 
 	/**
-	 * Generate a TreeChunk for the content in the given field cursor.
+	 * Generate a TreeChunk[] for the current field (and its children) of cursor.
 	 * This can be used to chunk data that is then inserted into the forest.
 	 *
 	 * @remarks
-	 * Like {@link chunkField}, but forces the results into a single TreeChunk.
-	 * While any TreeChunk is compatible with any forest, this method creates one optimized for this specific forest.
+	 * Similar to {@link chunkField} but it creates chunks optimized for this specific forest by using its compression policy.
 	 * The provided data must be compatible with the forest's current schema.
 	 */
-	chunkField(cursor: ITreeCursorSynchronous): TreeChunk;
+	chunkField(cursor: ITreeCursorSynchronous): TreeChunk[];
 
 	/**
 	 * Allocates a cursor in the "cleared" state.

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

@@ -197,6 +197,15 @@ export function chunkFieldSingle(
 	policy: ChunkCompressor,
 ): TreeChunk {
 	const chunks = chunkField(cursor, policy);
+	return combineChunks(chunks);
+}
+
+/**
+ * Create a single TreeChunk from an array of TreeChunks.
+ * @remarks
+ * This takes ownership of the provided TreeChunk references, and returns an owned referenced.
+ */
+export function combineChunks(chunks: TreeChunk[]): TreeChunk {
 	if (chunks.length === 1) {
 		return chunks[0] ?? oob();
 	}

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

@@ -44,7 +44,7 @@ import {
 } from "../../util/index.js";
 
 import { BasicChunk, BasicChunkCursor, type SiblingsOrKey } from "./basicChunk.js";
-import { type IChunker, basicChunkTree, chunkFieldSingle, chunkTree } from "./chunkTree.js";
+import { type IChunker, basicChunkTree, chunkField, chunkTree } from "./chunkTree.js";
 
 function makeRoot(): BasicChunk {
 	return new BasicChunk(aboveRootPlaceholder, new Map());
@@ -90,8 +90,8 @@ export class ChunkedForest implements IEditableForest {
 		return new ChunkedForest(this.roots, schema, this.chunker.clone(schema), anchors);
 	}
 
-	public chunkField(cursor: ITreeCursorSynchronous): TreeChunk {
-		return chunkFieldSingle(cursor, { idCompressor: this.idCompressor, policy: this.chunker });
+	public chunkField(cursor: ITreeCursorSynchronous): TreeChunk[] {
+		return chunkField(cursor, { idCompressor: this.idCompressor, policy: this.chunker });
 	}
 
 	public forgetAnchor(anchor: Anchor): void {

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

@@ -36,7 +36,7 @@ import {
 	SpecialField,
 	version,
 } from "./format.js";
-import type { ChunkReferenceId, IncrementalEncoder } from "./codecs.js";
+import type { IncrementalEncoder } from "./codecs.js";
 
 /**
  * Encode data from `FieldBatch` into an `EncodedFieldBatch`.
@@ -223,8 +223,8 @@ export const anyNodeEncoder: NodeEncoder = {
 		outputBuffer: BufferFormat,
 	): void {
 		// TODO: Fast path uniform chunk content.
-		const shape = context.nodeEncoderFromSchema(cursor.type);
-		AnyShape.encodeNode(cursor, context, outputBuffer, shape);
+		const nodeEncoder = context.nodeEncoderFromSchema(cursor.type);
+		AnyShape.encodeNode(cursor, context, outputBuffer, nodeEncoder);
 	},
 
 	shape: AnyShape.instance,
@@ -350,7 +350,7 @@ export class InlineArrayEncoder
 }
 
 /**
- * Encodes the shape for a nested array as {@link EncodedNestedArray} shape.
+ * Encodes the shape for a nested array as {@link EncodedNestedArrayShape} shape.
  */
 export class NestedArrayShape extends ShapeGeneric<EncodedChunkShape> {
 	/**
@@ -420,30 +420,9 @@ export class NestedArrayEncoder implements FieldEncoder {
 }
 
 /**
- * Encodes a chunk with the {@link EncodedIncrementalChunkShape} shape.
- * This chunks will be encoded separately, i.e., the contents of the chunk will not be part of the main buffer.
- * A reference to the chunk will be stored in the main buffer as an {@link ChunkReferenceId}.
+ * Encodes the shape for an incremental chunk as {@link EncodedIncrementalChunkShape} shape.
  */
 export class IncrementalChunkShape extends ShapeGeneric<EncodedChunkShape> {
-	/**
-	 * Encodes all the nodes in the chunk at the cursor position using `InlineArrayShape`.
-	 */
-	public static encodeChunk(chunk: TreeChunk, context: EncoderContext): BufferFormat {
-		const chunkOutputBuffer: BufferFormat = [];
-		const nodesEncoder = asNodesEncoder(anyNodeEncoder);
-		const chunkCursor = chunk.cursor();
-		chunkCursor.firstNode();
-		const chunkLength = chunkCursor.chunkLength;
-		for (let index = 0; index < chunkLength; index++) {
-			nodesEncoder.encodeNodes(chunkCursor, context, chunkOutputBuffer);
-		}
-		assert(
-			chunkCursor.mode === CursorLocationType.Fields,
-			0xc29 /* should return to fields mode when finished encoding */,
-		);
-		return chunkOutputBuffer;
-	}
-
 	public encodeShape(
 		identifiers: DeduplicationTable<string>,
 		shapes: DeduplicationTable<Shape>,
@@ -464,9 +443,10 @@ export class IncrementalChunkShape extends ShapeGeneric<EncodedChunkShape> {
 }
 
 /**
- * Encodes an incremental field whose chunks are encoded separately and referenced by their {@link ChunkReferenceId}.
- * The shape of the content of this field is {@link NestedShape} where the items in the array are
- * the {@link ChunkReferenceId}s of the encoded chunks.
+ * Encodes an incremental field whose tree chunks are encoded separately and referenced by their {@link ChunkReferenceId}.
+ * The shape of the content of this field is {@link NestedArrayShape}.
+ * The inner items of the array have shape {@link IncrementalChunkShape} and are {@link ChunkReferenceId}s
+ * of the encoded chunks.
  */
 export const incrementalFieldEncoder: FieldEncoder = {
 	encodeField(
@@ -475,12 +455,13 @@ export const incrementalFieldEncoder: FieldEncoder = {
 		outputBuffer: BufferFormat,
 	): void {
 		assert(
-			context.shouldEncodeIncrementally,
-			0xc2a /* incremental encoding must be enabled to use IncrementalFieldShape */,
+			context.incrementalEncoder !== undefined,
+			"incremental encoder must be defined to use incrementalFieldEncoder",
 		);
 
-		const chunkReferenceIds = context.encodeIncrementalField(cursor, (chunk: TreeChunk) =>
-			IncrementalChunkShape.encodeChunk(chunk, context),
+		const chunkReferenceIds = context.incrementalEncoder.encodeIncrementalField(
+			cursor,
+			(chunk: TreeChunk) => compressedEncode([chunk.cursor()], context),
 		);
 		outputBuffer.push(chunkReferenceIds);
 	},
@@ -540,7 +521,12 @@ export class EncoderContext implements NodeEncodeBuilder, FieldEncodeBuilder {
 		private readonly fieldEncoderFromPolicy: FieldEncoderPolicy,
 		public readonly fieldShapes: ReadonlyMap<FieldKindIdentifier, FlexFieldKind>,
 		public readonly idCompressor: IIdCompressor,
-		private readonly incrementalEncoder: IncrementalEncoder | undefined,
+		/**
+		 * To be used to encode incremental chunks, if any.
+		 * @remarks
+		 * See {@link IncrementalEncoder} for more information.
+		 */
+		public readonly incrementalEncoder: IncrementalEncoder | undefined,
 	) {}
 
 	public nodeEncoderFromSchema(schemaName: TreeNodeSchemaIdentifier): NodeEncoder {
@@ -556,30 +542,6 @@ export class EncoderContext implements NodeEncodeBuilder, FieldEncodeBuilder {
 	public nestedArrayEncoder(inner: NodeEncoder): NestedArrayEncoder {
 		return getOrCreate(this.nestedArrayEncoders, inner, () => new NestedArrayEncoder(inner));
 	}
-
-	public get shouldEncodeIncrementally(): boolean {
-		return this.incrementalEncoder !== undefined;
-	}
-
-	/**
-	 * {@link IncrementalEncoder.encodeIncrementalField}
-	 */
-	public encodeIncrementalField(
-		cursor: ITreeCursorSynchronous,
-		encoder: (chunk: TreeChunk) => BufferFormat,
-	): ChunkReferenceId[] {
-		assert(
-			this.incrementalEncoder !== undefined,
-			0xc2b /* incremental encoding must be enabled */,
-		);
-		// Encoder for the chunk that encodes its data using the provided encoder function and
-		// updates the encoded data for shapes and identifiers.
-		const chunkEncoder = (chunk: TreeChunk): EncodedFieldBatch => {
-			const chunkOutputBuffer = encoder(chunk);
-			return updateShapesAndIdentifiersEncoding(version, [chunkOutputBuffer]);
-		};
-		return this.incrementalEncoder.encodeIncrementalField(cursor, chunkEncoder);
-	}
 }
 
 export interface NodeEncodeBuilder {

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

@@ -12,6 +12,7 @@ export {
 	type IChunker,
 	chunkFieldSingle,
 	chunkField,
+	combineChunks,
 } from "./chunkTree.js";
 export { buildChunkedForest } from "./chunkedForest.js";
 export {

packages/dds/tree/src/feature-libraries/flex-tree/lazyField.ts

@@ -53,6 +53,7 @@ import {
 import { LazyEntity } from "./lazyEntity.js";
 import { type LazyTreeNode, getOrCreateHydratedFlexTreeNode } from "./lazyNode.js";
 import { indexForAt, treeStatusFromAnchorCache } from "./utilities.js";
+import { combineChunks } from "../chunked-forest/index.js";
 
 /**
  * Reuse fields.
@@ -247,7 +248,8 @@ export abstract class LazyField extends LazyEntity<FieldAnchor> implements FlexT
 	protected getEditor(): IDefaultEditBuilder<ITreeCursorSynchronous> {
 		return new MappedEditBuilder(
 			this.context.checkout.editor,
-			(cursor: ITreeCursorSynchronous) => this.context.checkout.forest.chunkField(cursor),
+			(cursor: ITreeCursorSynchronous) =>
+				combineChunks(this.context.checkout.forest.chunkField(cursor)),
 		);
 	}
 }

packages/dds/tree/src/feature-libraries/forest-summary/incrementalSummaryBuilder.ts

@@ -288,7 +288,7 @@ export class ForestIncrementalSummaryBuilder implements IncrementalEncoderDecode
 
 	public constructor(
 		private readonly enableIncrementalSummary: boolean,
-		private readonly getChunkAtCursor: (cursor: ITreeCursorSynchronous) => TreeChunk,
+		private readonly getChunkAtCursor: (cursor: ITreeCursorSynchronous) => TreeChunk[],
 		public readonly shouldEncodeIncrementally: IncrementalEncodingPolicy,
 		private readonly initialSequenceNumber: number,
 	) {}
@@ -392,80 +392,74 @@ export class ForestIncrementalSummaryBuilder implements IncrementalEncoderDecode
 		// Validate that a summary is currently being tracked and that the tracked summary properties are defined.
 		validateTrackingSummary(this.forestSummaryState, this.trackedSummaryProperties);
 
-		if (cursor.getFieldLength() === 0) {
-			return [];
-		}
-
-		let chunkReferenceId: ChunkReferenceId;
-		let chunkProperties: ChunkSummaryProperties;
+		const chunkReferenceIds: ChunkReferenceId[] = [];
+		const chunks = this.getChunkAtCursor(cursor);
+		for (const chunk of chunks) {
+			let chunkProperties: ChunkSummaryProperties;
+
+			// Try and get the properties of the chunk from the latest successful summary.
+			// If it exists and the summary is not a full tree, use the properties to generate a summary handle.
+			// If it does not exist, encode the chunk and generate new properties for it.
+			const previousChunkProperties = tryGetFromNestedMap(
+				this.chunkTrackingPropertiesMap,
+				this.latestSummarySequenceNumber,
+				chunk,
+			);
+			if (previousChunkProperties !== undefined && !this.trackedSummaryProperties.fullTree) {
+				chunkProperties = previousChunkProperties;
+				this.trackedSummaryProperties.parentSummaryBuilder.addHandle(
+					`${chunkProperties.referenceId}`,
+					SummaryType.Tree,
+					`${this.trackedSummaryProperties.latestSummaryBasePath}/${chunkProperties.summaryPath}`,
+				);
+			} else {
+				// Generate a new reference ID for the chunk.
+				const newReferenceId: ChunkReferenceId = brand(this.nextReferenceId++);
+
+				// Add the reference ID of this chunk to the chunk summary path and use the path as the summary path
+				// for the chunk in its summary properties.
+				// This is done before encoding the chunk so that the summary path is updated correctly when encoding
+				// any incremental chunks that are under this chunk.
+				this.trackedSummaryProperties.chunkSummaryPath.push(newReferenceId);
+
+				chunkProperties = {
+					referenceId: newReferenceId,
+					summaryPath: this.trackedSummaryProperties.chunkSummaryPath.join("/"),
+				};
+
+				const parentSummaryBuilder = this.trackedSummaryProperties.parentSummaryBuilder;
+				// Create a new summary builder for this chunk to build its summary tree which will be stored in the
+				// parent's summary tree under its reference ID.
+				// Before encoding the chunk, set the parent summary builder to this chunk's summary builder so that
+				// any incremental chunks in the subtree of this chunk will use that as their parent summary builder.
+				const chunkSummaryBuilder = new SummaryTreeBuilder();
+				this.trackedSummaryProperties.parentSummaryBuilder = chunkSummaryBuilder;
+				chunkSummaryBuilder.addBlob(
+					chunkContentsBlobKey,
+					this.trackedSummaryProperties.stringify(chunkEncoder(chunk)),
+				);
 
-		// An additional ref-count must be added to these chunks representing a reference from the summary tree to the chunk.
-		// This will ensure that the blob's content never change and thus the reference stays accurate: instead of modifying it,
-		// a copy will be created without the blob reference.
-		// The "getChunkAtCursor" adds this additional ref-count.
-		const chunk = this.getChunkAtCursor(cursor);
+				// Add this chunk's summary tree to the parent's summary tree. The summary tree contains its encoded
+				// contents and the summary trees of any incremental chunks under it.
+				parentSummaryBuilder.addWithStats(
+					`${newReferenceId}`,
+					chunkSummaryBuilder.getSummaryTree(),
+				);
 
-		// Try and get the properties of the chunk from the latest successful summary.
-		// If it exists and the summary is not a full tree, use the properties to generate a summary handle.
-		// If it does not exist, encode the chunk and generate new properties for it.
-		const previousChunkProperties = tryGetFromNestedMap(
-			this.chunkTrackingPropertiesMap,
-			this.latestSummarySequenceNumber,
-			chunk,
-		);
-		if (previousChunkProperties !== undefined && !this.trackedSummaryProperties.fullTree) {
-			chunkProperties = previousChunkProperties;
-			chunkReferenceId = previousChunkProperties.referenceId;
-			this.trackedSummaryProperties.parentSummaryBuilder.addHandle(
-				`${chunkReferenceId}`,
-				SummaryType.Tree,
-				`${this.trackedSummaryProperties.latestSummaryBasePath}/${previousChunkProperties.summaryPath}`,
-			);
-		} else {
-			// Generate a new reference ID for the chunk.
-			chunkReferenceId = brand(this.nextReferenceId++);
-			// Add the reference ID of this chunk to the chunk summary path and use the path as the summary path
-			// for the chunk in its summary properties.
-			// This is done before encoding the chunk so that the summary path is updated correctly when encoding
-			// any incremental chunks that are under this chunk.
-			this.trackedSummaryProperties.chunkSummaryPath.push(chunkReferenceId);
-
-			chunkProperties = {
-				referenceId: chunkReferenceId,
-				summaryPath: this.trackedSummaryProperties.chunkSummaryPath.join("/"),
-			};
-
-			const parentSummaryBuilder = this.trackedSummaryProperties.parentSummaryBuilder;
-			// Create a new summary builder for this chunk to build its summary tree which will be stored in the
-			// parent's summary tree under its reference ID.
-			// Before encoding the chunk, set the parent summary builder to this chunk's summary builder so that
-			// any incremental chunks in the subtree of this chunk will use that as their parent summary builder.
-			const chunkSummaryBuilder = new SummaryTreeBuilder();
-			this.trackedSummaryProperties.parentSummaryBuilder = chunkSummaryBuilder;
-			chunkSummaryBuilder.addBlob(
-				chunkContentsBlobKey,
-				this.trackedSummaryProperties.stringify(chunkEncoder(chunk)),
-			);
+				// Restore the parent summary builder and chunk summary path.
+				this.trackedSummaryProperties.parentSummaryBuilder = parentSummaryBuilder;
+				this.trackedSummaryProperties.chunkSummaryPath.pop();
+			}
 
-			// Add this chunk's summary tree to the parent's summary tree. The summary tree contains its encoded
-			// contents and the summary trees of any incremental chunks under it.
-			parentSummaryBuilder.addWithStats(
-				`${chunkReferenceId}`,
-				chunkSummaryBuilder.getSummaryTree(),
+			setInNestedMap(
+				this.chunkTrackingPropertiesMap,
+				this.trackedSummaryProperties.summarySequenceNumber,
+				chunk,
+				chunkProperties,
 			);
-
-			// Restore the parent summary builder and chunk summary path.
-			this.trackedSummaryProperties.parentSummaryBuilder = parentSummaryBuilder;
-			this.trackedSummaryProperties.chunkSummaryPath.pop();
+			chunkReferenceIds.push(chunkProperties.referenceId);
 		}
-
-		setInNestedMap(
-			this.chunkTrackingPropertiesMap,
-			this.trackedSummaryProperties.summarySequenceNumber,
-			chunk,
-			chunkProperties,
-		);
-		return [chunkReferenceId];
+		return chunkReferenceIds;
 	}
 
 	/**

packages/dds/tree/src/feature-libraries/index.ts

@@ -110,6 +110,7 @@ export {
 	fluidVersionToFieldBatchCodecWriteVersion,
 	type FieldBatchEncodingContext,
 	emptyChunk,
+	combineChunks,
 	type IncrementalEncodingPolicy,
 	defaultIncrementalEncodingPolicy,
 } from "./chunked-forest/index.js";

packages/dds/tree/src/feature-libraries/object-forest/objectForest.ts

@@ -49,7 +49,7 @@ import {
 	type Breakable,
 	type WithBreakable,
 } from "../../util/index.js";
-import { chunkFieldSingle, defaultChunkPolicy } from "../chunked-forest/index.js";
+import { chunkField, defaultChunkPolicy } from "../chunked-forest/index.js";
 import { cursorForMapTreeNode, mapTreeFromCursor } from "../mapTreeCursor.js";
 import { type CursorWithNode, SynchronousCursor } from "../treeCursorUtils.js";
 import {
@@ -128,8 +128,8 @@ export class ObjectForest implements IEditableForest, WithBreakable {
 		return new ObjectForest(this.breaker, schema, anchors, this.additionalAsserts, this.roots);
 	}
 
-	public chunkField(cursor: ITreeCursorSynchronous): TreeChunk {
-		return chunkFieldSingle(cursor, { idCompressor: undefined, policy: defaultChunkPolicy });
+	public chunkField(cursor: ITreeCursorSynchronous): TreeChunk[] {
+		return chunkField(cursor, { idCompressor: undefined, policy: defaultChunkPolicy });
 	}
 
 	public forgetAnchor(anchor: Anchor): void {