Add HeapSet._sorted internal flag (#6949)

Author: gjoseph92

distributed/collections.py

@@ -38,18 +38,20 @@ class HeapSet(MutableSet[T]):
     Values must be compatible with :mod:`weakref`.
     """
 
-    __slots__ = ("key", "_data", "_heap", "_inc")
+    __slots__ = ("key", "_data", "_heap", "_inc", "_sorted")
     key: Callable[[T], Any]
     _data: set[T]
     _inc: int
     _heap: list[tuple[Any, int, weakref.ref[T]]]
+    _sorted: bool
 
     def __init__(self, *, key: Callable[[T], Any]):
         # FIXME https://github.com/python/mypy/issues/708
         self.key = key  # type: ignore
         self._data = set()
         self._inc = 0
         self._heap = []
+        self._sorted = True
 
     def __repr__(self) -> str:
         return f"<{type(self).__name__}: {len(self)} items>"
@@ -68,6 +70,7 @@ def _unpickle(
         self._inc = inc
         self._heap = [(k, i, weakref.ref(v)) for k, i, v in heap]
         heapq.heapify(self._heap)
+        self._sorted = not heap
         return self
 
     def __contains__(self, value: object) -> bool:
@@ -82,13 +85,14 @@ def add(self, value: T) -> None:
         k = self.key(value)  # type: ignore
         vref = weakref.ref(value)
         heapq.heappush(self._heap, (k, self._inc, vref))
+        self._sorted = False
         self._data.add(value)
         self._inc += 1
 
     def discard(self, value: T) -> None:
         self._data.discard(value)
         if not self._data:
-            self._heap.clear()
+            self.clear()
 
     def peek(self) -> T:
         """Return the smallest element without removing it"""
@@ -99,6 +103,7 @@ def peek(self) -> T:
             if value in self._data:
                 return value
             heapq.heappop(self._heap)
+            self._sorted = False
 
     def peekn(self, n: int) -> Iterator[T]:
         "Iterator over the N smallest elements. This is O(1) for n == 1, O(n*logn) otherwise."
@@ -119,9 +124,12 @@ def pop(self) -> T:
             raise KeyError("pop from an empty set")
         while True:
             _, _, vref = heapq.heappop(self._heap)
+            self._sorted = False
             value = vref()
             if value in self._data:
                 self._data.discard(value)
+                if not self._data:
+                    self.clear()
                 return value
 
     def peekright(self) -> T:
@@ -147,6 +155,8 @@ def popright(self) -> T:
             value = vref()
             if value in self._data:
                 self._data.discard(value)
+                if not self._data:
+                    self.clear()
                 return value
 
     def __iter__(self) -> Iterator[T]:
@@ -160,11 +170,15 @@ def sorted(self) -> Iterator[T]:
         elements in order, from smallest to largest according to the key and insertion
         order.
         """
-        for _, _, vref in sorted(self._heap):
+        if not self._sorted:
+            self._heap.sort()  # A sorted list maintains the heap invariant
+            self._sorted = True
+        for _, _, vref in self._heap:
             value = vref()
             if value in self._data:
                 yield value
 
     def clear(self) -> None:
         self._data.clear()
         self._heap.clear()
+        self._sorted = True

distributed/tests/test_collections.py

@@ -176,6 +176,91 @@ def __init__(self, i):
     assert set(heap) == {cx}
 
 
+def assert_heap_sorted(heap: HeapSet) -> None:
+    assert heap._sorted
+    assert heap._heap == sorted(heap._heap)
+
+
+def test_heapset_sorted_flag_left():
+    heap = HeapSet(key=operator.attrgetter("i"))
+    assert heap._sorted
+    c1 = C("1", 1)
+    c2 = C("2", 2)
+    c3 = C("3", 3)
+    c4 = C("4", 4)
+
+    heap.add(c4)
+    assert not heap._sorted
+    heap.add(c3)
+    heap.add(c2)
+    heap.add(c1)
+
+    list(heap.sorted())
+    assert_heap_sorted(heap)
+
+    # `peek` maintains sort if first element is not discarded
+    assert heap.peek() is c1
+    assert_heap_sorted(heap)
+
+    # `pop` always de-sorts
+    assert heap.pop() is c1
+    assert not heap._sorted
+
+    list(heap.sorted())
+
+    # discard first element
+    heap.discard(c2)
+    assert heap.peek() is c3
+    assert not heap._sorted
+
+    # popping the last element resets the sorted flag
+    assert heap.pop() is c3
+    assert heap.pop() is c4
+    assert not heap
+    assert_heap_sorted(heap)
+
+    # discarding`` the last element resets the sorted flag
+    heap.add(c1)
+    heap.add(c2)
+    assert not heap._sorted
+    heap.discard(c1)
+    assert not heap._sorted
+    heap.discard(c2)
+    assert not heap
+    assert_heap_sorted(heap)
+
+
+def test_heapset_sorted_flag_right():
+    "Verify right operations don't affect sortedness"
+    heap = HeapSet(key=operator.attrgetter("i"))
+    c1 = C("1", 1)
+    c2 = C("2", 2)
+    c3 = C("3", 3)
+
+    heap.add(c2)
+    heap.add(c3)
+    heap.add(c1)
+
+    assert not heap._sorted
+    list(heap.sorted())
+    assert_heap_sorted(heap)
+
+    assert heap.peekright() is c3
+    assert_heap_sorted(heap)
+    assert heap.popright() is c3
+    assert_heap_sorted(heap)
+    assert heap.popright() is c2
+    assert_heap_sorted(heap)
+
+    heap.add(c2)
+    assert not heap._sorted
+    assert heap.popright() is c2
+    assert not heap._sorted
+    assert heap.popright() is c1
+    assert not heap
+    assert_heap_sorted(heap)
+
+
 @pytest.mark.parametrize("peek", [False, True])
 def test_heapset_popright(peek):
     heap = HeapSet(key=operator.attrgetter("i"))
@@ -231,8 +316,11 @@ def test_heapset_pickle():
         if random.random() > 0.7:
             heap.remove(c)
 
+    list(heap.sorted())  # trigger sort
+    assert heap._sorted
     heap2 = pickle.loads(pickle.dumps(heap))
     assert len(heap) == len(heap2)
+    assert not heap2._sorted  # re-heapification may have broken the sort
     # Test that the heap has been re-heapified upon unpickle
     assert len(heap2._heap) < len(heap._heap)
     while heap: