lib: add DECLARE_HEAP datastructure

This is an 8-ary heap (cacheline optimized.)  It works as a semi-sorted
kind of middle ground between unsorted and sorted datastructures;  pop()
always returns the lowest item but ordering is only loosely enforced.

Signed-off-by: David Lamparter <equinox@diac24.net>

2 files changed, 260 insertions, 0 deletions

diff --git a/lib/typesafe.c b/lib/typesafe.c
index 4b07c482e..47a6d0af4 100644
--- a/lib/typesafe.c
+++ b/lib/typesafe.c
@@ -22,6 +22,7 @@
 
 DEFINE_MTYPE_STATIC(LIB, TYPEDHASH_BUCKET, "Typed-hash bucket")
 DEFINE_MTYPE_STATIC(LIB, SKIPLIST_OFLOW, "Skiplist overflow")
+DEFINE_MTYPE_STATIC(LIB, HEAP_ARRAY, "Typed-heap array")
 
 #if 0
 static void hash_consistency_check(struct thash_head *head)
@@ -407,3 +408,148 @@ struct sskip_item *typesafe_skiplist_pop(struct sskip_head *head)
 
 	return item;
 }
+
+/* heap */
+
+#if 0
+static void heap_consistency_check(struct heap_head *head,
+				   int (*cmpfn)(const struct heap_item *a,
+						const struct heap_item *b),
+				   uint32_t pos)
+{
+	uint32_t rghtpos = pos + 1;
+	uint32_t downpos = HEAP_NARY * (pos + 1);
+
+	if (pos + 1 > ~0U / HEAP_NARY)
+		downpos = ~0U;
+
+	if ((pos & (HEAP_NARY - 1)) != HEAP_NARY - 1 && rghtpos < head->count) {
+		assert(cmpfn(head->array[rghtpos], head->array[pos]) >= 0);
+		heap_consistency_check(head, cmpfn, rghtpos);
+	}
+	if (downpos < head->count) {
+		assert(cmpfn(head->array[downpos], head->array[pos]) >= 0);
+		heap_consistency_check(head, cmpfn, downpos);
+	}
+}
+#else
+#define heap_consistency_check(head, cmpfn, pos)
+#endif
+
+void typesafe_heap_resize(struct heap_head *head, bool grow)
+{
+	uint32_t newsize;
+
+	if (grow) {
+		newsize = head->arraysz;
+		if (newsize <= 36)
+			newsize = 72;
+		else if (newsize < 262144)
+			newsize += newsize / 2;
+		else if (newsize < 0xaaaa0000)
+			newsize += newsize / 3;
+		else
+			assert(!newsize);
+	} else if (head->count > 0) {
+		newsize = head->count;
+	} else {
+		XFREE(MTYPE_HEAP_ARRAY, head->array);
+		head->arraysz = 0;
+		return;
+	}
+
+	newsize += HEAP_NARY - 1;
+	newsize &= ~(HEAP_NARY - 1);
+	if (newsize == head->arraysz)
+		return;
+	
+	head->array = XREALLOC(MTYPE_HEAP_ARRAY, head->array,
+			       newsize * sizeof(struct heap_item *));
+	head->arraysz = newsize;
+}
+
+void typesafe_heap_pushdown(struct heap_head *head, uint32_t pos,
+		struct heap_item *item,
+		int (*cmpfn)(const struct heap_item *a,
+			     const struct heap_item *b))
+{
+	uint32_t rghtpos, downpos, moveto;
+
+	while (1) {
+		/* rghtpos: neighbor to the "right", inside block of NARY.
+		 *          may be invalid if last in block, check nary_last()
+		 * downpos: first neighbor in the "downwards" block further
+		 *          away from the root
+		 */
+		rghtpos = pos + 1;
+
+		/* make sure we can use the full 4G items */
+		downpos = HEAP_NARY * (pos + 1);
+		if (pos + 1 > ~0U / HEAP_NARY)
+			/* multiplication overflowed.  ~0U is guaranteed
+			 * to be an invalid index; size limit is enforced in
+			 * resize()
+			 */
+			downpos = ~0U;
+
+		/* only used on break */
+		moveto = pos;
+
+#define nary_last(x) (((x) & (HEAP_NARY - 1)) == HEAP_NARY - 1)
+		if (downpos >= head->count
+		    || cmpfn(head->array[downpos], item) >= 0) {
+			/* not moving down; either at end or down is >= item */
+			if (nary_last(pos) || rghtpos >= head->count
+			    || cmpfn(head->array[rghtpos], item) >= 0)
+				/* not moving right either - got our spot */
+				break;
+
+			moveto = rghtpos;
+
+		/* else: downpos is valid and < item.  choose between down
+		 * or right (if the latter is an option) */
+		} else if (nary_last(pos) || cmpfn(head->array[rghtpos],
+						   head->array[downpos]) >= 0)
+			moveto = downpos;
+		else
+			moveto = rghtpos;
+#undef nary_last
+
+		head->array[pos] = head->array[moveto];
+		head->array[pos]->index = pos;
+		pos = moveto;
+	}
+
+	head->array[moveto] = item;
+	item->index = moveto;
+
+	heap_consistency_check(head, cmpfn, 0);
+}
+
+void typesafe_heap_pullup(struct heap_head *head, uint32_t pos,
+		struct heap_item *item,
+		int (*cmpfn)(const struct heap_item *a,
+			     const struct heap_item *b))
+{
+	uint32_t moveto;
+
+	while (pos != 0) {
+		if ((pos & (HEAP_NARY - 1)) == 0)
+			moveto = pos / HEAP_NARY - 1;
+		else
+			moveto = pos - 1;
+
+		if (cmpfn(head->array[moveto], item) <= 0)
+			break;
+
+		head->array[pos] = head->array[moveto];
+		head->array[pos]->index = pos;
+
+		pos = moveto;
+	}
+
+	head->array[pos] = item;
+	item->index = pos;
+
+	heap_consistency_check(head, cmpfn, 0);
+}
diff --git a/lib/typesafe.h b/lib/typesafe.h
index 28d18e09f..7783d0707 100644
--- a/lib/typesafe.h
+++ b/lib/typesafe.h
@@ -252,6 +252,120 @@ macro_pure size_t prefix ## _count(struct prefix##_head *h)                    \
 }                                                                              \
 /* ... */
 
