Implement bitfield iterator (#68)

* Add one bitfield test from dat/hypercore

* Rnd::choose is deprecated

* Fix warnings and clippy lints

* Implement bitfield iterator

* Modify data_iterate mask to better match usage

* Cleanup

* Rename Bitfield::length() to len() and add is_empty()

Author: khernyo

Cargo.toml

@@ -21,7 +21,7 @@ blake2-rfc = "0.2.18"
 byteorder = "1.2.6"
 ed25519-dalek = "0.8.1"
 failure = "0.1.2"
-flat-tree = "4.0.1"
+flat-tree = "4.1"
 lazy_static = "1.1.0"
 memory-pager = "0.8.0"
 merkle-tree-stream = "0.8.0"

src/bitfield/iterator.rs

@@ -4,23 +4,23 @@ use super::Bitfield;
 
 /// Iterate over a bitfield.
 #[derive(Debug)]
-pub struct<b>Iterator {
+pub struct Iterator<'a> {
   start: usize,
   end: usize,
   index_end: usize,
-  pos: usize,
-  byte: usize,
-  bitfield: Bitfield,
+  pos: Option<usize>,
+  byte: u8,
+  bitfield: &'a mut Bitfield,
 }
 
-impl Iterator {
+impl<'a> Iterator<'a> {
   /// Create a new instance.
-  pub fn new(bitfield: &mut Bitfield) -> Self {
+  pub fn new(bitfield: &'a mut Bitfield) -> Self {
     Self {
       start: 0,
       end: 0,
       index_end: 0,
-      pos: 0,
+      pos: Some(0),
       byte: 0,
       bitfield,
     }
@@ -30,33 +30,129 @@ impl Iterator {
   pub fn range(&mut self, start: usize, end: usize) {
     self.start = start;
     self.end = end;
-    self.index_end = 2 * (end as f32 / 32).ceil();
+    self.index_end = 2 * ((end + 31) / 32);
 
-    if self.end > self.bitfield.len() {
+    if self.end > self.bitfield.length {
       self.bitfield.expand(self.end);
     }
   }
 
-  pub fn seek(&mut self, offset: usize)  {
-  offset += self.start;
-  if offset < self.start {
-    offset = self.start
+  /// Seek to `offset`
+  pub fn seek(&mut self, mut offset: usize) -> &mut Self {
+    offset += self.start;
+    // FIXME This is fishy. Offset and start is unsigned, so `offset < self.start` can only
+    //  be true when the previous addition overflows. The overflow would cause a panic, so,
+    //  either the addition should be a wrapping_add, or rather, the original offset should
+    //  be checked to ensure it is less than `self.end - self.start`.
+    if offset < self.start {
+      offset = self.start;
     }
 
-  if offset >= self.end {
-    self.pos = -1
-    return;
+    if offset >= self.end {
+      self.pos = None;
+      return self;
+    }
+
+    let o = offset % 8;
+
+    let pos = offset / 8;
+    self.pos = Some(pos);
+
+    self.byte =
+      self.bitfield.data.get_byte(pos) | self.bitfield.masks.data_iterate[o];
+
+    self
+  }
+
+  pub fn next(&mut self) -> Option<usize> {
+    let mut pos = self.pos?;
+
+    let mut free = self.bitfield.masks.next_data_0_bit[self.byte as usize];
+
+    while free == -1 {
+      pos += 1;
+      self.byte = self.bitfield.data.get_byte(pos);
+      free = self.bitfield.masks.next_data_0_bit[self.byte as usize];
+
+      if free == -1 {
+        pos = self.skip_ahead(pos)?;
+
+        self.byte = self.bitfield.data.get_byte(pos);
+        free = self.bitfield.masks.next_data_0_bit[self.byte as usize];
+      }
+    }
+    self.pos = Some(pos);
+
+    self.byte |= self.bitfield.masks.data_iterate[free as usize + 1];
+
+    let n = 8 * pos + free as usize;
+    if n < self.end {
+      Some(n)
+    } else {
+      None
+    }
   }
 
-  let o = offset & 7;
+  pub fn skip_ahead(&mut self, start: usize) -> Option<usize> {
+    let bitfield_index = &self.bitfield.index;
+    let tree_end = self.index_end;
+    let iter = &mut self.bitfield.iterator;
+    let o = start & 3;
+
+    iter.seek(2 * (start / 4));
 
-  self.pos = (offset - o) / 8;
-  let left = self.bitfield.data.get_byte(self.pos);
-  let right = if o == 0 {
-    0
-  } else {
-    self.bitfield.masks.data_iterate_mask[o - 1]
-  };
+    let mut tree_byte = bitfield_index.get_byte(iter.index())
+      | self.bitfield.masks.index_iterate[o];
+
+    while self.bitfield.masks.next_index_0_bit[tree_byte as usize] == -1 {
+      if iter.is_left() {
+        iter.next();
+      } else {
+        iter.next();
+        iter.parent();
+      }
+
+      if right_span(iter) >= tree_end {
+        while right_span(iter) >= tree_end && is_parent(iter) {
+          iter.left_child();
+        }
+        if right_span(iter) >= tree_end {
+          return None;
+        }
+      }
+
+      tree_byte = bitfield_index.get_byte(iter.index());
+    }
+
+    while iter.factor() > 2 {
+      if self.bitfield.masks.next_index_0_bit[tree_byte as usize] < 2 {
+        iter.left_child();
+      } else {
+        iter.right_child();
+      }
+
+      tree_byte = bitfield_index.get_byte(iter.index());
+    }
+
+    let mut free = self.bitfield.masks.next_index_0_bit[tree_byte as usize];
+    if free == -1 {
+      free = 4;
+    }
+
+    let next = iter.index() * 2 + free as usize;
+
+    if next <= start {
+      Some(start + 1)
+    } else {
+      Some(next)
+    }
+  }
+}
+
+fn right_span(iter: &flat_tree::Iterator) -> usize {
+  iter.index() + iter.factor() / 2 - 1
+}
 
-  self.byte = left | right;
+fn is_parent(iter: &flat_tree::Iterator) -> bool {
+  iter.index() & 1 == 1
 }

src/bitfield/masks.rs

@@ -34,6 +34,7 @@ impl Masks {
     ];
 
     let data_iterate = vec![
+      0b00_00_00_00, //   0
       0b10_00_00_00, // 128
       0b11_00_00_00, // 192
       0b11_10_00_00, // 224

src/bitfield/mod.rs

@@ -17,6 +17,7 @@
 //! We need to make sure the performance impact of this stays well within
 //! bounds.
 
+mod iterator;
 mod masks;
 
 use self::masks::Masks;
@@ -57,6 +58,16 @@ impl Bitfield {
     }
   }
 
+  /// Get the current length
+  pub fn len(&self) -> usize {
+    self.length
+  }
+
+  /// Returns `true` if the bitfield is empty
+  pub fn is_empty(&self) -> bool {
+    self.length == 0
+  }
+
   /// Set a value at an index.
   pub fn set(&mut self, index: usize, value: bool) -> Change {
     let o = mask_8b(index);
@@ -113,17 +124,8 @@ impl Bitfield {
     let pos = (start - o) / 8;
     let last = (end - e) / 8;
 
-    let left_mask = if o == 0 {
-      255
-    } else {
-      255 - self.masks.data_iterate[o - 1]
-    };
-
-    let right_mask = if e == 0 {
-      0
-    } else {
-      self.masks.data_iterate[e - 1]
-    };
+    let left_mask = 255 - self.masks.data_iterate[o];
+    let right_mask = self.masks.data_iterate[e];
 
     let byte = self.data.get_byte(pos);
     if pos == last {
@@ -226,6 +228,25 @@ impl Bitfield {
       }
     }
   }
+
+  /// Constructs an iterator from start to end
+  pub fn iterator(&mut self) -> iterator::Iterator<'_> {
+    let len = self.length;
+    self.iterator_with_range(0, len)
+  }
+
+  /// Constructs an iterator from `start` to `end`
+  pub fn iterator_with_range(
+    &mut self,
+    start: usize,
+    end: usize,
+  ) -> iterator::Iterator<'_> {
+    let mut iter = iterator::Iterator::new(self);
+    iter.range(start, end);
+    iter.seek(0);
+
+    iter
+  }
 }
 
 // NOTE: can we move this into `sparse_bitfield`?

src/feed.rs

@@ -535,8 +535,8 @@ where
       }
     }
     Ok(Audit {
-      valid_blocks: valid_blocks,
-      invalid_blocks: invalid_blocks,
+      valid_blocks,
+      invalid_blocks,
     })
   }