Optimize anchorset and flex tree traversal (#22717)

## Description

Optimizations based on a profile of BubbleBench.

This resulted in about 5% more bubbles (80 vs 76).

Given this application's read costs scale with bubbles squared (for the
collision logic), contains writes, and most reads are leaf values which
are not impacted by this, this seems like a pretty good win. Some real
apps could probably show significantly larger gains if they are less
leaf heavy or less write heavy.

.changeset/spicy-candles-smash.md

@@ -0,0 +1,13 @@
+---
+"@fluidframework/tree": minor
+"fluid-framework": minor
+---
+---
+"section": tree
+---
+
+Optimize non-leaf field access
+
+When reading non-leaf children which have been read previously, they are retrieved from cache faster.
+Several operations on subtrees under arrays have been optimized, including reading of non-leaf nodes for the first time.
+Overall this showed a roughly 5% speed up in a read heavy test application (the BubbleBench example) but gains are expected to vary a lot based on use-case.

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

@@ -233,6 +233,13 @@ export interface AnchorNode extends UpPath<AnchorNode>, Listenable<AnchorEvents>
 	 */
 	child(key: FieldKey, index: number): UpPath<AnchorNode>;
 
+	/**
+	 * Gets the child AnchorNode if already exists.
+	 *
+	 * Does NOT add a ref, so the returned AnchorNode must be used with care.
+	 */
+	childIfAnchored(key: FieldKey, index: number): AnchorNode | undefined;
+
 	/**
 	 * Gets a child AnchorNode (creating it if needed), and an Anchor owning a ref to it.
 	 * Caller is responsible for freeing the returned Anchor, and must not use the AnchorNode after that.
@@ -413,16 +420,18 @@ export class AnchorSet implements Listenable<AnchorSetRootEvents>, AnchorLocator
 
 	/**
 	 * Finds a path node if it already exists.
+	 *
+	 * Does not add a ref!
 	 */
