Remove custom heap and use bintrees

3 files changed, 14 insertions, 102 deletions

diff --git a/packages/0x.js/src/order_watcher/expiration_watcher.ts b/packages/0x.js/src/order_watcher/expiration_watcher.ts
index 71199e75f..7d6f8df65 100644
--- a/packages/0x.js/src/order_watcher/expiration_watcher.ts
+++ b/packages/0x.js/src/order_watcher/expiration_watcher.ts
@@ -1,9 +1,9 @@
 import * as _ from 'lodash';
 import {BigNumber} from 'bignumber.js';
+import {RBTree} from 'bintrees';
 import {utils} from '../utils/utils';
 import {intervalUtils} from '../utils/interval_utils';
 import {SignedOrder, ZeroExError} from '../types';
-import {Heap} from '../utils/heap';
 import {ZeroEx} from '../0x';
 
 const DEFAULT_EXPIRATION_MARGIN_MS = 0;
@@ -14,7 +14,7 @@ const DEFAULT_ORDER_EXPIRATION_CHECKING_INTERVAL_MS = 50;
  * It stores them in a min heap by expiration time and checks for expired ones every `orderExpirationCheckingIntervalMs`
  */
 export class ExpirationWatcher {
-    private orderHashHeapByExpiration: Heap<string>;
+    private orderHashRBTreeByExpiration: RBTree<string>;
     private expiration: {[orderHash: string]: BigNumber} = {};
     private callbackIfExists?: (orderHash: string) => void;
     private orderExpirationCheckingIntervalMs: number;
@@ -27,7 +27,8 @@ export class ExpirationWatcher {
         this.orderExpirationCheckingIntervalMs = expirationMarginIfExistsMs ||
                                                  DEFAULT_ORDER_EXPIRATION_CHECKING_INTERVAL_MS;
         const scoreFunction = (orderHash: string) => this.expiration[orderHash].toNumber();
-        this.orderHashHeapByExpiration = new Heap(scoreFunction);
+        const comparator = (lhs: string, rhs: string) => scoreFunction(lhs) - scoreFunction(rhs);
+        this.orderHashRBTreeByExpiration = new RBTree(comparator);
     }
     public subscribe(callback: (orderHash: string) => void): void {
         if (!_.isUndefined(this.callbackIfExists)) {
@@ -48,20 +49,23 @@ export class ExpirationWatcher {
     }
     public addOrder(orderHash: string, expirationUnixTimestampMs: BigNumber): void {
         this.expiration[orderHash] = expirationUnixTimestampMs;
-        // We don't remove hashes from the heap on order remove because it's slow (linear).
-        // We just skip them later if the order was already removed from the order watcher.
-        this.orderHashHeapByExpiration.push(orderHash);
+        this.orderHashRBTreeByExpiration.insert(orderHash);
+    }
+    public removeOrder(orderHash: string): void {
+        this.orderHashRBTreeByExpiration.remove(orderHash);
+        delete this.expiration[orderHash];
     }
     private pruneExpiredOrders(): void {
         const currentUnixTimestampMs = utils.getCurrentUnixTimestampMs();
         while (
-            this.orderHashHeapByExpiration.size() !== 0 &&
-            this.expiration[this.orderHashHeapByExpiration.head()].lessThan(
+            this.orderHashRBTreeByExpiration.size !== 0 &&
+            this.expiration[this.orderHashRBTreeByExpiration.min()].lessThan(
                 currentUnixTimestampMs.plus(this.expirationMarginMs),
             ) &&
             !_.isUndefined(this.callbackIfExists)
         ) {
-            const orderHash = this.orderHashHeapByExpiration.pop();
+            const orderHash = this.orderHashRBTreeByExpiration.min();
+            this.orderHashRBTreeByExpiration.remove(orderHash);
             delete this.expiration[orderHash];
             this.callbackIfExists(orderHash);
         }
diff --git a/packages/0x.js/src/order_watcher/order_state_watcher.ts b/packages/0x.js/src/order_watcher/order_state_watcher.ts
index 3659fc6e2..975679f57 100644
--- a/packages/0x.js/src/order_watcher/order_state_watcher.ts
+++ b/packages/0x.js/src/order_watcher/order_state_watcher.ts
@@ -95,7 +95,6 @@ export class OrderStateWatcher {
         assert.isValidSignature(orderHash, signedOrder.ecSignature, signedOrder.maker);
         this._orderByOrderHash[orderHash] = signedOrder;
         this.addToDependentOrderHashes(signedOrder, orderHash);
-        // We don't remove orders from expirationWatcher because heap removal is linear. We just skip it later
         const expirationUnixTimestampMs = signedOrder.expirationUnixTimestampSec.times(1000);
         this._expirationWatcher.addOrder(orderHash, expirationUnixTimestampMs);
     }
@@ -111,6 +110,7 @@ export class OrderStateWatcher {
         }
         delete this._orderByOrderHash[orderHash];
         this.removeFromDependentOrderHashes(signedOrder.maker, signedOrder.makerTokenAddress, orderHash);
+        this._expirationWatcher.removeOrder(orderHash);
     }
     /**
      * Starts an orderStateWatcher subscription. The callback will be called every time a watched order's
diff --git a/packages/0x.js/src/utils/heap.ts b/packages/0x.js/src/utils/heap.ts
deleted file mode 100644
index e262266d4..000000000
--- a/packages/0x.js/src/utils/heap.ts
+++ /dev/null
@@ -1,92 +0,0 @@
-// Based on Original JavaScript Code from  Marijn Haverbeke (http://eloquentjavascript.net/1st_edition/appendix2.html)
-export class Heap<T> {
-    private content: T[];
-    private scoreFunction: (x: T) => number;
-    constructor(scoreFunction: (x: T) => number) {
-        this.content = [];
-        this.scoreFunction = scoreFunction;
-    }
-    public push(element: T) {
-        this.content.push(element);
-        this.bubbleUp(this.content.length - 1);
-    }
-    public size(): number {
-        const size = this.content.length;
-        return size;
-    }
-    public head(): T {
-        const head = this.content[0];
-        return head;
-    }
-    public pop(): T {
-        const head = this.content[0];
-        const end = this.content.pop();
-        if (this.content.length > 0) {
-            this.content[0] = end as T;
-            this.sinkDown(0);
-        }
-        return head;
-    }
-    private bubbleUp(n: number) {
-        // Fetch the element that has to be moved.
-        const element = this.content[n];
-        const score = this.scoreFunction(element);
-        // When at 0, an element can not go up any further.
-        while (n > 0) {
-            // Compute the parent element's index, and fetch it.
-            const parentN = Math.floor((n + 1) / 2) - 1;
-            const parent = this.content[parentN];
-            // If the parent has a lesser score, things are in order and we
-            // are done.
-            if (score >= this.scoreFunction(parent)) {
-                break;
-            }
-
-            // Otherwise, swap the parent with the current element and
-            // continue.
-            this.content[parentN] = element;
-            this.content[n] = parent;
-            n = parentN;
-        }
-    }
-    private sinkDown(n: number) {
-        // Look up the target element and its score.
-        const length = this.content.length;
-        const element = this.content[n];
-        const elemScore = this.scoreFunction(element);
-
-        while (true) {
-            // Compute the indices of the child elements.
-            const child2N = (n + 1) * 2;
-            const child1N = child2N - 1;
-            // This is used to store the new position of the element, if any.
-            let swap = n;
-            let child1Score;
-            let child2Score;
-            // If the first child exists (is inside the array)...
-            if (child1N < length) {
-                // Look it up and compute its score.
-                const child1 = this.content[child1N];
-                child1Score = this.scoreFunction(child1);
-                // If the score is less than our element's, we need to swap.
-                if (child1Score < elemScore) {
-                    swap = child1N;
-                }
-                // Do the same checks for the other child.
-                if (child2N < length) {
-                    const child2 = this.content[child2N];
-                    child2Score = this.scoreFunction(child2);
-                    if (child2Score < (swap == null ? elemScore : child1Score)) {
-                        swap = child2N;
-                    }
-                }
-            }
-            if (swap === n) {
-                break;
-            }
-            this.content[n] = this.content[swap];
-            this.content[swap] = element;
-            n = swap;
-        }
-    }
-}