bcachefs: Consolidate bch2_btree_iter_peek() and peek_with_updates()

Ideally we'll be getting rid of peek_with_updates(), but the callers
will need to be checked.

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

diff --git a/fs/bcachefs/btree_iter.c b/fs/bcachefs/btree_iter.c
index d036ace70552e8eeef1e590c0acb3a2f141cfbdf..43885f9078375a87dd6551ee276e32440213c6ab 100644 (file)
--- a/fs/bcachefs/btree_iter.c
+++ b/fs/bcachefs/btree_iter.c
@@ -1531,12 +1531,28 @@ static inline bool btree_iter_set_pos_to_prev_leaf(struct btree_iter *iter)
        return ret;
 }
 
-/**
- * bch2_btree_iter_peek: returns first key greater than or equal to iterator's
- * current position
- */
-struct bkey_s_c bch2_btree_iter_peek(struct btree_iter *iter)
+static struct bkey_i *btree_trans_peek_updates(struct btree_trans *trans,
+                                              enum btree_id btree_id, struct bpos pos)
 {
+       struct btree_insert_entry *i;
+
+       trans_for_each_update2(trans, i)
+               if ((cmp_int(btree_id,  i->iter->btree_id) ?:
+                    bkey_cmp(pos,      i->k->k.p)) <= 0) {
+                       if (btree_id == i->iter->btree_id)
+                               return i->k;
+                       break;
+               }
+
+       return NULL;
+}
+
+static inline struct bkey_s_c __btree_iter_peek(struct btree_iter *iter, bool with_updates)
+{
+       struct bpos search_key = btree_iter_search_key(iter);
+       struct bkey_i *next_update = with_updates
+               ? btree_trans_peek_updates(iter->trans, iter->btree_id, search_key)
+               : NULL;
        struct bkey_s_c k;
        int ret;
 
@@ -1544,7 +1560,7 @@ struct bkey_s_c bch2_btree_iter_peek(struct btree_iter *iter)
        bch2_btree_iter_verify(iter);
        bch2_btree_iter_verify_entry_exit(iter);
 
-       btree_iter_set_search_pos(iter, btree_iter_search_key(iter));
+       btree_iter_set_search_pos(iter, search_key);
 
        while (1) {
                ret = bch2_btree_iter_traverse(iter);
@@ -1552,16 +1568,28 @@ struct bkey_s_c bch2_btree_iter_peek(struct btree_iter *iter)
                        return bkey_s_c_err(ret);
 
                k = btree_iter_level_peek(iter, &iter->l[0]);
-               if (likely(k.k))
+
+               if (next_update &&
+                   bkey_cmp(next_update->k.p, iter->real_pos) <= 0)
+                       k = bkey_i_to_s_c(next_update);
+
+               if (likely(k.k)) {
+                       if (bkey_deleted(k.k)) {
+                               btree_iter_set_search_pos(iter,
+                                               bkey_successor(k.k->p));
+                               continue;
+                       }
+
                        break;
+               }
 
                if (!btree_iter_set_pos_to_next_leaf(iter))
                        return bkey_s_c_null;
        }
 
        /*
-        * iter->pos should always be equal to the key we just
-        * returned - except extents can straddle iter->pos:
+        * iter->pos should be mononotically increasing, and always be equal to
+        * the key we just returned - except extents can straddle iter->pos:
         */
        if (bkey_cmp(bkey_start_pos(k.k), iter->pos) > 0)
                iter->pos = bkey_start_pos(k.k);
@@ -1571,6 +1599,15 @@ struct bkey_s_c bch2_btree_iter_peek(struct btree_iter *iter)
        return k;
 }
 
+/**
+ * bch2_btree_iter_peek: returns first key greater than or equal to iterator's
+ * current position
+ */
+struct bkey_s_c bch2_btree_iter_peek(struct btree_iter *iter)
+{
+       return __btree_iter_peek(iter, false);
+}
+
 /**
  * bch2_btree_iter_next: returns first key greater than iterator's current
  * position
@@ -1583,79 +1620,9 @@ struct bkey_s_c bch2_btree_iter_next(struct btree_iter *iter)
        return bch2_btree_iter_peek(iter);
 }
 
-static struct bkey_s_c __btree_trans_updates_peek(struct btree_iter *iter)
-{
-       struct bpos pos = btree_iter_search_key(iter);
-       struct btree_trans *trans = iter->trans;
-       struct btree_insert_entry *i;
-
-       trans_for_each_update2(trans, i)
-               if ((cmp_int(iter->btree_id,    i->iter->btree_id) ?:
-                    bkey_cmp(pos,              i->k->k.p)) <= 0)
-                       break;
-
-       return i < trans->updates2 + trans->nr_updates2 &&
-               iter->btree_id == i->iter->btree_id
-               ? bkey_i_to_s_c(i->k)
-               : bkey_s_c_null;
-}
-
-static struct bkey_s_c __bch2_btree_iter_peek_with_updates(struct btree_iter *iter)
-{
-       struct btree_iter_level *l = &iter->l[0];
-       struct bkey_s_c k = btree_iter_level_peek(iter, l);
-       struct bkey_s_c u = __btree_trans_updates_peek(iter);
-
-       if (k.k && (!u.k || bkey_cmp(k.k->p, u.k->p) < 0))
-               return k;
-       if (u.k && bkey_cmp(u.k->p, l->b->key.k.p) <= 0) {
-               iter->k = *u.k;
-               return u;
-       }
-       return bkey_s_c_null;
-}
-
 struct bkey_s_c bch2_btree_iter_peek_with_updates(struct btree_iter *iter)
 {
-       struct bkey_s_c k;
-       int ret;
-
-       EBUG_ON(btree_iter_type(iter) != BTREE_ITER_KEYS);
-       bch2_btree_iter_verify(iter);
-       bch2_btree_iter_verify_entry_exit(iter);
-
-       btree_iter_set_search_pos(iter, btree_iter_search_key(iter));
-
-       while (1) {
-               ret = bch2_btree_iter_traverse(iter);
-               if (unlikely(ret))
-                       return bkey_s_c_err(ret);
-
-               k = __bch2_btree_iter_peek_with_updates(iter);
-
-               if (k.k && bkey_deleted(k.k)) {
-                       if (!bch2_btree_iter_advance(iter))
-                               return bkey_s_c_null;
-                       continue;
-               }
-
-               if (likely(k.k))
-                       break;
-
-               if (!btree_iter_set_pos_to_next_leaf(iter))
-                       return bkey_s_c_null;
-       }
-
-       /*
-        * iter->pos should be mononotically increasing, and always be equal to
-        * the key we just returned - except extents can straddle iter->pos:
-        */
-       if (bkey_cmp(bkey_start_pos(k.k), iter->pos) > 0)
-               iter->pos = bkey_start_pos(k.k);
-
-       bch2_btree_iter_verify_entry_exit(iter);
-       bch2_btree_iter_verify(iter);
-       return k;
+       return __btree_iter_peek(iter, true);
 }
 
 struct bkey_s_c bch2_btree_iter_next_with_updates(struct btree_iter *iter)