refactor: make SortedMap more efficient

Author: dirkmc

src/utils/sorted-map.js

@@ -6,50 +6,76 @@
 class SortedMap extends Map {
   /**
    * @param {Array<k, v>} [entries]
-   * @param {function(a, b)} [cmp]
+   * @param {function(a, b)} [cmp] compares [k1, v1] to [k2, v2]
    */
   constructor (entries, cmp) {
-    super(entries)
-    this._cmp = cmp
+    super()
+    this._cmp = cmp || this._defaultSort
     this._keys = []
-  }
-
-  // Performance is O(N log(N)).
-  // Could be O(log(N)) if we use a sorted array. Fix this if performance isn't
-  // good enough.
-  set (k, data) {
-    super.set(k, data)
-    const i = this._keys.indexOf(k)
-    if (i < 0) {
-      this._keys.push(k)
-      this._keys.sort(this._kvCmp.bind(this))
+    for (const [k, v] of entries || []) {
+      this.set(k, v)
     }
   }
 
-  clear () {
-    super.clear()
-    this._keys = []
-  }
-
   /**
-   * Call update to manually trigger a sort.
+   * Call update to update the position of the key when it should change.
    * For example if the compare function sorts by the priority field, and the
    * priority changes, call update.
    *
    * @param {Object} [k] the key corresponding to the entry whose position
    * should be updated.
    */
   update (k) {
-    this._keys.sort(this._kvCmp.bind(this))
+    if (this.has(k)) {
+      this.set(k, this.get(k))
+    }
   }
 
-  // Same performance comments as set()
-  delete (k) {
-    const i = this._keys.indexOf(k)
-    if (i >= 0) {
+  set (k, v) {
+    // If the key is already in the map, remove it from the ordering and
+    // re-insert it below
+    if (this.has(k)) {
+      const i = this._find(k)
       this._keys.splice(i, 1)
-      super.delete(k)
     }
+
+    // Update / insert the k/v into the map
+    super.set(k, v)
+
+    // Find the correct position of the newly inserted k/v in the order
+    const i = this._find(k)
+    this._keys.splice(i, 0, k)
+  }
+
+  clear () {
+    super.clear()
+    this._keys = []
+  }
+
+  delete (k) {
+    if (!this.has(k)) {
+      return
+    }
+    const i = this._find(k)
+    this._keys.splice(i, 1)
+    super.delete(k)
+  }
+
+  _find (k) {
+    let lower = 0
+    let upper = this._keys.length
+    while (lower < upper) {
+      const pivot = (lower + upper) >>> 1 // lower + (upper - lower) / 2
+      const cmp = this._kCmp(this._keys[pivot], k)
+      if (cmp < 0) { // pivot < k
+        lower = pivot + 1
+      } else if (cmp > 0) { // pivot > k
+        upper = pivot
+      } else { // pivot == k
+        return pivot
+      }
+    }
+    return lower
   }
 
   * keys () {
@@ -84,14 +110,13 @@ class SortedMap extends Map {
     }
   }
 
-  _kvCmp (a, b) {
-    // By default sort by built-in compare
-    if (!this._cmp) {
-      if (a < b) return -1
-      if (b < a) return 1
-      return 0
-    }
+  _defaultSort (a, b) {
+    if (a[0] < b[0]) return -1
+    if (b[0] < a[0]) return 1
+    return 0
+  }
 
+  _kCmp (a, b) {
     return this._cmp(
       [a, this.get(a)],
       [b, this.get(b)]

test/utils.spec.js

@@ -200,6 +200,9 @@ describe('utils spec', function () {
 
       expect(sm.get('two')).to.be.undefined()
       expect(sm.size).to.eql(2)
+
+      sm.delete('two')
+      expect(sm.size).to.eql(2)
     })
 
     it('clear', () => {
@@ -235,26 +238,100 @@ describe('utils spec', function () {
       expect(collected).to.eql([...sm])
     })
 
-    it('custom order', () => {
-      const sm = new SortedMap([], (a, b) => b[1].priority - a[1].priority)
-
-      const data1 = { k: 'v1', priority: 1 }
-      const data2 = { k: 'v2', priority: 3 }
-      const data3 = { k: 'v3', priority: 2 }
-      sm.set('one', data1)
-      sm.set('two', data2)
-      sm.set('three', data3)
-
-      expect([...sm.keys()]).to.eql(['two', 'three', 'one'])
-      expect([...sm.values()].map(v => v.k)).to.eql(['v2', 'v3', 'v1'])
-
-      // After changing data that affects the sort order, need to call update
-      // to actually trigger the sort
-      data3.priority = 5
-      sm.update('three')
-
-      expect([...sm.keys()]).to.eql(['three', 'two', 'one'])
-      expect([...sm.values()].map(v => v.k)).to.eql(['v3', 'v2', 'v1'])
+    describe('custom order', () => {
+      const prioritySort = (a, b) => b[1].priority - a[1].priority
+
+      it('forward', () => {
+        const sm = new SortedMap([
+          ['low', { priority: 1 }],
+          ['high', { priority: 2 }]
+        ], prioritySort)
+        expect([...sm.keys()]).to.eql(['high', 'low'])
+      })
+
+      it('backward', () => {
+        const sm = new SortedMap([
+          ['high', { priority: 2 }],
+          ['low', { priority: 1 }]
+        ], prioritySort)
+        expect([...sm.keys()]).to.eql(['high', 'low'])
+      })
+
+      it('insert start', () => {
+        const sm = new SortedMap([
+          ['mid', { priority: 2 }],
+          ['low', { priority: 1 }],
+          ['high', { priority: 3 }]
+        ], prioritySort)
+        expect([...sm.keys()]).to.eql(['high', 'mid', 'low'])
+      })
+
+      it('insert end', () => {
+        const sm = new SortedMap([
+          ['low', { priority: 1 }],
+          ['mid', { priority: 2 }],
+          ['high', { priority: 3 }]
+        ], prioritySort)
+        expect([...sm.keys()]).to.eql(['high', 'mid', 'low'])
+      })
+
+      it('insert middle', () => {
+        const sm = new SortedMap([
+          ['low', { priority: 1 }],
+          ['high', { priority: 3 }],
+          ['mid', { priority: 2 }]
+        ], prioritySort)
+        expect([...sm.keys()]).to.eql(['high', 'mid', 'low'])
+      })
+
+      it('insert same priority start', () => {
+        const sm = new SortedMap([
+          ['low', { priority: 1 }],
+          ['high-a', { priority: 3 }],
+          ['high-b', { priority: 3 }]
+        ], prioritySort)
+        expect([...sm.keys()].map(s => s.substring(0, 4))).to.eql(['high', 'high', 'low'])
+      })
+
+      it('insert same priority end', () => {
+        const sm = new SortedMap([
+          ['hi', { priority: 3 }],
+          ['low-a', { priority: 1 }],
+          ['low-b', { priority: 1 }]
+        ], prioritySort)
+        expect([...sm.keys()].map(s => s.substring(0, 3))).to.eql(['hi', 'low', 'low'])
+      })
+
+      it('insert same key', () => {
+        const sm = new SortedMap([
+          ['low', { priority: 1 }],
+          ['high', { priority: 3 }],
+          ['high', { priority: 4 }]
+        ], prioritySort)
+        expect([...sm.keys()]).to.eql(['high', 'low'])
+      })
+
+      it('update', () => {
+        const sm = new SortedMap([], prioritySort)
+
+        const data1 = { k: 'v1', priority: 1 }
+        const data2 = { k: 'v2', priority: 3 }
+        const data3 = { k: 'v3', priority: 2 }
+        sm.set('one', data1)
+        sm.set('two', data2)
+        sm.set('three', data3)
+
+        expect([...sm.keys()]).to.eql(['two', 'three', 'one'])
+        expect([...sm.values()].map(v => v.k)).to.eql(['v2', 'v3', 'v1'])
+
+        // After changing data that affects the sort order, need to call update
+        // to actually trigger the sort
+        data3.priority = 5
+        sm.update('three')
+
+        expect([...sm.keys()]).to.eql(['three', 'two', 'one'])
+        expect([...sm.values()].map(v => v.k)).to.eql(['v3', 'v2', 'v1'])
+      })
     })
   })
 })