refactor(tree): minor map tree to unhydrated flex tree node rename an…

…d cleanup (#23651) - renames a weak map to be more accurate - removes tryGetTreeNodeFromMapNode and replaces it with a restricted version of the original map
microsoft · Jan 24, 2025 · 472f93a · 472f93a
1 parent 768cc1e
commit 472f93a
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 19 deletions.
diff --git a/packages/dds/tree/src/simple-tree/core/getOrCreateNode.ts b/packages/dds/tree/src/simple-tree/core/getOrCreateNode.ts
@@ -6,7 +6,11 @@
 import type { TreeValue } from "../../core/index.js";
 import type { FlexTreeNode } from "../../feature-libraries/index.js";
 import { fail } from "../../util/index.js";
-import { type InnerNode, mapTreeNodeToProxy, proxySlot } from "./treeNodeKernel.js";
+import {
+	type InnerNode,
+	unhydratedFlexTreeNodeToTreeNode,
+	proxySlot,
+} from "./treeNodeKernel.js";
 import { getSimpleNodeSchemaFromInnerNode } from "./schemaCaching.js";
 import type { TreeNode, InternalTreeNode } from "./types.js";
 import { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
@@ -20,7 +24,7 @@ import { UnhydratedFlexTreeNode } from "./unhydratedFlexTree.js";
 export function getOrCreateNodeFromInnerNode(flexNode: InnerNode): TreeNode | TreeValue {
 	const cached =
 		flexNode instanceof UnhydratedFlexTreeNode
-			? mapTreeNodeToProxy.get(flexNode)
+			? unhydratedFlexTreeNodeToTreeNode.get(flexNode)
 			: flexNode.anchorNode.slots.get(proxySlot);
 
 	if (cached !== undefined) {

diff --git a/packages/dds/tree/src/simple-tree/core/index.ts b/packages/dds/tree/src/simple-tree/core/index.ts
@@ -10,7 +10,7 @@ export {
 	tryGetTreeNodeSchema,
 	type InnerNode,
 	tryDisposeTreeNode,
-	tryGetTreeNodeFromMapNode,
+	unhydratedFlexTreeNodeToTreeNode,
 	getOrCreateInnerNode,
 	treeNodeFromAnchor,
 } from "./treeNodeKernel.js";

diff --git a/packages/dds/tree/src/simple-tree/core/treeNodeKernel.ts b/packages/dds/tree/src/simple-tree/core/treeNodeKernel.ts
@@ -146,7 +146,7 @@ export class TreeNodeKernel {
 
 		if (innerNode instanceof UnhydratedFlexTreeNode) {
 			// Unhydrated case
-			mapTreeNodeToProxy.set(innerNode, node);
+			unhydratedFlexTreeNodeToTreeNodeInternal.set(innerNode, node);
 			// Register for change events from the unhydrated flex node.
 			// These will be fired if the unhydrated node is edited, and will also be forwarded later to the hydrated node.
 			this.#hydrationState = {
@@ -158,7 +158,7 @@ export class TreeNodeKernel {
 
 					let n: UnhydratedFlexTreeNode | undefined = innerNode;
 					while (n !== undefined) {
-						const treeNode = mapTreeNodeToProxy.get(n);
+						const treeNode = unhydratedFlexTreeNodeToTreeNodeInternal.get(n);
 						if (treeNode !== undefined) {
 							const kernel = getKernel(treeNode);
 							kernel.#unhydratedEvents.value.emit("subtreeChangedAfterBatch");
@@ -200,7 +200,7 @@ export class TreeNodeKernel {
 	private hydrate(anchorNode: AnchorNode): void {
 		assert(!this.disposed, 0xa2a /* cannot hydrate a disposed node */);
 		assert(!isHydrated(this.#hydrationState), 0xa2b /* hydration should only happen once */);
-		mapTreeNodeToProxy.delete(this.#hydrationState.innerNode);
+		unhydratedFlexTreeNodeToTreeNodeInternal.delete(this.#hydrationState.innerNode);
 		this.#hydrationState = this.createHydratedState(anchorNode);
 
 		// If needed, register forwarding emitters for events from before hydration
@@ -389,9 +389,20 @@ type KernelEvents = Pick<AnchorEvents, (typeof kernelEvents)[number]>;
 export type InnerNode = FlexTreeNode | UnhydratedFlexTreeNode;
 
 /**
- * {@inheritdoc proxyToMapTreeNode}
+ * Associates a given {@link UnhydratedFlexTreeNode} with a {@link TreeNode}.
  */
-export const mapTreeNodeToProxy = new WeakMap<UnhydratedFlexTreeNode, TreeNode>();
+const unhydratedFlexTreeNodeToTreeNodeInternal = new WeakMap<
+	UnhydratedFlexTreeNode,
+	TreeNode
+>();
+/**
+ * Retrieves the {@link TreeNode} associated with the given {@link UnhydratedFlexTreeNode} if any.
+ */
+export const unhydratedFlexTreeNodeToTreeNode =
+	unhydratedFlexTreeNodeToTreeNodeInternal as Pick<
+		WeakMap<UnhydratedFlexTreeNode, TreeNode>,
+		"get"
+	>;
 
 /**
  * An anchor slot which associates an anchor with its corresponding TreeNode, if there is one.
@@ -401,15 +412,6 @@ export const mapTreeNodeToProxy = new WeakMap<UnhydratedFlexTreeNode, TreeNode>(
  */
 export const proxySlot = anchorSlot<TreeNode>();
 
-/**
- * Retrieves the node associated with the given MapTreeNode node if any.
- */
-export function tryGetTreeNodeFromMapNode(
-	flexNode: UnhydratedFlexTreeNode,
-): TreeNode | undefined {
-	return mapTreeNodeToProxy.get(flexNode);
-}
-
 export function tryDisposeTreeNode(anchorNode: AnchorNode): void {
 	const treeNode = anchorNode.slots.get(proxySlot);
 	if (treeNode !== undefined) {

diff --git a/packages/dds/tree/src/simple-tree/proxies.ts b/packages/dds/tree/src/simple-tree/proxies.ts
@@ -24,9 +24,9 @@ import { type Mutable, fail, isReadonlyArray } from "../util/index.js";
 import {
 	getKernel,
 	type TreeNode,
-	tryGetTreeNodeFromMapNode,
 	getOrCreateNodeFromInnerNode,
 	tryUnhydratedFlexTreeNode,
+	unhydratedFlexTreeNodeToTreeNode,
 } from "./core/index.js";
 
 /**
@@ -146,7 +146,7 @@ function walkMapTree(
 		const [p, m] = next;
 		const mapTreeNode = tryUnhydratedFlexTreeNode(m);
 		if (mapTreeNode !== undefined) {
-			const treeNode = tryGetTreeNodeFromMapNode(mapTreeNode);
+			const treeNode = unhydratedFlexTreeNodeToTreeNode.get(mapTreeNode);
 			if (treeNode !== undefined) {
 				onVisitTreeNode(p, treeNode);
 			}