Enhance pqueue so we can remove items from the middle.

This changes the pqueue API by requiring an additional int in every
structure that we store in a pqueue to hold the index of that structure
within the heap.

src/common/container.c

@@ -605,6 +605,38 @@ smartlist_uniq_strings(smartlist_t *sl)
 /* Heap-based priority queue implementation for O(lg N) insert and remove.
  * Recall that the heap property is that, for every index I, h[I] <
  * H[LEFT_CHILD[I]] and h[I] < H[RIGHT_CHILD[I]].
+ *
+ * For us to remove items other than the topmost item, each item must store
+ * its own index within the heap.  When calling the pqueue functions, tell
+ * them about the offset of the field that stores the index within the item.
+ *
+ * Example:
+ *
+ *   typedef struct timer_t {
+ *     struct timeval tv;
+ *     int heap_index;
+ *   } timer_t;
+ *
+ *   static int compare(const void *p1, const void *p2) {
+ *     const timer_t *t1 = p1, *t2 = p2;
+ *     if (t1->tv.tv_sec < t2->tv.tv_sec) {
+ *        return -1;
+ *     } else if (t1->tv.tv_sec > t2->tv.tv_sec) {
+ *        return 1;
+ *     } else {
+ *        return t1->tv.tv_usec - t2->tv_usec;
+ *     }
+ *   }
+ *
+ *   void timer_heap_insert(smartlist_t *heap, timer_t *timer) {
+ *      smartlist_pqueue_add(heap, compare, STRUCT_OFFSET(timer_t, heap_index),
+ *         timer);
+ *   }
+ *
+ *   void timer_heap_pop(smartlist_t *heap) {
+ *      return smartlist_pqueue_pop(heap, compare,
+ *         STRUCT_OFFSET(timer_t, heap_index));
+ *   }
  */
 
 /* For a 1-indexed array, we would use LEFT_CHILD[x] = 2*x and RIGHT_CHILD[x]
@@ -616,12 +648,22 @@ smartlist_uniq_strings(smartlist_t *sl)
 #define RIGHT_CHILD(i) ( 2*(i) + 2 )
 #define PARENT(i)      ( ((i)-1) / 2 )
 
+#define IDXP(p) ((int*)STRUCT_VAR_P(p, idx_field_offset))
+
+#define UPDATE_IDX(i)  do {                            \
+    void *updated = sl->list[i];                       \
+    *IDXP(updated) = i;                                \
+  } while (0)
+
+#define IDX_OF_ITEM(p) (*IDXP(p))
+
 /** Helper. <b>sl</b> may have at most one violation of the heap property:
  * the item at <b>idx</b> may be greater than one or both of its children.
  * Restore the heap property. */
 static INLINE void
 smartlist_heapify(smartlist_t *sl,
                   int (*compare)(const void *a, const void *b),
+                  int idx_field_offset,
                   int idx)
 {
   while (1) {
@@ -644,21 +686,28 @@ smartlist_heapify(smartlist_t *sl,
       void *tmp = sl->list[idx];
       sl->list[idx] = sl->list[best_idx];
       sl->list[best_idx] = tmp;
+      UPDATE_IDX(idx);
+      UPDATE_IDX(best_idx);
 
       idx = best_idx;
     }
   }
 }
 
