feat(Heap): enhance element search methods to support custom comparison functions

Should fix:
#669

Author: ignlg

src/Heap.ts

@@ -439,7 +439,7 @@ export class Heap<T> implements Iterable<T> {
   }
 
   /**
-   * Get the index of the first occurrence of the element in the heap (using the comparator).
+   * Get the index of the first occurrence of the element in the heap.
    * @param  {any}      element    Element to be found
    * @param  {Function} callbackFn Optional comparison function, receives (element, needle)
    * @return {Number}              Index or -1 if not found
@@ -448,6 +448,17 @@ export class Heap<T> implements Iterable<T> {
     if (this.heapArray.length === 0) {
       return -1;
     }
+    // When a custom callback is provided, we must search all elements
+    // because the callback may not be consistent with this.compare
+    if (callbackFn !== Heap.defaultIsEqual) {
+      for (let i = 0; i < this.heapArray.length; i++) {
+        if (callbackFn(this.heapArray[i], element)) {
+          return i;
+        }
+      }
+      return -1;
+    }
+    // Default case: use heap structure optimization
     const indexes: number[] = [];
     let currentIndex = 0;
     while (currentIndex < this.heapArray.length) {
@@ -463,7 +474,7 @@ export class Heap<T> implements Iterable<T> {
   }
 
   /**
-   * Get the indexes of the every occurrence of the element in the heap (using the comparator).
+   * Get the indexes of every occurrence of the element in the heap.
    * @param  {any}      element    Element to be found
    * @param  {Function} callbackFn Optional comparison function, receives (element, needle)
    * @return {Array}               Array of indexes or empty array if not found
@@ -472,6 +483,18 @@ export class Heap<T> implements Iterable<T> {
     if (this.heapArray.length === 0) {
       return [];
     }
+    // When a custom callback is provided, we must search all elements
+    // because the callback may not be consistent with this.compare
+    if (callbackFn !== Heap.defaultIsEqual) {
+      const foundIndexes: number[] = [];
+      for (let i = 0; i < this.heapArray.length; i++) {
+        if (callbackFn(this.heapArray[i], element)) {
+          foundIndexes.push(i);
+        }
+      }
+      return foundIndexes;
+    }
+    // Default case: use heap structure optimization
     const indexes: number[] = [];
     const foundIndexes: number[] = [];
     let currentIndex = 0;
@@ -618,6 +641,25 @@ export class Heap<T> implements Iterable<T> {
       this.pop();
       return true;
     }
+    // When a custom callback is provided, we must search all elements
+    // because the callback may not be consistent with this.compare
+    if (callbackFn !== Heap.defaultIsEqual) {
+      const idx = this.indexOf(o, callbackFn);
+      if (idx === -1) {
+        return false;
+      }
+      if (idx === 0) {
+        this.pop();
+      } else if (idx === this.heapArray.length - 1) {
+        this.heapArray.pop();
+      } else {
+        this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
+        this._sortNodeUp(idx);
+        this._sortNodeDown(idx);
+      }
+      return true;
+    }
+    // Default case: use heap structure optimization
     const queue = [0];
     while (queue.length) {
       const idx = queue.shift() as number;

tests/heap/heap-public-methods.test.ts

@@ -641,3 +641,86 @@ describe('Heap instances', function () {
     });
   });
 });
+
+describe('Heap with custom IsEqual function', function () {
+  describe('when using object heap with custom IsEqual', function () {
+    let heap: Heap<{ value: number }>;
+
+    beforeEach(function () {
+      heap = new Heap<{ value: number }>((a, b) => b.value - a.value);
+      heap.add({ value: 1 });
+      heap.add({ value: 2 });
+      heap.add({ value: 3 });
+    });
+
+    describe('#contains(element, fn)', function () {
+      it('should find element using custom IsEqual that ignores the needle parameter', function () {
+        // Using a custom IsEqual that only checks the element, ignoring the needle
+        expect(heap.contains(0 as any, (a) => a.value === 2)).toBe(true);
+        expect(heap.contains(0 as any, (a) => a.value === 1)).toBe(true);
+        expect(heap.contains(0 as any, (a) => a.value === 3)).toBe(true);
+        expect(heap.contains(0 as any, (a) => a.value === 99)).toBe(false);
+      });
+    });
+
+    describe('#indexOf(element, fn)', function () {
+      it('should find index using custom IsEqual that ignores the needle parameter', function () {
+        // Should find the element regardless of what needle is passed
+        expect(heap.indexOf(0 as any, (a) => a.value === 2)).not.toBe(-1);
+        expect(heap.indexOf(0 as any, (a) => a.value === 1)).not.toBe(-1);
+        expect(heap.indexOf(0 as any, (a) => a.value === 3)).not.toBe(-1);
+        expect(heap.indexOf(0 as any, (a) => a.value === 99)).toBe(-1);
+      });
+
+      it('should find correct element at returned index', function () {
+        const idx = heap.indexOf(0 as any, (a) => a.value === 2);
+        expect(heap.heapArray[idx].value).toBe(2);
+      });
+    });
+
+    describe('#indexOfEvery(element, fn)', function () {
+      it('should find all indexes using custom IsEqual', function () {
+        heap.add({ value: 2 }); // Add duplicate
+        const indexes = heap.indexOfEvery(0 as any, (a) => a.value === 2);
+        expect(indexes.length).toBe(2);
+        for (const idx of indexes) {
+          expect(heap.heapArray[idx].value).toBe(2);
+        }
+      });
+    });
+
+    describe('#remove(element, fn)', function () {
+      it('should remove element using custom IsEqual that ignores the needle parameter', function () {
+        // This was the failing case from the issue
+        expect(heap.remove(0 as any, (a) => a.value === 2)).toBe(true);
+        expect(heap.length).toBe(2);
+        expect(heap.contains(0 as any, (a) => a.value === 2)).toBe(false);
+        expect(heap.check()).not.toBeDefined();
+      });
+
+      it('should remove any element regardless of heap position', function () {
+        // Test removing elements at different positions
+        const heap2 = new Heap<{ value: number }>((a, b) => b.value - a.value);
+        heap2.add({ value: 1 });
+        heap2.add({ value: 2 });
+        heap2.add({ value: 3 });
+
+        expect(heap2.remove(0 as any, (a) => a.value === 1)).toBe(true);
+        expect(heap2.check()).not.toBeDefined();
+
+        const heap3 = new Heap<{ value: number }>((a, b) => b.value - a.value);
+        heap3.add({ value: 1 });
+        heap3.add({ value: 2 });
+        heap3.add({ value: 3 });
+
+        expect(heap3.remove(0 as any, (a) => a.value === 3)).toBe(true);
+        expect(heap3.check()).not.toBeDefined();
+      });
+
+      it('should return false when element is not found', function () {
+        expect(heap.remove(0 as any, (a) => a.value === 99)).toBe(false);
+        expect(heap.length).toBe(3);
+      });
+    });
+  });
+});