Add tree context API

.changeset/green-sides-watch.md

@@ -0,0 +1,49 @@
+---
+"@fluidframework/tree": minor
+"__section": tree
+---
+Added `TreeAlpha.context(node)` to provide context-aware APIs for any SharedTree node, plus a new `TreeContextAlpha` surface for transactions and branch checks.
+
+This release introduces a node-scoped context that works for nodes inserted into the tree as well as new nodes that are not yet inserted.
+The new `TreeContextAlpha` interface exposes `runTransaction` / `runTransactionAsync` and an `isBranch()` type guard.
+`TreeBranchAlpha` now extends `TreeContextAlpha`, so you can keep using branch APIs when available.
+
+### Migration
+If you previously used `TreeAlpha.branch(node)` to discover a branch, switch to `TreeAlpha.context(node)` and check `isBranch()`:
+
+```ts
+import { TreeAlpha } from "@fluidframework/tree/alpha";
+
+const context = TreeAlpha.context(node);
+if (context.isBranch()) {
+	// Same branch APIs as before
+	context.fork();
+}
+```
+
+`TreeAlpha.branch(node)` is now deprecated.
+Prefer the context API above.
+
+### New transaction entry point
+You can now run transactions from a node context, regardless of whether the node is hydrated:
+
+```ts
+const context = TreeAlpha.context(node);
+
+// No return value
+const result = context.runTransaction(() => {
+	node.count += 1;
+});
+
+// Return a value by wrapping it in `{ value }`
+const resultWithValue = context.runTransaction(() => ({ value: node.count }));
+```
+
+For asynchronous work:
+
+```ts
+const result = await context.runTransactionAsync(async () => {
+	await doWork();
+	return { value: node.count };
+});
+```

packages/dds/tree/api-report/tree.alpha.api.md