-/** Insert <b>item</b> into the heap stored in <b>sl</b>, where order
- * is determined by <b>compare</b>. */
+/** Insert <b>item</b> into the heap stored in <b>sl</b>, where order is
+ * determined by <b>compare</b> and the offset of the item in the heap is
+ * stored in an int-typed field at position <b>idx_field_offset</b> within
+ * item.
+ */
 void
 smartlist_pqueue_add(smartlist_t *sl,
                      int (*compare)(const void *a, const void *b),
+                     int idx_field_offset,
                      void *item)
 {
   int idx;
   smartlist_add(sl,item);
+  UPDATE_IDX(sl->num_used-1);
 
   for (idx = sl->num_used - 1; idx; ) {
     int parent = PARENT(idx);
@@ -666,6 +715,8 @@ smartlist_pqueue_add(smartlist_t *sl,
       void *tmp = sl->list[parent];
       sl->list[parent] = sl->list[idx];
       sl->list[idx] = tmp;
+      UPDATE_IDX(parent);
+      UPDATE_IDX(idx);
       idx = parent;
     } else {
       return;
@@ -674,32 +725,63 @@ smartlist_pqueue_add(smartlist_t *sl,
 }
 
 /** Remove and return the top-priority item from the heap stored in <b>sl</b>,
- * where order is determined by <b>compare</b>.  <b>sl</b> must not be
- * empty. */
+ * where order is determined by <b>compare</b> and the item's position in is
+ * stored at position <b>idx_field_offset</b> within the item.  <b>sl</b> must
+ * not be empty. */
 void *
 smartlist_pqueue_pop(smartlist_t *sl,
-                     int (*compare)(const void *a, const void *b))
+                     int (*compare)(const void *a, const void *b),
+                     int idx_field_offset)
 {
   void *top;
   tor_assert(sl->num_used);
 
   top = sl->list[0];
+  *IDXP(top)=-1;
   if (--sl->num_used) {
     sl->list[0] = sl->list[sl->num_used];
-    smartlist_heapify(sl, compare, 0);
+    UPDATE_IDX(0);
+    smartlist_heapify(sl, compare, idx_field_offset, 0);
   }
   return top;
 }
 
+/** Remove the item <b>item</b> from the heap stored in <b>sl</b>,
+ * where order is determined by <b>compare</b> and the item's position in is
+ * stored at position <b>idx_field_offset</b> within the item.  <b>sl</b> must
+ * not be empty. */
+void
+smartlist_pqueue_remove(smartlist_t *sl,
+                        int (*compare)(const void *a, const void *b),
+                        int idx_field_offset,
+                        void *item)
+{
+  int idx = IDX_OF_ITEM(item);
+  tor_assert(idx >= 0);
+  tor_assert(sl->list[idx] == item);
+  --sl->num_used;
+  *IDXP(item) = -1;
+  if (idx == sl->num_used) {
+    return;
+  } else {
+    sl->list[idx] = sl->list[sl->num_used];
+    UPDATE_IDX(idx);
+    smartlist_heapify(sl, compare, idx_field_offset, idx);
+  }
+}
+
 /** Assert that the heap property is correctly maintained by the heap stored
  * in <b>sl</b>, where order is determined by <b>compare</b>. */
 void
 smartlist_pqueue_assert_ok(smartlist_t *sl,
-                           int (*compare)(const void *a, const void *b))
+                           int (*compare)(const void *a, const void *b),
+                           int idx_field_offset)
 {
   int i;
-  for (i = sl->num_used - 1; i > 0; --i) {
-    tor_assert(compare(sl->list[PARENT(i)], sl->list[i]) <= 0);
+  for (i = sl->num_used - 1; i >= 0; --i) {
+    if (i>0)
+      tor_assert(compare(sl->list[PARENT(i)], sl->list[i]) <= 0);
+    tor_assert(IDX_OF_ITEM(sl->list[i]) == i);
   }
 }
 

src/common/container.h

@@ -118,11 +118,18 @@ int smartlist_bsearch_idx(const smartlist_t *sl, const void *key,
 
 void smartlist_pqueue_add(smartlist_t *sl,
                           int (*compare)(const void *a, const void *b),
+                          int idx_field_offset,
                           void *item);
 void *smartlist_pqueue_pop(smartlist_t *sl,
-                           int (*compare)(const void *a, const void *b));
+                           int (*compare)(const void *a, const void *b),
+                           int idx_field_offset);
+void smartlist_pqueue_remove(smartlist_t *sl,
+                             int (*compare)(const void *a, const void *b),
+                             int idx_field_offset,
+                             void *item);
 void smartlist_pqueue_assert_ok(smartlist_t *sl,
-                                int (*compare)(const void *a, const void *b));
+                                int (*compare)(const void *a, const void *b),
+                                int idx_field_offset);
 
 #define SPLIT_SKIP_SPACE   0x01
 #define SPLIT_IGNORE_BLANK 0x02

src/or/dns.c

@@ -128,6 +128,8 @@ typedef struct cached_resolve_t {
   uint32_t ttl; /**< What TTL did the nameserver tell us? */
   /** Connections that want to know when we get an answer for this resolve. */
   pending_connection_t *pending_connections;
+  /** Position of this element in the heap*/
+  int minheap_idx;
 } cached_resolve_t;
 
 static void purge_expired_resolves(time_t now);
@@ -344,6 +346,7 @@ set_expiry(cached_resolve_t *resolve, time_t expires)
   resolve->expire = expires;
   smartlist_pqueue_add(cached_resolve_pqueue,
                        _compare_cached_resolves_by_expiry,
+                       STRUCT_OFFSET(cached_resolve_t, minheap_idx),
                        resolve);
 }
 
