Fix Store iteration bug

crates/fj-kernel/src/storage/blocks.rs

@@ -55,8 +55,13 @@ impl<T> Blocks<T> {
  block.insert(index.object_index, object)
  }
 
- pub fn get(&self, index: usize) -> Option<&Block<T>> {
- self.inner.get(index)
+ pub fn get_and_inc(&self, index: &mut Index) -> Option<&Option<T>> {
+ let block = self.inner.get(index.block_index.0)?;
+ let object = block.get(index.object_index);
+
+ index.inc(block);
+
+ Some(object)
  }
 
  #[cfg(test)]
@@ -105,8 +110,8 @@ impl<T> Block<T> {
  &self.objects[index.0]
  }
 
- pub fn get(&self, index: usize) -> &Option<T> {
- &self.objects[index]
+ pub fn get(&self, index: ObjectIndex) -> &Option<T> {
+ &self.objects[index.0]
  }
 
  pub fn len(&self) -> usize {
@@ -121,7 +126,7 @@ impl<T> Block<T> {
  return None;
  }
 
- let object = self.get(i).as_ref()?;
+ let object = self.get(ObjectIndex(i)).as_ref()?;
  i += 1;
 
  Some(object)
@@ -135,6 +140,23 @@ pub struct Index {
  object_index: ObjectIndex,
 }
 
+impl Index {
+ pub fn zero() -> Self {
+ Self {
+ block_index: BlockIndex(0),
+ object_index: ObjectIndex(0),
+ }
+ }
+
+ pub fn inc<T>(&mut self, block: &Block<T>) {
+ self.object_index.0 += 1;
+ if self.object_index.0 >= block.len() {
+ self.block_index.0 += 1;
+ self.object_index.0 = 0;
+ }
+ }
+}
+
 #[derive(Clone, Copy, Debug)]
 pub struct BlockIndex(usize);
 

crates/fj-kernel/src/storage/store.rs

@@ -38,9 +38,17 @@ pub struct Store<T> {
 
 impl<T> Store<T> {
  /// Construct a new instance of `Store`
+ ///
+ /// Equivalent to calling [`Store::with_block_size`] with a default block
+ /// size.
  pub fn new() -> Self {
+ Self::with_block_size(16384)
+ }
+
+ /// Construct a new instance of `Store` using the provided block size
+ pub fn with_block_size(block_size: usize) -> Self {
  let inner = Arc::new(RwLock::new(StoreInnerInner {
- blocks: Blocks::new(BLOCK_SIZE),
+ blocks: Blocks::new(block_size),
  }));
 
  Self { inner }
@@ -61,8 +69,7 @@ impl<T> Store<T> {
  pub fn iter(&self) -> Iter<T> {
  Iter {
  store: self.inner.clone(),
- next_block: 0,
- next_object: 0,
+ next_index: Index::zero(),
  _a: PhantomData,
  }
  }
@@ -103,8 +110,7 @@ impl<'a, T> IntoIterator for &'a Store<T> {
 /// An iterator over objects in a [`Store`]
 pub struct Iter<'a, T> {
  store: StoreInner<T>,
- next_block: usize,
- next_object: usize,
+ next_index: Index,
  _a: PhantomData<&'a ()>,
 }
 
@@ -114,13 +120,7 @@ impl<'a, T: 'a> Iterator for Iter<'a, T> {
  fn next(&mut self) -> Option<Self::Item> {
  let inner = self.store.read();
 
- let block = inner.blocks.get(self.next_block)?;
- let object = block.get(self.next_object);
-
- self.next_object += 1;
- if self.next_object >= block.len() {
- self.next_block += 1;
- }
+ let object = inner.blocks.get_and_inc(&mut self.next_index)?;
 
  Some(Handle {
  store: self.store.clone(),
@@ -175,15 +175,13 @@ pub struct StoreInnerInner<T> {
  blocks: Blocks<T>,
 }
 
-const BLOCK_SIZE: usize = 16384;
-
 #[cfg(test)]
 mod tests {
  use super::Store;
 
  #[test]
  fn insert_and_handle() {
- let store = Store::new();
+ let store = Store::with_block_size(1);
 
  let object = 0;
  let handle = store.insert(object);
@@ -193,7 +191,7 @@ mod tests {
 
  #[test]
  fn insert_and_iter() {
- let store = Store::new();
+ let store = Store::with_block_size(1);
 
  let a = store.insert(0);
  let b = store.insert(1);