bcachefs: add functionality for heaps to update backpointers

Signed-off-by: Kent Overstreet <kent.overstreet@linux.dev>

1 files changed, 37 insertions, 22 deletions

diff --git a/fs/bcachefs/util.h b/fs/bcachefs/util.h
index 44e2c96b6509..9caf2487ee63 100644
--- a/fs/bcachefs/util.h
+++ b/fs/bcachefs/util.h
@@ -127,7 +127,19 @@ do {									\
 	(heap)->data = NULL;						\
 } while (0)
 
-#define heap_swap(h, i, j)	swap((h)->data[i], (h)->data[j])
+#define heap_set_backpointer(h, i, _fn)					\
+do {									\
+	void (*fn)(typeof(h), size_t) = _fn;				\
+	if (fn)								\
+		fn(h, i);						\
+} while (0)
+
+#define heap_swap(h, i, j, set_backpointer)				\
+do {									\
+	swap((h)->data[i], (h)->data[j]);				\
+	heap_set_backpointer(h, i, set_backpointer);			\
+	heap_set_backpointer(h, j, set_backpointer);			\
+} while (0)
 
 #define heap_peek(h)							\
 ({									\
@@ -137,7 +149,7 @@ do {									\
 
 #define heap_full(h)	((h)->used == (h)->size)
 
-#define heap_sift_down(h, i, cmp)					\
+#define heap_sift_down(h, i, cmp, set_backpointer)			\
 do {									\
 	size_t _c, _j = i;						\
 									\
@@ -149,72 +161,75 @@ do {									\
 									\
 		if (cmp(h, (h)->data[_c], (h)->data[_j]) >= 0)		\
 			break;						\
-		heap_swap(h, _c, _j);					\
+		heap_swap(h, _c, _j, set_backpointer);			\
 	}								\
 } while (0)
 
-#define heap_sift_up(h, i, cmp)						\
+#define heap_sift_up(h, i, cmp, set_backpointer)			\
 do {									\
 	while (i) {							\
 		size_t p = (i - 1) / 2;					\
 		if (cmp(h, (h)->data[i], (h)->data[p]) >= 0)		\
 			break;						\
-		heap_swap(h, i, p);					\
+		heap_swap(h, i, p, set_backpointer);			\
 		i = p;							\
 	}								\
 } while (0)
 
-#define __heap_add(h, d, cmp)						\
-do {									\
+#define __heap_add(h, d, cmp, set_backpointer)				\
+({									\
 	size_t _i = (h)->used++;					\
 	(h)->data[_i] = d;						\
+	heap_set_backpointer(h, _i, set_backpointer);			\
 									\
-	heap_sift_up(h, _i, cmp);					\
-} while (0)
+	heap_sift_up(h, _i, cmp, set_backpointer);			\
+	_i;								\
+})
 
-#define heap_add(h, d, cmp)						\
+#define heap_add(h, d, cmp, set_backpointer)				\
 ({									\
 	bool _r = !heap_full(h);					\
 	if (_r)								\
-		__heap_add(h, d, cmp);					\
+		__heap_add(h, d, cmp, set_backpointer);			\
 	_r;								\
 })
 
-#define heap_add_or_replace(h, new, cmp)				\
+#define heap_add_or_replace(h, new, cmp, set_backpointer)		\
 do {									\
-	if (!heap_add(h, new, cmp) &&					\
+	if (!heap_add(h, new, cmp, set_backpointer) &&			\
 	    cmp(h, new, heap_peek(h)) >= 0) {				\
 		(h)->data[0] = new;					\
-		heap_sift_down(h, 0, cmp);				\
+		heap_set_backpointer(h, 0, set_backpointer);		\
+		heap_sift_down(h, 0, cmp, set_backpointer);		\
 	}								\
 } while (0)
 
-#define heap_del(h, i, cmp)						\
+#define heap_del(h, i, cmp, set_backpointer)				\
 do {									\
 	size_t _i = (i);						\
 									\
 	BUG_ON(_i >= (h)->used);					\
 	(h)->used--;							\
-	heap_swap(h, _i, (h)->used);					\
-	heap_sift_up(h, _i, cmp);					\
-	heap_sift_down(h, _i, cmp);					\
+	heap_swap(h, _i, (h)->used, set_backpointer);			\
+	heap_sift_up(h, _i, cmp, set_backpointer);			\
+	heap_sift_down(h, _i, cmp, set_backpointer);			\
 } while (0)
 
-#define heap_pop(h, d, cmp)						\
+#define heap_pop(h, d, cmp, set_backpointer)				\
 ({									\
 	bool _r = (h)->used;						\
 	if (_r) {							\
 		(d) = (h)->data[0];					\
-		heap_del(h, 0, cmp);					\
+		heap_del(h, 0, cmp, set_backpointer);			\
 	}								\
 	_r;								\
 })
 
-#define heap_resort(heap, cmp)						\
+#define heap_resort(heap, cmp, set_backpointer)				\
 do {									\
 	ssize_t _i;							\
 	for (_i = (ssize_t) (heap)->used / 2 -  1; _i >= 0; --_i)	\
-		heap_sift_down(heap, _i, cmp);				\
+		heap_sift_down(heap, _i, cmp, set_backpointer);		\
 } while (0)
 
 #define ANYSINT_MAX(t)							\