+/* note: heap currently caps out at 4G items */
+
+#define HEAP_NARY 8U
+typedef uint32_t heap_index_i;
+
+struct heap_item {
+	uint32_t index;
+};
+
+struct heap_head {
+	struct heap_item **array;
+	uint32_t arraysz, count;
+};
+
+#define HEAP_RESIZE_TRESH_UP(h) \
+	(h->hh.count + 1 >= h->hh.arraysz)
+#define HEAP_RESIZE_TRESH_DN(h) \
+	(h->hh.count == 0 || \
+	 h->hh.arraysz - h->hh.count > (h->hh.count + 1024) / 2)
+
+#define PREDECL_HEAP(prefix)                                                   \
+struct prefix ## _head { struct heap_head hh; };                               \
+struct prefix ## _item { struct heap_item hi; };
+
+#define INIT_HEAP(var)		{ }
+
+#define DECLARE_HEAP(prefix, type, field, cmpfn)                               \
+                                                                               \
+macro_inline void prefix ## _init(struct prefix##_head *h)                     \
+{                                                                              \
+	memset(h, 0, sizeof(*h));                                              \
+}                                                                              \
+macro_inline void prefix ## _fini(struct prefix##_head *h)                     \
+{                                                                              \
+	assert(h->hh.count == 0);                                              \
+	memset(h, 0, sizeof(*h));                                              \
+}                                                                              \
+macro_inline int prefix ## __cmp(const struct heap_item *a,                    \
+		const struct heap_item *b)                                     \
+{                                                                              \
+	return cmpfn(container_of(a, type, field.hi),                          \
+			container_of(b, type, field.hi));                      \
+}                                                                              \
+macro_inline type *prefix ## _add(struct prefix##_head *h, type *item)         \
+{                                                                              \
+	if (HEAP_RESIZE_TRESH_UP(h))                                           \
+		typesafe_heap_resize(&h->hh, true);                            \
+	typesafe_heap_pullup(&h->hh, h->hh.count, &item->field.hi,             \
+			     prefix ## __cmp);                                 \
+	h->hh.count++;                                                         \
+	return NULL;                                                           \
+}                                                                              \
+macro_inline void prefix ## _del(struct prefix##_head *h, type *item)          \
+{                                                                              \
+	struct heap_item *other;                                               \
+	uint32_t index = item->field.hi.index;                                 \
+	assert(h->hh.array[index] == &item->field.hi);                         \
+	h->hh.count--;                                                         \
+	other = h->hh.array[h->hh.count];                                      \
+	if (cmpfn(container_of(other, type, field.hi), item) < 0)              \
+		typesafe_heap_pullup(&h->hh, index, other, prefix ## __cmp);   \
+	else                                                                   \
+		typesafe_heap_pushdown(&h->hh, index, other, prefix ## __cmp); \
+	if (HEAP_RESIZE_TRESH_DN(h))                                           \
+		typesafe_heap_resize(&h->hh, false);                           \
+}                                                                              \
+macro_inline type *prefix ## _pop(struct prefix##_head *h)                     \
+{                                                                              \
+	struct heap_item *hitem, *other;                                       \
+	if (h->hh.count == 0)                                                  \
+		return NULL;                                                   \
+	hitem = h->hh.array[0];                                                \
+	h->hh.count--;                                                         \
+	other = h->hh.array[h->hh.count];                                      \
+	typesafe_heap_pushdown(&h->hh, 0, other, prefix ## __cmp);             \
+	if (HEAP_RESIZE_TRESH_DN(h))                                           \
+		typesafe_heap_resize(&h->hh, false);                           \
+	return container_of(hitem, type, field.hi);                            \
+}                                                                              \
+macro_pure type *prefix ## _first(struct prefix##_head *h)                     \
+{                                                                              \
+	if (h->hh.count == 0)                                                  \
+		return NULL;                                                   \
+	return container_of(h->hh.array[0], type, field.hi);                   \
+}                                                                              \
+macro_pure type *prefix ## _next(struct prefix##_head *h, type *item)          \
+{                                                                              \
+	uint32_t idx = item->field.hi.index + 1;                               \
+	if (idx >= h->hh.count)                                                \
+		return NULL;                                                   \
+	return container_of(h->hh.array[idx], type, field.hi);                 \
+}                                                                              \
+macro_pure type *prefix ## _next_safe(struct prefix##_head *h, type *item)     \
+{                                                                              \
+	if (!item)                                                             \
+		return NULL;                                                   \
+	return prefix ## _next(h, item);                                       \
+}                                                                              \
+macro_pure size_t prefix ## _count(struct prefix##_head *h)                    \
+{                                                                              \
+	return h->hh.count;                                                    \
+}                                                                              \
+/* ... */
+
+extern void typesafe_heap_resize(struct heap_head *head, bool grow);
+extern void typesafe_heap_pushdown(struct heap_head *head, uint32_t index,
+		struct heap_item *item,
+		int (*cmpfn)(const struct heap_item *a,
+			     const struct heap_item *b));
+extern void typesafe_heap_pullup(struct heap_head *head, uint32_t index,
+		struct heap_item *item,
+		int (*cmpfn)(const struct heap_item *a,
+			     const struct heap_item *b));
+
 /* single-linked list, sorted.
  * can be used as priority queue with add / pop
  */