@@ -1453,13 +1453,11 @@ export namespace TableSchema {
 export function trackDirtyNodes(view: TreeViewAlpha<ImplicitFieldSchema>, dirty: DirtyTreeMap): () => void;
 
 // @alpha
-export type TransactionCallbackStatus<TSuccessValue, TFailureValue> = ({
+export type TransactionCallbackStatus<TSuccessValue, TFailureValue> = ((WithValue<TSuccessValue> & {
     rollback?: false;
-    value: TSuccessValue;
-} | {
+}) | (WithValue<TFailureValue> & {
     rollback: true;
-    value: TFailureValue;
-}) & {
+})) & {
     preconditionsOnRevert?: readonly TransactionConstraintAlpha[];
 };
 
@@ -1476,15 +1474,13 @@ export type TransactionResult = Omit<TransactionResultSuccess<unknown>, "value">
 export type TransactionResultExt<TSuccessValue, TFailureValue> = TransactionResultSuccess<TSuccessValue> | TransactionResultFailed<TFailureValue>;
 
 // @alpha
-export interface TransactionResultFailed<TFailureValue> {
+export interface TransactionResultFailed<TFailureValue> extends WithValue<TFailureValue> {
     success: false;
-    value: TFailureValue;
 }
 
 // @alpha
-export interface TransactionResultSuccess<TSuccessValue> {
+export interface TransactionResultSuccess<TSuccessValue> extends WithValue<TSuccessValue> {
     success: true;
-    value: TSuccessValue;
 }
 
 // @public @sealed
@@ -1498,9 +1494,11 @@ export const Tree: Tree;
 
 // @alpha @sealed
 export interface TreeAlpha {
+    // @deprecated
     branch(node: TreeNode): TreeBranchAlpha | undefined;
     child(node: TreeNode, key: string | number): TreeNode | TreeLeafValue | undefined;
     children(node: TreeNode): Iterable<[propertyKey: string | number, child: TreeNode | TreeLeafValue]>;
+    context(node: TreeNode): TreeContextAlpha;
     create<const TSchema extends ImplicitFieldSchema | UnsafeUnknownSchema>(schema: UnsafeUnknownSchema extends TSchema ? ImplicitFieldSchema : TSchema & ImplicitFieldSchema, data: InsertableField<TSchema>): Unhydrated<TSchema extends ImplicitFieldSchema ? TreeFieldFromImplicitField<TSchema> : TreeNode | TreeLeafValue | undefined>;
     exportCompressed(tree: TreeNode | TreeLeafValue, options: {
         idCompressor?: IIdCompressor;
@@ -1573,7 +1571,7 @@ export interface TreeBranch extends IDisposable {
 }
 
 // @alpha @sealed
-export interface TreeBranchAlpha extends TreeBranch {
+export interface TreeBranchAlpha extends TreeBranch, TreeContextAlpha {
     applyChange(change: JsonCompatibleReadOnly): void;
     readonly events: Listenable_2<TreeBranchEvents>;
     // (undocumented)
@@ -1614,6 +1612,15 @@ export enum TreeCompressionStrategy {
     Uncompressed = 1
 }
 
+// @alpha
+export interface TreeContextAlpha {
+    isBranch(): this is TreeBranchAlpha;
+    runTransaction<TValue>(transaction: () => WithValue<TValue>): TransactionResultExt<TValue, TValue>;
+    runTransaction(transaction: () => void): TransactionResult;
+    runTransactionAsync<TValue>(transaction: () => Promise<WithValue<TValue>>): Promise<TransactionResultExt<TValue, TValue>>;
+    runTransactionAsync(transaction: () => Promise<void>): Promise<TransactionResult>;
+}
+
 // @beta @input
 export interface TreeEncodingOptions<TKeyOptions = KeyEncodingOptions> {
     readonly keys?: TKeyOptions;
@@ -1911,6 +1918,11 @@ export interface WithType<out TName extends string = string, out TKind extends N
     get [typeSchemaSymbol](): TreeNodeSchemaClass<TName, TKind, TreeNode, never, boolean, TInfo>;
 }
 
+// @alpha
+export interface WithValue<TValue> {
+    value: TValue;
+}
+
 // (No @packageDocumentation comment for this package)
 
 ```

packages/dds/tree/src/index.ts

@@ -305,6 +305,8 @@ export {
 	type SchemaCompatibilitySnapshotsOptions,
 	type ArrayPlaceAnchor,
 	createArrayInsertionAnchor,
+	type WithValue,
+	type TreeContextAlpha,
 } from "./simple-tree/index.js";
 export {
 	SharedTree,

packages/dds/tree/src/shared-tree/schematizingTreeView.ts

@@ -166,6 +166,10 @@ export class SchematizingSimpleTreeView<
 		);
 	}
 
+	public isBranch(): this is TreeBranchAlpha {
+		return true;
+	}
+
 	public applyChange(change: JsonCompatibleReadOnly): void {
 		this.checkout.applySerializedChange(change);
 	}

packages/dds/tree/src/shared-tree/treeAlpha.ts

@@ -80,9 +80,13 @@ import {
 	exportConcise,
 	borrowCursorFromTreeNodeOrValue,
 	contentSchemaSymbol,
+	type TreeContextAlpha,
 	type TreeNodeSchema,
 	getUnhydratedContext,
 	type TreeBranchAlpha,
+	type TransactionResult,
+	type TransactionResultExt,
+	type WithValue,
 } from "../simple-tree/index.js";
 import { brand, extractFromOpaque, type JsonCompatible } from "../util/index.js";
 
@@ -237,9 +241,17 @@ export interface TreeAlpha {
 	 *
 	 * This does not fork a new branch, but rather retrieves the _existing_ branch for the node.
 	 * To create a new branch, use e.g. {@link TreeBranch.fork | `myBranch.fork()`}.
+	 *
+	 * @deprecated To obtain a {@link TreeBranchAlpha | branch }, use `TreeAlpha.context(node)` to obtain a {@link TreeContextAlpha | context} and then check {@link TreeContextAlpha.isBranch | isBranch()}.
 	 */
 	branch(node: TreeNode): TreeBranchAlpha | undefined;
 
+	/**
+	 * Retrieve the {@link TreeContextAlpha | context} for the given node.
+	 * @param node - The node to query
+	 */
+	context(node: TreeNode): TreeContextAlpha;
+
 	/**
 	 * Construct tree content that is compatible with the field defined by the provided `schema`.
 	 * @param schema - The schema for what to construct. As this is an {@link ImplicitFieldSchema}, a {@link FieldSchema}, {@link TreeNodeSchema} or {@link AllowedTypes} array can be provided.
@@ -763,6 +775,10 @@ export const TreeAlpha: TreeAlpha = {
 		return result;
 	},
 
+	context(node: TreeNode): TreeContextAlpha {
+		return this.branch(node) ?? UnhydratedTreeContext.instance;
+	},
+
 	branch(node: TreeNode): TreeBranchAlpha | undefined {
 		const kernel = getKernel(node);
 		if (!kernel.isHydrated()) {
@@ -1079,3 +1095,41 @@ function borrowFieldCursorFromTreeNodeOrValue(
 	const mapTree = mapTreeFromCursor(cursor);
 	return cursorForMapTreeField([mapTree]);
 }
+
+class UnhydratedTreeContext implements TreeContextAlpha {
+	public static instance = new UnhydratedTreeContext();
+	private constructor() {}
+
+	public isBranch(): this is TreeBranchAlpha {
+		return false;
+	}
+
+	public runTransaction<TValue>(
+		t: () => WithValue<TValue>,
+	): TransactionResultExt<TValue, TValue>;
+	public runTransaction(t: () => void): TransactionResult;
+	public runTransaction(
+		t: () => WithValue<unknown> | void,
+	): TransactionResultExt<unknown, unknown> | TransactionResult {
+		return UnhydratedTreeContext.wrapTransactionResult(t());
+	}
+
+	public runTransactionAsync<TValue>(
+		t: () => Promise<WithValue<TValue>>,
+	): Promise<TransactionResultExt<TValue, TValue>>;
+	public runTransactionAsync(t: () => Promise<void>): Promise<TransactionResult>;
+	public async runTransactionAsync(
+		t: () => Promise<WithValue<unknown> | void>,
+	): Promise<TransactionResultExt<unknown, unknown> | TransactionResult> {
+		return UnhydratedTreeContext.wrapTransactionResult(await t());
+	}
+
+	private static wrapTransactionResult<TValue>(
+		value: WithValue<TValue> | void,
+	): TransactionResultExt<TValue, TValue> | TransactionResult {
+		if (value?.value !== undefined) {
+			return { success: true, value: value.value };
+		}
+		return { success: true };
+	}
+}

packages/dds/tree/src/simple-tree/api/index.ts

@@ -22,6 +22,7 @@ export type {
 	TreeBranch,
 	TreeBranchAlpha,
 	TreeBranchEvents,
+	TreeContextAlpha,
 	ITreeAlpha,
 } from "./tree.js";
 export { asTreeViewAlpha } from "./tree.js";
@@ -164,6 +165,7 @@ export {
 	type TransactionResultExt,
 	type TransactionResultSuccess,
 	type TransactionResultFailed,
+	type WithValue,
 	rollback,
 } from "./transactionTypes.js";
 

packages/dds/tree/src/simple-tree/api/transactionTypes.ts

@@ -48,23 +48,28 @@ export interface NoChangeConstraint {
 	readonly type: "noChange";
 }
 
+/**
+ * An interface representing a value associated with a transaction.
+ * @alpha
+ */
+export interface WithValue<TValue> {
+	/** The user defined value. */
+	value: TValue;
+}
+
 /**
  * The status of the transaction callback in the {@link RunTransaction | RunTransaction} API.
  * @alpha
  */
 export type TransactionCallbackStatus<TSuccessValue, TFailureValue> = (
-	| {
+	| (WithValue<TSuccessValue> & {
 			/** Indicates that the transaction callback ran successfully. */
 			rollback?: false;
-			/** The user defined value when the transaction ran successfully. */
-			value: TSuccessValue;
-	  }
-	| {
+	  })
+	| (WithValue<TFailureValue> & {
 			/** Indicates that the transaction callback failed and the transaction should be rolled back. */
 			rollback: true;
-			/** The user defined value when the transaction failed. */
-			value: TFailureValue;
-	  }
+	  })
 ) & {
 	/**
 	 * An optional list of {@link TransactionConstraintAlpha | constraints} that will be checked when the commit corresponding
@@ -78,7 +83,7 @@ export type TransactionCallbackStatus<TSuccessValue, TFailureValue> = (
 
 /**
  * The status of a the transaction callback in the {@link RunTransaction | RunTransaction} API where the transaction doesn't
- * need to return a value. This is the same as {@link TransactionCallbackStatus} but with the `value` field omitted. This
+ * need to return a value. This is the same as {@link TransactionCallbackStatus} but with the `value` field omitted.
  * @alpha
  */
 export type VoidTransactionCallbackStatus = Omit<
@@ -90,22 +95,18 @@ export type VoidTransactionCallbackStatus = Omit<
  * The result of the {@link RunTransaction | RunTransaction} API when it was successful.
  * @alpha
  */
-export interface TransactionResultSuccess<TSuccessValue> {
+export interface TransactionResultSuccess<TSuccessValue> extends WithValue<TSuccessValue> {
 	/** Indicates that the transaction was successful. */
 	success: true;
-	/** The user defined value when the transaction was successful. */
-	value: TSuccessValue;
 }
 
 /**
  * The result of the {@link RunTransaction | RunTransaction} API when it failed.
  * @alpha
  */
-export interface TransactionResultFailed<TFailureValue> {
+export interface TransactionResultFailed<TFailureValue> extends WithValue<TFailureValue> {
 	/** Indicates that the transaction failed. */
 	success: false;
-	/** The user defined value when the transaction failed. */
-	value: TFailureValue;
 }
 
 /**