-	private find(path: UpPath): PathNode | undefined {
+	public find(path: UpPath): PathNode | undefined {
 		if (path instanceof PathNode) {
 			if (path.anchorSet === this) {
 				return path;
 			}
 		}
 		const parent = path.parent ?? this.root;
 		const parentPath = this.find(parent);
-		return parentPath?.tryGetChild(path.parentField, path.parentIndex);
+		return parentPath?.childIfAnchored(path.parentField, path.parentIndex);
 	}
 
 	/**
@@ -452,7 +461,7 @@ export class AnchorSet implements Listenable<AnchorSetRootEvents>, AnchorLocator
 		while ((wrapWith = stack.pop()) !== undefined) {
 			if (path === undefined || path instanceof PathNode) {
 				// If path already has an anchor, get an anchor for it's child if there is one:
-				const child = (path ?? this.root).tryGetChild(
+				const child = (path ?? this.root).childIfAnchored(
 					wrapWith.parentField,
 					wrapWith.parentIndex,
 				);
@@ -1210,7 +1219,11 @@ class PathNode extends ReferenceCountedBase implements UpPath<PathNode>, AnchorN
 	public child(key: FieldKey, index: number): UpPath<AnchorNode> {
 		// Fast path: if child exists, return it.
 		return (
-			this.tryGetChild(key, index) ?? { parent: this, parentField: key, parentIndex: index }
+			this.childIfAnchored(key, index) ?? {
+				parent: this,
+				parentField: key,
+				parentIndex: index,
+			}
 		);
 	}
 
@@ -1271,8 +1284,7 @@ class PathNode extends ReferenceCountedBase implements UpPath<PathNode>, AnchorN
 			field = [];
 			this.children.set(key, field);
 		}
-		// TODO: should do more optimized search (ex: binary search).
-		let child = field.find((c) => c.parentIndex === index);
+		let child = binaryFind(field, index);
 		if (child === undefined) {
 			child = new PathNode(this.anchorSet, key, index, this);
 			field.push(child);
@@ -1284,16 +1296,11 @@ class PathNode extends ReferenceCountedBase implements UpPath<PathNode>, AnchorN
 		return child;
 	}
 
-	/**
-	 * Gets a child if it exists.
-	 * Does NOT add a ref.
-	 */
-	public tryGetChild(key: FieldKey, index: number): PathNode | undefined {
+	public childIfAnchored(key: FieldKey, index: number): PathNode | undefined {
 		assert(this.status === Status.Alive, 0x40d /* PathNode must be alive */);
 		const field = this.children.get(key);
 
-		// TODO: should do more optimized search (ex: binary search or better) using index.
-		return field?.find((c) => c.parentIndex === index);
+		return field === undefined ? undefined : binaryFind(field, index);
 	}
 
 	/**
@@ -1343,3 +1350,46 @@ class PathNode extends ReferenceCountedBase implements UpPath<PathNode>, AnchorN
 		}
 	}
 }
+
+/**
+ * Find a child PathNode by index using a binary search.
+ * @param sorted - array of PathNode's sorted by parentIndex.
+ * @param index - index being looked for.
+ * @returns child with the requested parentIndex, or undefined.
+ * @privateRemarks
+ * This function is very commonly used with small arrays (length 0 or one for all non sequence fields),
+ * and is currently a hot path due to how flex tree leaves to excessive cursor to anchor and anchor to cursor translations,
+ * both of which walk paths down the AnchorSet.
+ * Additionally current usages tends to fully populate the anchor tree leading the correct array index to be the requested parent index.
+ * This makes the performance of this performance both important in small cases and easy to overly tune to the current usage patterns.
+ * This lead to not implementing a general purpose reusable binary search.
+ * Once this function is not so heavily overused due to inefficient patterns in flex-tree,
+ * replacing it with a standard binary search is likely fine.
+ * Until then, care and benchmarking should be used when messing with this function.
+ */
+function binaryFind(sorted: readonly PathNode[], index: number): PathNode | undefined {
+	// Try guessing the list is not sparse as a starter:
+	const guess = sorted[index];
+	if (guess !== undefined && guess.parentIndex === index) {
+		return guess;
+	}
+
+	// inclusive
+	let min = 0;
+	// exclusive
+	let max = sorted.length;
+
+	while (min !== max) {
+		const mid = Math.floor((min + max) / 2);
+		const item = sorted[mid]!;
+		const found = item.parentIndex;
+		if (found === index) {
+			return item; // Found the target, return it.
+		} else if (found > index) {
+			max = mid; // Continue search on lower half.
+		} else {
+			min = mid + 1; // Continue search on left half.
+		}
+	}
+	return undefined; // If we reach here, target is not in array (or array was not sorted)
+}

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

@@ -4,8 +4,10 @@
  */
 
 import { assert } from "@fluidframework/core-utils/internal";
+import { UsageError } from "@fluidframework/telemetry-utils/internal";
 
 import {
+	type AnchorNode,
 	CursorLocationType,
 	type ExclusiveMapTree,
 	type FieldAnchor,
@@ -51,9 +53,7 @@ import {
 	tryMoveCursorToAnchorSymbol,
 } from "./lazyEntity.js";
 import { type LazyTreeNode, makeTree } from "./lazyNode.js";
-import { unboxedFlexNode } from "./unboxed.js";
 import { indexForAt, treeStatusFromAnchorCache } from "./utilities.js";
-import { UsageError } from "@fluidframework/telemetry-utils/internal";
 import { cursorForMapTreeField, cursorForMapTreeNode } from "../mapTreeCursor.js";
 
 /**
@@ -196,7 +196,7 @@ export abstract class LazyField extends LazyEntity<FieldAnchor> implements FlexT
 
 	public atIndex(index: number): FlexTreeUnknownUnboxed {
 		return inCursorNode(this[cursorSymbol], index, (cursor) =>
-			unboxedFlexNode(this.context, cursor),
+			unboxedFlexNode(this.context, cursor, this[anchorSymbol]),
 		);
 	}
 
@@ -222,7 +222,7 @@ export abstract class LazyField extends LazyEntity<FieldAnchor> implements FlexT
 
 	public [Symbol.iterator](): IterableIterator<FlexTreeUnknownUnboxed> {
 		return iterateCursorField(this[cursorSymbol], (cursor) =>
-			unboxedFlexNode(this.context, cursor),
+			unboxedFlexNode(this.context, cursor, this[anchorSymbol]),
 		);
 	}
 
@@ -259,9 +259,7 @@ export class LazySequence extends LazyField implements FlexTreeSequenceField {
 			return undefined;
 		}
 
-		return inCursorNode(this[cursorSymbol], finalIndex, (cursor) =>
-			unboxedFlexNode(this.context, cursor),
-		);
+		return this.atIndex(finalIndex);
 	}
 	public get asArray(): readonly FlexTreeUnknownUnboxed[] {
 		return this.map((x) => x);
@@ -352,3 +350,42 @@ const builderList: [FieldKindIdentifier, Builder][] = [
 ];
 
 const kindToClass: ReadonlyMap<FieldKindIdentifier, Builder> = new Map(builderList);
+
+/**
+ * Returns the flex tree node, or the value if it has one.
+ */
+export function unboxedFlexNode(
+	context: Context,
+	cursor: ITreeSubscriptionCursor,
+	fieldAnchor: FieldAnchor,
+): FlexTreeUnknownUnboxed {
+	const value = cursor.value;
+	if (value !== undefined) {
+		return value;
+	}
+
+	// Try accessing cached child node via anchors.
+	// This avoids O(depth) related costs from makeTree in the cached case.
+	const anchor = fieldAnchor.parent;
+	let child: AnchorNode | undefined;
+	if (anchor !== undefined) {
+		const anchorNode = context.checkout.forest.anchors.locate(anchor);
+		assert(anchorNode !== undefined, "missing anchor");
+		child = anchorNode.childIfAnchored(fieldAnchor.fieldKey, cursor.fieldIndex);
+	} else {
+		child = context.checkout.forest.anchors.find({
+			parent: undefined,
+			parentField: fieldAnchor.fieldKey,
+			parentIndex: cursor.fieldIndex,
+		});
+	}
+
+	if (child !== undefined) {
+		const cached = child.slots.get(flexTreeSlot);
+		if (cached !== undefined) {
+			return cached;
+		}
+	}
+
+	return makeTree(context, cursor);
+}