have aux search tree lookup return a range

1 files changed, 45 insertions, 15 deletions

diff --git a/fs/bcachefs/bset.c b/fs/bcachefs/bset.c
index 1c7318c6e46f..7a016f9590f4 100644
--- a/fs/bcachefs/bset.c
+++ b/fs/bcachefs/bset.c
@@ -1167,8 +1167,12 @@ void bch2_bset_delete(struct btree *b,
 
 /* Lookup */
 
+struct bkey_range {
+	struct bkey_packed *l, *r;
+};
+
 __flatten
-static struct bkey_packed *bset_search_write_set(const struct btree *b,
+static struct bkey_range bset_search_write_set(const struct btree *b,
 				struct bset_tree *t,
 				struct bpos *search,
 				const struct bkey_packed *packed_search)
@@ -1184,7 +1188,10 @@ static struct bkey_packed *bset_search_write_set(const struct btree *b,
 			r = m;
 	}
 
-	return rw_aux_to_bkey(b, t, l);
+	return (struct bkey_range) {
+		.l = rw_aux_to_bkey(b, t, l),
+		.r = r < t->size ? rw_aux_to_bkey(b, t, r) : btree_bkey_last(b, t),
+	};
 }
 
 static inline void prefetch_four_cachelines(void *p)
@@ -1222,7 +1229,7 @@ static inline bool bkey_mantissa_bits_dropped(const struct btree *b,
 }
 
 __flatten
-static struct bkey_packed *bset_search_tree(const struct btree *b,
+static struct bkey_range bset_search_tree(const struct btree *b,
 				const struct bset_tree *t,
 				const struct bpos *search,
 				const struct bkey_packed *packed_search)
@@ -1256,26 +1263,39 @@ slowpath:
 		k = tree_to_bkey(b, t, n);
 		cmp = bkey_cmp_p_or_unp(b, k, packed_search, search);
 		if (!cmp)
-			return k;
+			return (struct bkey_range) { k, k };
 
 		n = n * 2 + (cmp < 0);
 	} while (n < t->size);
 
 	inorder = __eytzinger1_to_inorder(n >> 1, t->size, t->extra);
+	k = cacheline_to_bkey(b, t, inorder, f->key_offset);
 
 	/*
 	 * n would have been the node we recursed to - the low bit tells us if
 	 * we recursed left or recursed right.
 	 */
-	if (likely(!(n & 1))) {
-		--inorder;
-		if (unlikely(!inorder))
-			return btree_bkey_first(b, t);
-
-		f = &base->f[eytzinger1_prev(n >> 1, t->size)];
+	if (n & 1) {
+		n = eytzinger1_next(n >> 1, t->size);
+
+		return (struct bkey_range) {
+			.l = k,
+			.r = (n
+			      ? cacheline_to_bkey(b, t, inorder + 1,
+						  base->f[n].key_offset)
+			      : btree_bkey_last(b, t)),
+		};
+	} else {
+		n = eytzinger1_prev(n >> 1, t->size);
+
+		return (struct bkey_range) {
+			.l = (n
+			      ? cacheline_to_bkey(b, t, inorder - 1,
+						  base->f[n].key_offset)
+			      : btree_bkey_first(b, t)),
+			.r = k,
+		};
 	}
-
-	return cacheline_to_bkey(b, t, inorder, f->key_offset);
 }
 
 static __always_inline __flatten
@@ -1284,6 +1304,7 @@ struct bkey_packed *__bch2_bset_search(struct btree *b,
 				struct bpos *search,
 				const struct bkey_packed *lossy_packed_search)
 {
+	struct bkey_range r;
 
 	/*
 	 * First, we search for a cacheline, then lastly we do a linear search
@@ -1302,9 +1323,12 @@ struct bkey_packed *__bch2_bset_search(struct btree *b,
 
 	switch (bset_aux_tree_type(t)) {
 	case BSET_NO_AUX_TREE:
-		return btree_bkey_first(b, t);
+		r.l = btree_bkey_first(b, t);
+		r.r = btree_bkey_last(b, t);
+		break;
 	case BSET_RW_AUX_TREE:
-		return bset_search_write_set(b, t, search, lossy_packed_search);
+		r = bset_search_write_set(b, t, search, lossy_packed_search);
+		break;
 	case BSET_RO_AUX_TREE:
 		/*
 		 * Each node in the auxiliary search tree covers a certain range
@@ -1316,10 +1340,16 @@ struct bkey_packed *__bch2_bset_search(struct btree *b,
 		if (bkey_cmp(*search, t->max_key) > 0)
 			return btree_bkey_last(b, t);
 
-		return bset_search_tree(b, t, search, lossy_packed_search);
+		r = bset_search_tree(b, t, search, lossy_packed_search);
+		break;
 	default:
 		unreachable();
 	}
+
+	WARN_ON(r.r != btree_bkey_last(b, t) &&
+		bkey_cmp_left_packed(b, r.r, search) < 0);
+
+	return r.l;
 }
 
 static __always_inline __flatten