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..abb8cec27fb7 100644 --- a/packages/dds/tree/src/feature-libraries/chunked-forest/basicChunk.ts +++ b/packages/dds/tree/src/feature-libraries/chunked-forest/basicChunk.ts @@ -169,7 +169,7 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor // 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; + const halfHeight = this.siblingStack.length >> 1; assert( this.indexOfChunkStack.length === halfHeight, 0x51c /* unexpected indexOfChunkStack */, @@ -203,8 +203,15 @@ export class BasicChunkCursor extends SynchronousCursor implements ChunkedCursor return this.indexStack[height] ?? oob(); } - private getStackedNode(height: number): BasicChunk { - const index = this.getStackedNodeIndex(height); + private getStackedChunkIndex(height: number): number { + assert(height % 2 === 1, "must be node height"); + assert(height >= 0, "must not be above root"); + // eslint-disable-next-line no-bitwise + return this.indexOfChunkStack[height >> 1] ?? oob(); + } + + private getStackedChunk(height: number): BasicChunk { + const index = this.getStackedChunkIndex(height); return (this.siblingStack[height] as readonly TreeChunk[])[index] as BasicChunk; } @@ -322,6 +329,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, @@ -355,6 +367,8 @@ 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 return true; @@ -422,8 +436,6 @@ 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; @@ -486,6 +498,10 @@ 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"); } public exitNode(): void { @@ -502,16 +518,15 @@ 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; } 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; + return (this.siblings as TreeChunk[])[this.indexOfChunk] as BasicChunk; } private getField(): readonly TreeChunk[] { @@ -522,7 +537,7 @@ 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); + 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..32d8bd6e1781 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,75 @@ 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, + ); + + // Logical node index 2 maps to chunk array index 1 (the trailing BasicChunk). + cursor.enterNode(2); + + // enterField drives getField, which on main looks up the parent at the + // logical node index (2) instead of the chunk array index (1). siblings[2] + // is out of bounds in the 2-chunk array, so the cast returns undefined and + // the subsequent field read throws. + cursor.enterField(subField); + assert.doesNotThrow(() => cursor.getFieldLength()); + }); + + 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, + ); + + // Logical node index 4 maps to chunk array index 2 (the trailing BasicChunk). + cursor.enterNode(4); + cursor.enterField(subField); + assert.doesNotThrow(() => cursor.getFieldLength()); + }); + describe("SequenceChunk", () => { it("root", () => { validateChunkCursor(new SequenceChunk([numericBasicChunk(0)]), numberSequenceField(1));