fix: prevent TreeMultiSet comparator-equal object corruption

Author: zrwusa

CHANGELOG.md

@@ -8,7 +8,15 @@ All notable changes to this project will be documented in this file.
 - [Semantic Versioning](https://semver.org/spec/v2.0.0.html)
 - [`auto-changelog`](https://github.com/CookPete/auto-changelog)
 
-## [v2.6.0](https://github.com/zrwusa/data-structure-typed/compare/v2.5.3...main) (upcoming)
+## [v2.6.2](https://github.com/zrwusa/data-structure-typed/compare/v2.6.1...v2.6.2) (24 June 2026)
+
+### Changes
+
+- fix(TreeMultiSet): preserve counts for comparator-equal object keys [`#131`](https://github.com/zrwusa/data-structure-typed/issues/131)
+
+## [v2.6.1](https://github.com/zrwusa/data-structure-typed/compare/v2.6.0...v2.6.1) (2 April 2026)
+
+## [v2.6.0](https://github.com/zrwusa/data-structure-typed/compare/v2.5.3...v2.6.0) (31 March 2026)
 
 ## [v2.5.3](https://github.com/zrwusa/data-structure-typed/compare/v2.5.1...v2.5.3) (31 March 2026)
 

package.json

@@ -1,6 +1,6 @@
 {
   "name": "data-structure-typed",
-  "version": "2.6.1",
+  "version": "2.6.2",
   "description": "Production-ready TypeScript data structures: Heap, Deque, Trie, Graph, Red-Black Tree, TreeMap, TreeSet, and more. Zero dependencies, type-safe, with getRank/getByRank/rangeByRank support.",
   "browser": "dist/umd/data-structure-typed.min.js",
   "umd:main": "dist/umd/data-structure-typed.min.js",

src/data-structures/binary-tree/tree-multi-set.ts

@@ -113,7 +113,7 @@ export class TreeMultiSet<K = any, R = K> implements Iterable<K> {
    */
   count(key: K): number {
     this._validateKey(key);
-    return this.#core.get(key) ?? 0;
+    return this.#core.getNode(key)?.value ?? 0;
   }
 
   /**
@@ -133,9 +133,10 @@ export class TreeMultiSet<K = any, R = K> implements Iterable<K> {
     this._validateKey(key);
     this._validateCount(n);
     if (n === 0) return false;
-    const old = this.#core.get(key) ?? 0;
+    const node = this.#core.getNode(key);
+    const old = node?.value ?? 0;
     const next = old + n;
-    this.#core.set(key, next);
+    this.#core.set(node?.key ?? key, next);
     this._size += n;
     return true;
   }
@@ -153,12 +154,13 @@ export class TreeMultiSet<K = any, R = K> implements Iterable<K> {
   setCount(key: K, n: number): boolean {
     this._validateKey(key);
     this._validateCount(n);
-    const old = this.#core.get(key) ?? 0;
+    const node = this.#core.getNode(key);
+    const old = node?.value ?? 0;
     if (old === n) return false;
     if (n === 0) {
-      if (old !== 0) this.#core.delete(key);
+      if (node) this.#core.delete(node);
     } else {
-      this.#core.set(key, n);
+      this.#core.set(node?.key ?? key, n);
     }
     this._size += n - old;
     return true;
@@ -179,12 +181,13 @@ export class TreeMultiSet<K = any, R = K> implements Iterable<K> {
     this._validateKey(key);
     this._validateCount(n);
     if (n === 0) return false;
-    const old = this.#core.get(key) ?? 0;
+    const node = this.#core.getNode(key);
+    const old = node?.value ?? 0;
     if (old === 0) return false;
     const removed = Math.min(old, n);
     const next = old - removed;
-    if (next === 0) this.#core.delete(key);
-    else this.#core.set(key, next);
+    if (next === 0) this.#core.delete(node);
+    else this.#core.set(node?.key ?? key, next);
     this._size -= removed;
     return true;
   }
@@ -202,9 +205,10 @@ export class TreeMultiSet<K = any, R = K> implements Iterable<K> {
    */
   deleteAll(key: K): boolean {
     this._validateKey(key);
-    const old = this.#core.get(key) ?? 0;
+    const node = this.#core.getNode(key);
+    const old = node?.value ?? 0;
     if (old === 0) return false;
-    this.#core.delete(key);
+    this.#core.delete(node);
     this._size -= old;
     return true;
   }

test/unit/data-structures/binary-tree/tree-multi-set.test.ts

@@ -499,6 +499,32 @@ describe('TreeMultiSet', () => {
       expect(ms.size).toBe(3);
     });
 
+    it('handles comparator-equal non-identical objects without corrupting size', () => {
+      const ms = new TreeMultiSet<{ key: number }>([], {
+        comparator: (a, b) => a.key - b.key
+      });
+      const nodeA = { key: 1 };
+      const nodeB = { key: 1 };
+
+      ms.add(nodeA);
+      ms.add(nodeB);
+
+      expect(ms.size).toBe(2);
+      expect(ms.distinctSize).toBe(1);
+      expect(ms.count(nodeA)).toBe(2);
+      expect(ms.count(nodeB)).toBe(2);
+      expect(ms.toArray()).toEqual([nodeA, nodeA]);
+
+      while (ms.size > 0) {
+        const node = ms.first();
+        expect(node).toBe(nodeA);
+        expect(ms.delete(node!)).toBe(true);
+      }
+
+      expect(ms.first()).toBeUndefined();
+      expect(ms.distinctSize).toBe(0);
+    });
+
     it('rejects non-primitive types without comparator', () => {
       const ms = new TreeMultiSet<{ id: number }>();
       expect(() => ms.add({ id: 1 })).toThrow(TypeError);