Return DiskAddress instead of mutating (#206)

ava-labs · Aug 22, 2023 · da6bef3 · da6bef3
1 parent a981b6b
commit da6bef3
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 29 deletions.
diff --git a/firewood/benches/hashops.rs b/firewood/benches/hashops.rs
@@ -47,8 +47,7 @@ fn bench_insert(b: &mut Bencher) {
     )
     .unwrap();
     let mut merkle = Merkle::new(Box::new(store));
-    let mut root = DiskAddress::null();
-    merkle.init_root(&mut root).unwrap();
+    let root = merkle.init_root().unwrap();
 
     let mut rng = rand::rngs::StdRng::seed_from_u64(1234);
     const KEY_LEN: usize = 4;

diff --git a/firewood/src/db.rs b/firewood/src/db.rs
@@ -737,8 +737,9 @@ impl Db {
             db_header_ref
                 .write(|r| {
                     err = (|| {
-                        merkle.init_root(&mut r.acc_root)?;
-                        merkle.init_root(&mut r.kv_root)
+                        r.acc_root = merkle.init_root()?;
+                        r.kv_root = merkle.init_root()?;
+                        Ok(())
                     })();
                 })
                 .unwrap();

diff --git a/firewood/src/merkle.rs b/firewood/src/merkle.rs
@@ -878,9 +878,8 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
         Self { store }
     }
 
-    pub fn init_root(&self, root: &mut DiskAddress) -> Result<(), MerkleError> {
-        *root = self
-            .store
+    pub fn init_root(&self) -> Result<DiskAddress, MerkleError> {
+        self.store
             .put_item(
                 Node::new(NodeType::Branch(BranchNode {
                     chd: [None; NBRANCH],
@@ -889,9 +888,8 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
                 })),
                 Node::max_branch_node_size(),
             )
-            .map_err(MerkleError::Shale)?
-            .as_ptr();
-        Ok(())
+            .map_err(MerkleError::Shale)
+            .map(|node| node.as_ptr())
     }
 
     pub fn get_store(&self) -> &dyn ShaleStore<Node> {

diff --git a/firewood/src/merkle_util.rs b/firewood/src/merkle_util.rs
@@ -142,10 +142,9 @@ pub fn new_merkle(
     let space =
         shale::compact::CompactSpace::new(mem_meta, mem_payload, compact_header, cache, 10, 16)
             .expect("CompactSpace init fail");
-    let mut root = DiskAddress::null();
 
     let merkle = Merkle::new(Box::new(space));
-    merkle.init_root(&mut root).unwrap();
+    let root = merkle.init_root().unwrap();
 
     MerkleSetup { root, merkle }
 }
diff --git a/shale/src/compact.rs b/shale/src/compact.rs
@@ -558,6 +558,16 @@ impl<T: Storable + Send + Sync, M: CachedStore> CompactSpaceInner<T, M> {
         f.write(|f| f.payload_size = length).unwrap();
         Ok((offset + CompactHeader::MSIZE as usize).0.unwrap().get() as u64)
     }
+
+    fn alloc(&mut self, length: u64) -> Result<u64, ShaleError> {
+        self.alloc_from_freed(length).and_then(|addr| {
+            if let Some(addr) = addr {
+                Ok(addr)
+            } else {
+                self.alloc_new(length)
+            }
+        })
+    }
 }
 
 #[derive(Debug)]
@@ -591,27 +601,20 @@ impl<T: Storable + Send + Sync, M: CachedStore> CompactSpace<T, M> {
 impl<T: Storable + Send + Sync + Debug + 'static, M: CachedStore + Send + Sync> ShaleStore<T>
     for CompactSpace<T, M>
 {
-    fn put_item<'a>(&'a self, item: T, extra: u64) -> Result<ObjRef<'a, T>, ShaleError> {
+    fn put_item(&self, item: T, extra: u64) -> Result<ObjRef<'_, T>, ShaleError> {
         let size = item.dehydrated_len() + extra;
-        let ptr: DiskAddress = {
-            let mut inner = self.inner.write().unwrap();
-            let addr = if let Some(addr) = inner.alloc_from_freed(size)? {
-                addr
-            } else {
-                inner.alloc_new(size)?
-            };
+        let addr = self.inner.write().unwrap().alloc(size)?;
 
-            DiskAddress::from(addr as usize)
-        };
-        let inner = self.inner.read().unwrap();
+        let mut u = {
+            let inner = self.inner.read().unwrap();
+            let compact_space = inner.compact_space.as_ref();
+            let view =
+                StoredView::item_to_obj(compact_space, addr.try_into().unwrap(), size, item)?;
 
-        let mut u: ObjRef<'a, T> = inner.obj_cache.put(StoredView::item_to_obj(
-            inner.compact_space.as_ref(),
-            ptr.into(),
-            size,
-            item,
-        )?);
+            inner.obj_cache.put(view)
+        };
 
+        // should this use a `?` instead of `unwrap`?
         u.write(|_| {}).unwrap();
 
         Ok(u)