Restructure of the Heap class

Author: isoos

app/lib/search/heap.dart

@@ -8,15 +8,35 @@
 /// than or equal to its children's values.
 ///
 /// The provided comparator decides which kind of heap is being built.
-class _Heap<T> {
+class Heap<T> {
   final Comparator<T> _compare;
   final _items = <T>[];
+  bool _isValid = true;
 
-  _Heap(this._compare);
+  Heap(this._compare);
 
   int get length => _items.length;
 
-  void _pushDown(int index) {
+  /// Collects [item] and adds it to the end of the internal list, marks the [Heap]
+  /// as non-valid.
+  ///
+  /// A separate operation may trigger the restoration of the heap proprerty.
+  void collect(T item) {
+    _items.add(item);
+    _isValid = false;
+  }
+
+  /// Collects [items] and adds them the end of the internal list, marks the [Heap]
+  /// as non-valid.
+  ///
+  /// A separate operation may trigger the restoration of the heap proprerty.
+  void collectAll(Iterable<T> items) {
+    _items.addAll(items);
+    _isValid = false;
+  }
+
+  /// Ensures that the tree structre below the [index] is a valid heap.
+  void _heapify(int index) {
     final maxLength = _items.length;
     final item = _items[index];
     while (index < maxLength) {
@@ -39,60 +59,54 @@ class _Heap<T> {
     }
   }
 
+  /// (Re-)builds the heap property if needed.
+  void _buildHeapIfNeeded() {
+    if (_isValid) {
+      assert(_isValidHeap());
+      return;
+    }
+
+    if (_items.isEmpty) {
+      _isValid = true;
+      return;
+    }
+    for (var i = (_items.length >> 1); i >= 0; i--) {
+      _heapify(i);
+    }
+
+    assert(_isValidHeap());
+    _isValid = true;
+  }
+
+  /// Verifies the heap property is true for all items.
   bool _isValidHeap() {
     for (var i = 1; i < _items.length; i++) {
       final parentIndex = (i - 1) >> 1;
       if (_compare(_items[parentIndex], _items[i]) > 0) {
-        print(parentIndex);
-        print(_items);
         return false;
       }
     }
     return true;
   }
-}
-
-/// Builds a sorted list of the top-k items using the provided comparator.
-///
-/// The algorithm collects all items, builds a max-heap in O(N) steps, and
-/// then selects the top-k items by removing the largest item from the heap
-/// and restoring the heap property again in O(k * log(N)) steps.
-class TopKSortedListBuilder<T> {
-  final int _k;
-  final _Heap<T> _heap;
-
-  TopKSortedListBuilder(this._k, Comparator<T> compare)
-      : _heap = _Heap<T>(compare);
-
-  void addAll(Iterable<T> items) {
-    for (final item in items) {
-      add(item);
-    }
-  }
-
-  void add(T item) {
-    _heap._items.add(item);
-  }
 
-  /// Gets and removes the top-k items from the current list.
-  Iterable<T> getTopK() sync* {
-    if (_heap._items.isEmpty) {
-      return;
-    }
-    for (var i = (_heap._items.length >> 1); i >= 0; i--) {
-      _heap._pushDown(i);
-    }
-    assert(_heap._isValidHeap());
-    var count = _k;
-    while (count > 0 && _heap._items.isNotEmpty) {
-      yield _heap._items[0];
-      count--;
-      final last = _heap._items.removeLast();
-      if (_heap._items.isEmpty) {
+  /// Creates a sorted list of the top-k items and removes them from the [Heap].
+  ///
+  /// The algorithm builds a max-heap in `O(N)` steps on the already collected items,
+  /// and then selects the top-k items by removing the largest item from the [Heap]
+  /// and restoring the heap property again in `O(k * log(N))` steps.
+  Iterable<T> getAndRemoveTopK(int k) sync* {
+    _buildHeapIfNeeded();
+    var remaining = k;
+    while (remaining > 0 && _items.isNotEmpty) {
+      yield _items[0];
+      remaining--;
+      final last = _items.removeLast();
+      if (_items.isEmpty) {
         break;
       }
-      _heap._items[0] = last;
-      _heap._pushDown(0);
+      _items[0] = last;
+      _heapify(0);
     }
+    assert(_isValidHeap());
   }
 }

app/lib/search/mem_index.dart

@@ -528,7 +528,7 @@ class InMemoryPackageIndex {
     /// Return (and sort) only the top-k results.
     required int topK,
   }) {
-    final builder = TopKSortedListBuilder<int>(topK, (aIndex, bIndex) {
+    final heap = Heap<int>((aIndex, bIndex) {
       if (aIndex == bestNameIndex) return -1;
       if (bIndex == bestNameIndex) return 1;
       final aScore = score.getValue(aIndex);
@@ -541,9 +541,9 @@ class InMemoryPackageIndex {
     for (var i = 0; i < score.length; i++) {
       final value = score.getValue(i);
       if (value <= 0.0 && i != bestNameIndex) continue;
-      builder.add(i);
+      heap.collect(i);
     }
-    return builder.getTopK().map((i) => IndexedPackageHit(
+    return heap.getAndRemoveTopK(topK).map((i) => IndexedPackageHit(
         i, PackageHit(package: score.keys[i], score: score.getValue(i))));
   }
 

app/lib/search/token_index.dart

@@ -315,15 +315,15 @@ class IndexedScore<K> {
 
   Map<K, double> top(int count, {double? minValue}) {
     minValue ??= 0.0;
-    final builder = TopKSortedListBuilder<int>(
-        count, (a, b) => -_values[a].compareTo(_values[b]));
+    final heap = Heap<int>((a, b) => -_values[a].compareTo(_values[b]));
     for (var i = 0; i < length; i++) {
       final v = _values[i];
       if (v < minValue) continue;
-      builder.add(i);
+      heap.collect(i);
     }
-    return Map.fromEntries(
-        builder.getTopK().map((i) => MapEntry(_keys[i], _values[i])));
+    return Map.fromEntries(heap
+        .getAndRemoveTopK(count)
+        .map((i) => MapEntry(_keys[i], _values[i])));
   }
 
   Map<K, double> toMap() {

app/test/search/heap_test.dart

@@ -12,32 +12,32 @@ void main() {
     int compare(int a, int b) => -a.compareTo(b);
 
     test('no items', () {
-      final builder = TopKSortedListBuilder(5, compare);
-      expect(builder.getTopK().toList(), []);
+      final heap = Heap(compare);
+      expect(heap.getAndRemoveTopK(5).toList(), []);
     });
 
     test('single item', () {
-      final builder = TopKSortedListBuilder(5, compare);
-      builder.add(1);
-      expect(builder.getTopK().toList(), [1]);
+      final heap = Heap(compare);
+      heap.collect(1);
+      expect(heap.getAndRemoveTopK(5).toList(), [1]);
     });
 
     test('three items ascending', () {
-      final builder = TopKSortedListBuilder(5, compare);
-      builder.addAll([1, 2, 3]);
-      expect(builder.getTopK().toList(), [3, 2, 1]);
+      final builder = Heap(compare);
+      builder.collectAll([1, 2, 3]);
+      expect(builder.getAndRemoveTopK(5).toList(), [3, 2, 1]);
     });
 
     test('three items descending', () {
-      final builder = TopKSortedListBuilder(5, compare);
-      builder.addAll([3, 2, 1]);
-      expect(builder.getTopK().toList(), [3, 2, 1]);
+      final heap = Heap(compare);
+      heap.collectAll([3, 2, 1]);
+      expect(heap.getAndRemoveTopK(5).toList(), [3, 2, 1]);
     });
 
     test('10 items + repeated', () {
-      final builder = TopKSortedListBuilder(5, compare);
-      builder.addAll([1, 10, 2, 9, 3, 8, 4, 7, 6, 5, 9]);
-      expect(builder.getTopK().toList(), [10, 9, 9, 8, 7]);
+      final heap = Heap(compare);
+      heap.collectAll([1, 10, 2, 9, 3, 8, 4, 7, 6, 5, 9]);
+      expect(heap.getAndRemoveTopK(5).toList(), [10, 9, 9, 8, 7]);
     });
 
     test('randomized verification', () {
@@ -46,13 +46,13 @@ void main() {
         final length = 1000 + r.nextInt(1000);
         final k = 10 + r.nextInt(200);
         final items = List.generate(length, (i) => i);
-        final b1 = TopKSortedListBuilder(k, compare)..addAll(items);
-        final r1 = b1.getTopK().toList();
+        final b1 = Heap(compare)..collectAll(items);
+        final r1 = b1.getAndRemoveTopK(k).toList();
         expect(r1, List.generate(k, (i) => length - 1 - i));
 
         items.shuffle(r);
-        final b2 = TopKSortedListBuilder(k, compare)..addAll(items);
-        final r2 = b2.getTopK().toList();
+        final b2 = Heap(compare)..collectAll(items);
+        final r2 = b2.getAndRemoveTopK(k).toList();
         expect(r2, r1);
       }
     });