@@ -389,7 +392,8 @@ purge_expired_resolves(time_t now)
     if (resolve->expire > now)
       break;
     smartlist_pqueue_pop(cached_resolve_pqueue,
-                         _compare_cached_resolves_by_expiry);
+                         _compare_cached_resolves_by_expiry,
+                         STRUCT_OFFSET(cached_resolve_t, minheap_idx));
 
     if (resolve->state == CACHE_STATE_PENDING) {
       log_debug(LD_EXIT,
@@ -751,6 +755,7 @@ dns_resolve_impl(edge_connection_t *exitconn, int is_resolve,
   resolve = tor_malloc_zero(sizeof(cached_resolve_t));
   resolve->magic = CACHED_RESOLVE_MAGIC;
   resolve->state = CACHE_STATE_PENDING;
+  resolve->minheap_idx = -1;
   resolve->is_reverse = is_reverse;
   strlcpy(resolve->address, exitconn->_base.address, sizeof(resolve->address));
 
@@ -1734,7 +1739,8 @@ _assert_cache_ok(void)
     return;
 
   smartlist_pqueue_assert_ok(cached_resolve_pqueue,
-                             _compare_cached_resolves_by_expiry);
+                             _compare_cached_resolves_by_expiry,
+                             STRUCT_OFFSET(cached_resolve_t, minheap_idx));
 
   SMARTLIST_FOREACH(cached_resolve_pqueue, cached_resolve_t *, res,
     {

src/test/test_containers.c

@@ -515,11 +515,17 @@ test_container_digestset(void)
   smartlist_free(included);
 }
 
-/** Helper: return a tristate based on comparing two strings. */
+typedef struct pq_entry_t {
+  const char *val;
+  int idx;
+} pq_entry_t;
+
+/** Helper: return a tristate based on comparing two pq_entry_t values. */
 static int
-_compare_strings_for_pqueue(const void *s1, const void *s2)
+_compare_strings_for_pqueue(const void *p1, const void *p2)
 {
-  return strcmp((const char*)s1, (const char*)s2);
+  const pq_entry_t *e1=p1, *e2=p2;
+  return strcmp(e1->val, e2->val);
 }
 
 /** Run unit tests for heap-based priority queue functions. */
@@ -528,50 +534,90 @@ test_container_pqueue(void)
 {
   smartlist_t *sl = smartlist_create();
   int (*cmp)(const void *, const void*);
-#define OK() smartlist_pqueue_assert_ok(sl, cmp)
+  const int offset = STRUCT_OFFSET(pq_entry_t, idx);
+#define ENTRY(s) pq_entry_t s = { #s, -1 }
+  ENTRY(cows);
+  ENTRY(zebras);
+  ENTRY(fish);
+  ENTRY(frogs);
+  ENTRY(apples);
+  ENTRY(squid);
+  ENTRY(daschunds);
+  ENTRY(eggplants);
+  ENTRY(weissbier);
+  ENTRY(lobsters);
+  ENTRY(roquefort);
+  ENTRY(chinchillas);
+  ENTRY(fireflies);
+
+#define OK() smartlist_pqueue_assert_ok(sl, cmp, offset)
 
   cmp = _compare_strings_for_pqueue;
-
-  smartlist_pqueue_add(sl, cmp, (char*)"cows");
-  smartlist_pqueue_add(sl, cmp, (char*)"zebras");
-  smartlist_pqueue_add(sl, cmp, (char*)"fish");
-  smartlist_pqueue_add(sl, cmp, (char*)"frogs");
-  smartlist_pqueue_add(sl, cmp, (char*)"apples");
-  smartlist_pqueue_add(sl, cmp, (char*)"squid");
-  smartlist_pqueue_add(sl, cmp, (char*)"daschunds");
-  smartlist_pqueue_add(sl, cmp, (char*)"eggplants");
-  smartlist_pqueue_add(sl, cmp, (char*)"weissbier");
-  smartlist_pqueue_add(sl, cmp, (char*)"lobsters");
-  smartlist_pqueue_add(sl, cmp, (char*)"roquefort");
+  smartlist_pqueue_add(sl, cmp, offset, &cows);
+  smartlist_pqueue_add(sl, cmp, offset, &zebras);
+  smartlist_pqueue_add(sl, cmp, offset, &fish);
+  smartlist_pqueue_add(sl, cmp, offset, &frogs);
+  smartlist_pqueue_add(sl, cmp, offset, &apples);
+  smartlist_pqueue_add(sl, cmp, offset, &squid);
+  smartlist_pqueue_add(sl, cmp, offset, &daschunds);
+  smartlist_pqueue_add(sl, cmp, offset, &eggplants);
+  smartlist_pqueue_add(sl, cmp, offset, &weissbier);
+  smartlist_pqueue_add(sl, cmp, offset, &lobsters);
+  smartlist_pqueue_add(sl, cmp, offset, &roquefort);
 
   OK();
 
   test_eq(smartlist_len(sl), 11);
-  test_streq(smartlist_get(sl, 0), "apples");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "apples");
+  test_eq_ptr(smartlist_get(sl, 0), &apples);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &apples);
   test_eq(smartlist_len(sl), 10);
   OK();
-  test_streq(smartlist_pqueue_pop(sl, cmp), "cows");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "daschunds");
-  smartlist_pqueue_add(sl, cmp, (char*)"chinchillas");
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &cows);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &daschunds);
+  smartlist_pqueue_add(sl, cmp, offset, &chinchillas);
+  OK();
+  smartlist_pqueue_add(sl, cmp, offset, &fireflies);
+  OK();
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &chinchillas);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &eggplants);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &fireflies);
+  OK();
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &fish);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &frogs);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &lobsters);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &roquefort);
+  OK();
+  test_eq(smartlist_len(sl), 3);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &squid);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &weissbier);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &zebras);
+  test_eq(smartlist_len(sl), 0);
   OK();
