diff --git a/packages/dds/tree/src/feature-libraries/chunked-forest/basicChunk.ts b/packages/dds/tree/src/feature-libraries/chunked-forest/basicChunk.ts index 2ed31df86b6f..4ebaa68996f5 100644 --- a/packages/dds/tree/src/feature-libraries/chunked-forest/basicChunk.ts +++ b/packages/dds/tree/src/feature-libraries/chunked-forest/basicChunk.ts @@ -3,7 +3,7 @@ * Licensed under the MIT License. */ -import { assert, oob, fail } from "@fluidframework/core-utils/internal"; +import { assert, oob, fail, debugAssert } from "@fluidframework/core-utils/internal"; import { CursorLocationType, @@ -166,10 +166,19 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor if (this.nestedCursor !== undefined) { return this.nestedCursor.mode; } - // Compute the number of nodes deep the current depth is. - // We want the floor of the result, which can computed using a bitwise shift assuming the depth is less than 2^31, which seems safe. - // eslint-disable-next-line no-bitwise - const halfHeight = (this.siblingStack.length + 1) >> 1; + this.assertChunkStacksMatchNodeDepth(); + return this.siblingStack.length % 2 === 0 + ? CursorLocationType.Fields + : CursorLocationType.Nodes; + } + + /** + * Asserts that the node-only stacks (`indexOfChunkStack` and `indexWithinChunkStack`) are in sync with `siblingStack`. + * Since `siblingStack` interleaves field and node levels while the node-only stacks are pushed/popped only on node-level transitions, + * their length should always equal the number of node levels traversed. + */ + private assertChunkStacksMatchNodeDepth(): void { + const halfHeight = this.getNodeOnlyHeightFromHeight(); assert( this.indexOfChunkStack.length === halfHeight, 0x51c /* unexpected indexOfChunkStack */, @@ -178,9 +187,6 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor this.indexWithinChunkStack.length === halfHeight, 0x51d /* unexpected indexWithinChunkStack */, ); - return this.siblingStack.length % 2 === 0 - ? CursorLocationType.Fields - : CursorLocationType.Nodes; } public getFieldKey(): FieldKey { @@ -203,9 +209,32 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor return this.indexStack[height] ?? oob(); } - private getStackedNode(height: number): BasicChunk { - const index = this.getStackedNodeIndex(height); - return (this.siblingStack[height] as readonly TreeChunk[])[index] as BasicChunk; + private getStackedChunkIndex(height: number): number { + assert(height % 2 === 1, "must be node height"); + assert(height >= 0, "must not be above root"); + return this.indexOfChunkStack[this.getNodeOnlyHeightFromHeight(height)] ?? oob(); + } + + private getStackedChunk(height: number): BasicChunk { + const index = this.getStackedChunkIndex(height); + const chunk = (this.siblingStack[height] as readonly TreeChunk[])[index]; + debugAssert(() => chunk instanceof BasicChunk || "only basic chunks are expected"); + return chunk as BasicChunk; + } + + /** + * Converts a {@link height}, which contains field and node levels, into the corresponding depth/index + * for the node-only stacks ({@link indexOfChunkStack} and {@link indexWithinChunkStack}), which are + * only pushed on node-level transitions. + * + * @param height - A depth in {@link siblingStack} to convert. Defaults to {@link siblingStack}'s + * current length, which gives the current depth of the node-only stacks. + * @returns `floor(height / 2)` — the number of node levels at or below the given stack height. + */ + private getNodeOnlyHeightFromHeight(height: number = this.siblingStack.length): number { + // The bitwise shift computes the floor, which is valid assuming the depth is less than 2^31, which seems safe. + // eslint-disable-next-line no-bitwise + return height >> 1; } public getFieldLength(): number { @@ -322,6 +351,11 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor assert(this.mode === CursorLocationType.Nodes, 0x528 /* must be in nodes mode */); this.siblingStack.push(this.siblings); this.indexStack.push(this.index); + // Save the chunk array position of the current node. When siblings contain + // multi node chunks, the flat node index diverges from the array position, + // so getField needs this to locate the parent in the sibling array. + this.indexOfChunkStack.push(this.indexOfChunk); + this.indexWithinChunkStack.push(this.indexWithinChunk); // For fields, siblings are only used for key lookup and // nextField and which has arbitrary iteration order, @@ -330,6 +364,7 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor // at the cost of an allocation here. this.index = 0; this.siblings = [key]; + this.assertChunkStacksMatchNodeDepth(); } public nextField(): boolean { @@ -355,8 +390,11 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor this.siblingStack.push(this.siblings); this.indexStack.push(this.index); + this.indexOfChunkStack.push(this.indexOfChunk); + this.indexWithinChunkStack.push(this.indexWithinChunk); this.index = 0; this.siblings = [...fields.keys()]; // TODO: avoid this copy + this.assertChunkStacksMatchNodeDepth(); return true; } @@ -422,12 +460,11 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor } this.siblingStack.push(this.siblings); this.indexStack.push(this.index); - this.indexOfChunkStack.push(this.indexOfChunk); - this.indexWithinChunkStack.push(this.indexWithinChunk); this.index = 0; this.siblings = siblings; this.indexOfChunk = 0; this.indexWithinChunk = 0; + this.assertChunkStacksMatchNodeDepth(); this.initNestedCursor(); return true; } @@ -486,6 +523,11 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor this.siblings = this.siblingStack.pop() ?? fail(0xaf0 /* Unexpected siblingStack.length */); this.index = this.indexStack.pop() ?? fail(0xaf1 /* Unexpected indexStack.length */); + this.indexOfChunk = + this.indexOfChunkStack.pop() ?? fail("Unexpected indexOfChunkStack.length"); + this.indexWithinChunk = + this.indexWithinChunkStack.pop() ?? fail("Unexpected indexWithinChunkStack.length"); + this.assertChunkStacksMatchNodeDepth(); } public exitNode(): void { @@ -502,18 +544,27 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor this.siblings = this.siblingStack.pop() ?? fail(0xaf2 /* Unexpected siblingStack.length */); this.index = this.indexStack.pop() ?? fail(0xaf3 /* Unexpected indexStack.length */); - this.indexOfChunk = - this.indexOfChunkStack.pop() ?? fail(0xaf4 /* Unexpected indexOfChunkStack.length */); - this.indexWithinChunk = - this.indexWithinChunkStack.pop() ?? - fail(0xaf5 /* Unexpected indexWithinChunkStack.length */); + // At the Fields level these aren't semantically used, but reset for consistent state + // (so a fully-iterated cursor matches a fresh cursor at the same logical position). + this.indexOfChunk = 0; + this.indexWithinChunk = 0; + this.assertChunkStacksMatchNodeDepth(); } private getNode(): BasicChunk { assert(this.mode === CursorLocationType.Nodes, 0x52f /* can only get node when in node */); - return (this.siblings as TreeChunk[])[this.index] as BasicChunk; + const chunk = (this.siblings as TreeChunk[])[this.indexOfChunk]; + debugAssert(() => chunk instanceof BasicChunk || "only basic chunks are expected"); + return chunk as BasicChunk; } + /** + * Resolves the chunks that make up the field the cursor is currently in. At the root, this is + * {@link root} directly. Otherwise, the cursor must be in {@link CursorLocationType.Fields} mode, + * and the result is looked up on the parent node using the current field key. + * + * @returns The chunks that make up the field the cursor is currently in. + */ private getField(): readonly TreeChunk[] { if (this.siblingStack.length === 0) { return this.root; @@ -522,7 +573,11 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor this.mode === CursorLocationType.Fields, 0x530 /* can only get field when in fields */, ); - const parent = this.getStackedNode(this.indexStack.length - 1); + // The parent node is the `BasicChunk` in the node array at the top of + // `siblingStack` while we are in `CursorLocationType.Fields` mode. We need the parent + // since a field's chunks are stored on the parent node's `BasicChunk.fields` map, not on + // the cursor itself. + const parent = this.getStackedChunk(this.siblingStack.length - 1); const key: FieldKey = this.getFieldKey(); const field = parent.fields.get(key) ?? []; return field; diff --git a/packages/dds/tree/src/test/feature-libraries/chunked-forest/basicChunk.spec.ts b/packages/dds/tree/src/test/feature-libraries/chunked-forest/basicChunk.spec.ts index 10e1ddca2a2c..d78f28cdab82 100644 --- a/packages/dds/tree/src/test/feature-libraries/chunked-forest/basicChunk.spec.ts +++ b/packages/dds/tree/src/test/feature-libraries/chunked-forest/basicChunk.spec.ts @@ -7,13 +7,17 @@ import { strict as assert } from "node:assert"; import { EmptyKey, + type FieldKey, type ITreeCursor, type ITreeCursorSynchronous, type JsonableTree, type ChunkedCursor, } from "../../../core/index.js"; -// eslint-disable-next-line import-x/no-internal-modules -import { BasicChunk } from "../../../feature-libraries/chunked-forest/basicChunk.js"; +import { + BasicChunk, + BasicChunkCursor, + // eslint-disable-next-line import-x/no-internal-modules +} from "../../../feature-libraries/chunked-forest/basicChunk.js"; import { basicChunkTree, basicOnlyChunkPolicy, @@ -22,9 +26,14 @@ import { // eslint-disable-next-line import-x/no-internal-modules } from "../../../feature-libraries/chunked-forest/chunkTree.js"; // eslint-disable-next-line import-x/no-internal-modules -import { uniformChunk } from "../../../feature-libraries/chunked-forest/index.js"; +import { dummyRoot, uniformChunk } from "../../../feature-libraries/chunked-forest/index.js"; // eslint-disable-next-line import-x/no-internal-modules import { SequenceChunk } from "../../../feature-libraries/chunked-forest/sequenceChunk.js"; +import { + TreeShape, + UniformChunk, + // eslint-disable-next-line import-x/no-internal-modules +} from "../../../feature-libraries/chunked-forest/uniformChunk.js"; import { type TreeChunk, chunkTree, @@ -155,6 +164,92 @@ describe("basic chunk", () => { assert(!cursor.atChunkRoot()); }); + it("getField resolves parent via chunk array index when preceded by a multi-node chunk", () => { + // A root field with 3 logical nodes split across 2 chunks: + // chunks[0]: UniformChunk with 2 number nodes. + // chunks[1]: BasicChunk holding a single subField leaf. + const numberShape = new TreeShape(brand(numberSchema.identifier), true, []); + const uniform = new UniformChunk(numberShape.withTopLevelLength(2), [10, 20]); + + const subField: FieldKey = brand("subField"); + const subLeaf = new BasicChunk(brand(numberSchema.identifier), new Map(), 42); + const trailingBasic = new BasicChunk(brand("Trailing"), new Map([[subField, [subLeaf]]])); + + const cursor = new BasicChunkCursor( + [uniform, trailingBasic], + // siblingStack, indexStack, indexOfChunkStack, indexWithinChunkStack + [], + [], + [], + [], + [dummyRoot], + // index, indexOfChunk, indexWithinChunk + 0, + 0, + 0, + undefined, + ); + + // Transitions from field to firstNode and then increments to nodeIndex 2, + // which maps to chunk array index 1 (the trailing BasicChunk). + cursor.enterNode(2); + + // Entering the subfield of the trailing BasicChunk calls enterField, which is followed by + // getField, which looks up the current node (whose field we should get). This must + // handle the case where the index of the node does not match the index of the chunk + // containing the node in the siblings array held at the top of `siblingStack`. + // In this test, the node index is 2 and the chunk index is 1. + cursor.enterField(subField); + // Validates the cursor is positioned on the subField holding the subLeaf, + // then enters it to confirm its value (42). + assert.equal(cursor.getFieldLength(), 1); + cursor.firstNode(); + assert.equal(cursor.value, 42); + }); + + it("getField resolves parent via chunk array index when preceded by two multi-node chunks", () => { + // Root field with 5 logical nodes across 3 chunks: + // chunks[0]: UniformChunk with 2 number nodes. + // chunks[1]: UniformChunk with 2 number nodes. + // chunks[2]: BasicChunk with a single subField leaf. + const numberShape = new TreeShape(brand(numberSchema.identifier), true, []); + const uniformA = new UniformChunk(numberShape.withTopLevelLength(2), [10, 20]); + const uniformB = new UniformChunk(numberShape.withTopLevelLength(2), [30, 40]); + + const subField: FieldKey = brand("subField"); + const subLeaf = new BasicChunk(brand(numberSchema.identifier), new Map(), 99); + const trailingBasic = new BasicChunk(brand("Trailing"), new Map([[subField, [subLeaf]]])); + + const cursor = new BasicChunkCursor( + [uniformA, uniformB, trailingBasic], + [], + [], + [], + [], + [dummyRoot], + 0, + 0, + 0, + undefined, + ); + + // Transitions from field to firstNode and then increments to nodeIndex 4, + // which maps to chunk array index 2 (the trailing BasicChunk). + cursor.enterNode(4); + + // Entering the subfield of the trailing BasicChunk calls enterField, which is followed by + // getField, which looks up the current node (whose field we should get). This must + // handle the case where the index of the node does not match the index of the chunk + // containing the node in the siblings array held at the top of `siblingStack`. + // In this test, the node index is 4 and the chunk index is 2. + cursor.enterField(subField); + // Validates the cursor is positioned on the subField holding the subLeaf, + // then enters it to confirm its value (99). + assert.equal(cursor.getFieldLength(), 1); + cursor.firstNode(); + assert.equal(cursor.value, 99); + }); + describe("SequenceChunk", () => { it("root", () => { validateChunkCursor(new SequenceChunk([numericBasicChunk(0)]), numberSequenceField(1));