Skip to content

fix(tree): track chunkIndex on field entry so getField finds parent #27084

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
brrichards wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix(tree): track chunkIndex on field entry so getField finds parent #27084

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0bbdda9
basicChunkCursor bug fix
brrichards Apr 17, 2026
8b79cd2
address feedback
brrichards Apr 20, 2026
6d112ef
Update docs
brrichards Apr 21, 2026
7c238af
Small fix
brrichards Apr 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 70 additions & 18 deletions packages/dds/tree/src/feature-libraries/chunked-forest/basicChunk.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 */,
Expand All @@ -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 {
Expand All @@ -203,11 +209,32 @@ 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");
return this.indexOfChunkStack[this.getNodeOnlyHeightFromHeight(height)] ?? oob();
}

private getStackedChunk(height: number): BasicChunk {
const index = this.getStackedChunkIndex(height);
return (this.siblingStack[height] as readonly TreeChunk[])[index] as BasicChunk;
Copy link
Copy Markdown
Contributor

@CraigMacomber CraigMacomber Apr 21, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should always be true, but for clarity/safety, can we:

Suggested change
return (this.siblingStack[height] as readonly TreeChunk[])[index] as BasicChunk;
const chunk = this.siblingStack[height] as readonly TreeChunk[])[index];
debugAssert(() => chunk instanceof BasicChunk && "Only basic chunks are expected");
return chunk as BasicChunk;

Copy link
Copy Markdown
Contributor

@CraigMacomber CraigMacomber Apr 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that as this are is rather performance sensitive, and very unlikely to fail due to some specific runtime data, I opted for debugAssert which can be compiled out of production builds. We will likely want to do the same for some other places in here.

}

/**
* 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 {
if (this.nestedCursor !== undefined) {
return this.nestedCursor.getFieldLength();
Expand Down Expand Up @@ -322,6 +349,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,
Expand All @@ -330,6 +362,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 {
Expand All @@ -355,8 +388,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);
Comment thread
brrichards marked this conversation as resolved.
this.index = 0;
this.siblings = [...fields.keys()]; // TODO: avoid this copy
this.assertChunkStacksMatchNodeDepth();
return true;
}

Expand Down Expand Up @@ -422,12 +458,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;
}
Expand Down Expand Up @@ -486,6 +521,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");
Comment thread
CraigMacomber marked this conversation as resolved.
this.assertChunkStacksMatchNodeDepth();
}

public exitNode(): void {
Expand All @@ -502,18 +542,30 @@ 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;
return (this.siblings as TreeChunk[])[this.indexOfChunk] 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.
*
* @remarks The parent node is the {@link BasicChunk} in the node array at the top of
* {@link siblingStack} while we are in {@link CursorLocationType.Fields} mode. We need the parent
* since a field's chunks are stored on the parent node's {@link BasicChunk.fields} map, not on
* the cursor itself.
*
* @returns The chunks that make up the field the cursor is currently in.
*/
private getField(): readonly TreeChunk[] {
Comment thread
brrichards marked this conversation as resolved.
if (this.siblingStack.length === 0) {
return this.root;
Expand All @@ -522,7 +574,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);
Comment on lines 562 to +577
Copy link
Copy Markdown
Contributor

@CraigMacomber CraigMacomber Apr 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Details about the parent should be a code comment inside the method, since its not relevant to callers.

Also, now that I see this explanation (thanks for that!), I think this code could be simplified.

this.siblingStack should, at the correct level, contain the TreeChunk array we want, so we can skip having to get the parent node in the first place.

const key: FieldKey = this.getFieldKey();
const field = parent.fields.get(key) ?? [];
return field;
Expand Down
101 changes: 98 additions & 3 deletions packages/dds/tree/src/test/feature-libraries/chunked-forest/basicChunk.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Comment thread
brrichards marked this conversation as resolved.
import {
basicChunkTree,
basicOnlyChunkPolicy,
Expand All @@ -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
Comment thread
brrichards marked this conversation as resolved.
} from "../../../feature-libraries/chunked-forest/uniformChunk.js";
import {
type TreeChunk,
chunkTree,
Expand Down Expand Up @@ -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));
Expand Down
Loading