fix(data-structures): ensure size consistency in RedBlackTree.from (#…

…6307)
denoland · Dec 29, 2024 · 08fe910 · 08fe910
1 parent c40f593
commit 08fe910
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 0 deletions.
diff --git a/data_structures/_binary_search_tree_internals.ts b/data_structures/_binary_search_tree_internals.ts
@@ -39,4 +39,5 @@ export const internals: {
     tree: BinarySearchTree<T>,
     node: BinarySearchNode<T>,
   ): BinarySearchNode<T> | null;
+  setSize<T>(tree: BinarySearchTree<T>, size: number): void;
 } = {} as typeof internals;
diff --git a/data_structures/binary_search_tree.ts b/data_structures/binary_search_tree.ts
@@ -140,6 +140,8 @@ export class BinarySearchTree<T> implements Iterable<T> {
       tree: BinarySearchTree<T>,
       node: BinarySearchNode<T>,
     ): BinarySearchNode<T> | null => tree.#removeNode(node);
+    internals.setSize = <T>(tree: BinarySearchTree<T>, size: number) =>
+      tree.#size = size;
   }
 
   /**

diff --git a/data_structures/red_black_tree.ts b/data_structures/red_black_tree.ts
@@ -14,6 +14,7 @@ const {
   rotateNode,
   insertNode,
   removeNode,
+  setSize,
 } = internals;
 
 /**
@@ -251,6 +252,7 @@ export class RedBlackTree<T> extends BinarySearchTree<T> {
             nodes.push(right);
           }
         }
+        setSize(result, collection.size);
       }
     } else {
       result = (options?.compare

diff --git a/data_structures/red_black_tree_test.ts b/data_structures/red_black_tree_test.ts
@@ -375,6 +375,7 @@ Deno.test("RedBlackTree.from() handles default ascend comparator", () => {
   const originalTree: RedBlackTree<number> = new RedBlackTree();
   for (const value of values) originalTree.insert(value);
   let tree: RedBlackTree<number> = RedBlackTree.from(originalTree);
+  assertEquals(originalTree.size, tree.size);
   assertEquals([...originalTree], expected);
   assertEquals([...tree], expected);
   assertEquals([...tree.nlrValues()], [...originalTree.nlrValues()]);
@@ -434,6 +435,7 @@ Deno.test("RedBlackTree.from() handles descend comparator", () => {
   const originalTree: RedBlackTree<number> = new RedBlackTree(descend);
   for (const value of values) originalTree.insert(value);
   let tree: RedBlackTree<number> = RedBlackTree.from(originalTree);
+  assertEquals(originalTree.size, tree.size);
   assertEquals([...originalTree], expected);
   assertEquals([...tree], expected);
   assertEquals([...tree.nlrValues()], [...originalTree.nlrValues()]);