-  smartlist_pqueue_add(sl, cmp, (char*)"fireflies");
+
+  /* Now test remove. */
+  smartlist_pqueue_add(sl, cmp, offset, &cows);
+  smartlist_pqueue_add(sl, cmp, offset, &fish);
+  smartlist_pqueue_add(sl, cmp, offset, &frogs);
+  smartlist_pqueue_add(sl, cmp, offset, &apples);
+  smartlist_pqueue_add(sl, cmp, offset, &squid);
+  smartlist_pqueue_add(sl, cmp, offset, &zebras);
+  test_eq(smartlist_len(sl), 6);
   OK();
-  test_streq(smartlist_pqueue_pop(sl, cmp), "chinchillas");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "eggplants");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "fireflies");
+  smartlist_pqueue_remove(sl, cmp, offset, &zebras);
+  test_eq(smartlist_len(sl), 5);
   OK();
-  test_streq(smartlist_pqueue_pop(sl, cmp), "fish");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "frogs");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "lobsters");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "roquefort");
+  smartlist_pqueue_remove(sl, cmp, offset, &cows);
+  test_eq(smartlist_len(sl), 4);
   OK();
+  smartlist_pqueue_remove(sl, cmp, offset, &apples);
   test_eq(smartlist_len(sl), 3);
-  test_streq(smartlist_pqueue_pop(sl, cmp), "squid");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "weissbier");
-  test_streq(smartlist_pqueue_pop(sl, cmp), "zebras");
+  OK();
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &fish);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &frogs);
+  test_eq_ptr(smartlist_pqueue_pop(sl, cmp, offset), &squid);
   test_eq(smartlist_len(sl), 0);
   OK();
+
 #undef OK
 
  done: