rust: Implement BtreeNodeIter

This implements BtreeNodeIter, and adds it to cmd_list.rs - the next
step in having a full replacement for cmd_list.c

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

2 files changed, 101 insertions, 2 deletions

diff --git a/rust-src/bch_bindgen/src/btree.rs b/rust-src/bch_bindgen/src/btree.rs
index c6d2fec2..7e748f9f 100644
--- a/rust-src/bch_bindgen/src/btree.rs
+++ b/rust-src/bch_bindgen/src/btree.rs
@@ -63,7 +63,7 @@ impl<'t> BtreeIter<'t> {
 
             c::bch2_trans_iter_init_outlined(
                 ptr::addr_of!(trans.raw).cast_mut(),
-                &mut (*iter.as_mut_ptr()),
+                iter.as_mut_ptr(),
                 btree as u32,
                 pos,
                 flags.bits as u32);
@@ -105,3 +105,57 @@ impl<'t> Drop for BtreeIter<'t> {
         unsafe { c::bch2_trans_iter_exit(self.raw.trans, &mut self.raw) }
     }             
 }
+
+pub struct BtreeNodeIter<'t> {
+    raw:    c::btree_iter,
+    trans:  PhantomData<&'t BtreeTrans<'t>>,
+}
+
+impl<'t> BtreeNodeIter<'t> {
+    pub fn new(trans: &'t BtreeTrans<'t>,
+        btree:      c::btree_id,
+        pos:        c::bpos,
+        locks_want: u32,
+        depth:      u32,
+        flags: BtreeIterFlags) -> BtreeNodeIter {
+        unsafe {
+            let mut iter: MaybeUninit<c::btree_iter> = MaybeUninit::uninit();
+            c::bch2_trans_node_iter_init(
+                ptr::addr_of!(trans.raw).cast_mut(),
+                iter.as_mut_ptr(),
+                btree,
+                pos,
+                locks_want,
+                depth,
+                flags.bits as u32);
+
+            BtreeNodeIter { raw: iter.assume_init(), trans: PhantomData }
+        }
+    }
+
+    pub fn peek<'i>(&'i mut self) -> Result<Option<&'i c::btree>, bch_errcode> {
+        unsafe {
+            let b = c::bch2_btree_iter_peek_node(&mut self.raw);
+            errptr_to_result_c(b).map(|b| if !b.is_null() { Some(&*b) } else { None })
+        }
+    }
+
+    pub fn advance<'i>(&'i mut self) {
+        unsafe {
+            c::bch2_btree_iter_next_node(&mut self.raw);
+        }
+    }
+
+    pub fn next<'i>(&'i mut self) -> Result<Option<&'i c::btree>, bch_errcode> {
+        unsafe {
+            let b = c::bch2_btree_iter_next_node(&mut self.raw);
+            errptr_to_result_c(b).map(|b| if !b.is_null() { Some(&*b) } else { None })
+        }
+    }
+}
+
+impl<'t> Drop for BtreeNodeIter<'t> {
+    fn drop(&mut self) {
+        unsafe { c::bch2_trans_iter_exit(self.raw.trans, &mut self.raw) }
+    }             
+}
diff --git a/rust-src/src/cmd_list.rs b/rust-src/src/cmd_list.rs
index c89ba4f0..fb7cb1b2 100644
--- a/rust-src/src/cmd_list.rs
+++ b/rust-src/src/cmd_list.rs
@@ -5,6 +5,7 @@ use bch_bindgen::opt_set;
 use bch_bindgen::fs::Fs;
 use bch_bindgen::btree::BtreeTrans;
 use bch_bindgen::btree::BtreeIter;
+use bch_bindgen::btree::BtreeNodeIter;
 use bch_bindgen::btree::BtreeIterFlags;
 use clap::Parser;
 use colored::Colorize;
@@ -31,24 +32,68 @@ fn list_keys(fs: &Fs, opt: Cli) -> anyhow::Result<()> {
 
 fn list_btree_formats(fs: &Fs, opt: Cli) -> anyhow::Result<()> {
     let trans = BtreeTrans::new(fs);
+    let mut iter = BtreeNodeIter::new(&trans, opt.btree, opt.start,
+        0, opt.level,
+        BtreeIterFlags::PREFETCH);
+
+    while let Some(b) = iter.peek()? {
+        if b.key.k.p > opt.end {
+            break;
+        }
+
+        iter.advance();
+    }
 
     Ok(())
 }
 
 fn list_btree_nodes(fs: &Fs, opt: Cli) -> anyhow::Result<()> {
     let trans = BtreeTrans::new(fs);
+    let mut iter = BtreeNodeIter::new(&trans, opt.btree, opt.start,
+        0, opt.level,
+        BtreeIterFlags::PREFETCH);
+
+    while let Some(b) = iter.peek()? {
+        if b.key.k.p > opt.end {
+            break;
+        }
+
+        iter.advance();
+    }
 
     Ok(())
 }
 
 fn list_nodes_ondisk(fs: &Fs, opt: Cli) -> anyhow::Result<()> {
     let trans = BtreeTrans::new(fs);
+    let mut iter = BtreeNodeIter::new(&trans, opt.btree, opt.start,
+        0, opt.level,
+        BtreeIterFlags::PREFETCH);
+
+    while let Some(b) = iter.peek()? {
+        if b.key.k.p > opt.end {
+            break;
+        }
+
+        iter.advance();
+    }
 
     Ok(())
 }
 
 fn list_nodes_keys(fs: &Fs, opt: Cli) -> anyhow::Result<()> {
     let trans = BtreeTrans::new(fs);
+    let mut iter = BtreeNodeIter::new(&trans, opt.btree, opt.start,
+        0, opt.level,
+        BtreeIterFlags::PREFETCH);
+
+    while let Some(b) = iter.peek()? {
+        if b.key.k.p > opt.end {
+            break;
+        }
+
+        iter.advance();
+    }
 
     Ok(())
 }
@@ -70,7 +115,7 @@ struct Cli {
 
     /// Btree depth to descend to (0 == leaves)
     #[arg(short, long, default_value_t=0)]
-    level:      u8,
+    level:      u32,
 
     /// Start position to list from
     #[arg(short, long, default_value="POS_MIN")]