Merge pull request #1383 from hannobraun/store

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);