From e4501cf1e9a12d3e8dc770394262e916681179db Mon Sep 17 00:00:00 2001
From: Jiang Sheng <jiangshengdev@outlook.com>
Date: Thu, 26 Dec 2024 20:11:57 +0800
Subject: [PATCH] fix(data-structures): ensure size consistency in
 RedBlackTree.from

---
 data_structures/_binary_search_tree_internals.ts | 1 +
 data_structures/binary_search_tree.ts            | 2 ++
 data_structures/red_black_tree.ts                | 2 ++
 data_structures/red_black_tree_test.ts           | 2 ++
 4 files changed, 7 insertions(+)

diff --git a/data_structures/_binary_search_tree_internals.ts b/data_structures/_binary_search_tree_internals.ts
index 18b3c4f80905..1ce9f300d791 100644
--- 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
index 4ac1c352538a..bb22801d58c2 100644
--- 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
index 8277e5b44c65..cd56445610be 100644
--- 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
index d7f83c4b58fc..f3c1863f140b 100644
--- 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()]);