From 4ad7092c0f99e164f8d51b7510fab6a043865e8f Mon Sep 17 00:00:00 2001 From: Buffrr Date: Wed, 31 Jul 2024 15:38:29 -0700 Subject: [PATCH] Consume self in write tx Necessary because insert/delete consume the in memory state and in case of error tx should be re-created --- src/tx.rs | 46 +++++++++++++++++++-------------------- tests/integration_test.rs | 25 ++++++++++----------- 2 files changed, 34 insertions(+), 37 deletions(-) diff --git a/src/tx.rs b/src/tx.rs index 5f12e15..6153ff8 100644 --- a/src/tx.rs +++ b/src/tx.rs @@ -373,28 +373,26 @@ impl<'db, H: NodeHasher> WriteTransaction<'db, H> { Ok(()) } - pub fn insert(&mut self, key: Hash, value: Vec) -> Result<()> { + pub fn insert(mut self, key: Hash, value: Vec) -> Result { if self.state.is_none() { self.state = Some(Node::from_leaf(Path(key), value)); - return Ok(()); + return Ok(self); } let mut state = self.state.take().unwrap(); state = self.insert_into_node(state, Path(key), value, 0)?; self.state = Some(state); - Ok(()) + Ok(self) } - pub fn delete(&mut self, key: Hash) -> Result>> { + pub fn delete(mut self, key: Hash) -> Result { if self.state.is_none() { - return Ok(None); + return Ok(self); } let state = self.state.take().unwrap(); - let (node, value) = self.delete_node(state, Path(key), 0)?; - self.state = node; - - Ok(value) + self.state = self.delete_node(state, Path(key), 0)?; + Ok(self) } fn insert_into_node( @@ -507,7 +505,7 @@ impl<'db, H: NodeHasher> WriteTransaction<'db, H> { node: Node, key: Path, depth: usize, - ) -> Result<(Option, Option>)> { + ) -> Result> { let inner = self.read_inner(node)?; return match inner { NodeInner::Leaf { @@ -515,9 +513,9 @@ impl<'db, H: NodeHasher> WriteTransaction<'db, H> { } => { if node_key != key { let node = Node::from_leaf(node_key, value); - return Ok((Some(node), None)); + return Ok(Some(node)); } - Ok((None, Some(value))) + Ok(None) } NodeInner::Internal { prefix, @@ -528,32 +526,32 @@ impl<'db, H: NodeHasher> WriteTransaction<'db, H> { // Traverse further based on the direction match key.direction(depth) { Direction::Right => { - let (right_subtree, value) = self.delete_node(*right, key, depth + 1)?; + let right_subtree = self.delete_node(*right, key, depth + 1)?; match right_subtree { None => { // Right subtree was deleted, move left subtree up let left_subtree = self.read_inner(*left)?; - Ok((Some(self.lift_node(prefix, left_subtree, Direction::Left)), value)) + Ok(Some(self.lift_node(prefix, left_subtree, Direction::Left))) } Some(right_subtree) => { - Ok((Some( + Ok(Some( Node::from_internal(prefix, left, Box::new(right_subtree)) - ), value)) + )) } } } Direction::Left => { - let (left_subtree, value) = self.delete_node(*left, key, depth + 1)?; + let left_subtree = self.delete_node(*left, key, depth + 1)?; return match left_subtree { None => { // Left subtree was deleted, move right subtree up let right_subtree = self.read_inner(*right)?; - Ok((Some(self.lift_node(prefix, right_subtree, Direction::Right)), value)) + Ok(Some(self.lift_node(prefix, right_subtree, Direction::Right))) } Some(left_subtree) => { - Ok((Some( + Ok(Some( Node::from_internal(prefix, Box::new(left_subtree), right) - ), value)) + )) } }; } @@ -774,10 +772,10 @@ mod tests { fn test_extended_proofs() { let db = Database::memory().unwrap(); let mut tx = db.begin_write().unwrap(); - tx.insert([0b1000_0000u8; 32], vec![1]).unwrap(); - tx.insert([0b1100_0000u8; 32], vec![2]).unwrap(); - tx.insert([0b0000_0000u8; 32], vec![3]).unwrap(); - tx.commit().unwrap(); + tx.insert([0b1000_0000u8; 32], vec![1]).unwrap() + .insert([0b1100_0000u8; 32], vec![2]).unwrap() + .insert([0b0000_0000u8; 32], vec![3]).unwrap() + .commit().unwrap(); let mut snapshot = db.begin_read().unwrap(); let standard_subtree = snapshot.prove(&[[0u8; 32]], ProofType::Standard).unwrap(); diff --git a/tests/integration_test.rs b/tests/integration_test.rs index ee726ae..1eece03 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -25,12 +25,11 @@ fn it_works_with_empty_trees() { #[test] fn it_inserts_into_tree() { let db = Database::memory().unwrap(); - let mut tx = db.begin_write().unwrap(); + let tx = db.begin_write().unwrap(); let key = db.hash(&[]); let value = "some data".as_bytes().to_vec(); - tx.insert(key.clone(), value.clone()).unwrap(); - tx.commit().unwrap(); + tx.insert(key.clone(), value.clone()).unwrap().commit().unwrap(); let mut tree = db.begin_read().unwrap(); @@ -59,7 +58,7 @@ fn it_inserts_many_items_into_tree() { keys.push(key.clone()); let value = format!("data{}", i).as_bytes().to_vec(); - tx.insert(key.clone(), value.clone()).unwrap(); + tx = tx.insert(key.clone(), value.clone()).unwrap(); subtree.insert(key, ValueOrHash::Value(value)).unwrap(); } @@ -94,7 +93,7 @@ fn it_should_iterate_over_tree() { for i in 0..n { let key = Sha256Hasher::hash(format!("key{}", i).as_bytes()); let value = format!("data{}", i).as_bytes().to_vec(); - tx.insert(key, value.clone()).unwrap(); + tx = tx.insert(key, value.clone()).unwrap(); inserted_values.insert(String::from_utf8(value).unwrap()); } @@ -127,7 +126,7 @@ fn it_returns_none_when_key_not_exists() { let key = db.hash(&[]); let value = "some data".as_bytes().to_vec(); - tx.insert(key.clone(), value.clone()).unwrap(); + tx = tx.insert(key.clone(), value.clone()).unwrap(); tx.commit().unwrap(); let mut tree = db.begin_read().unwrap(); @@ -163,7 +162,7 @@ fn it_should_delete_elements_from_snapshot() { let db = Database::memory().unwrap(); let mut tx = db.begin_write().unwrap(); for key in initial_set { - tx.insert(u32_to_key(key), vec![0]).unwrap(); + tx = tx.insert(u32_to_key(key), vec![0]).unwrap(); } tx.commit().unwrap(); @@ -172,7 +171,7 @@ fn it_should_delete_elements_from_snapshot() { // add all elements that we wish to delete let mut tx = db.begin_write().unwrap(); for key in &keys_to_delete { - tx.insert(u32_to_key(*key), vec![0]).unwrap(); + tx = tx.insert(u32_to_key(*key), vec![0]).unwrap(); } tx.commit().unwrap(); @@ -181,7 +180,7 @@ fn it_should_delete_elements_from_snapshot() { let mut tx = db.begin_write().unwrap(); for key in &keys_to_delete { - tx.delete(u32_to_key(*key)).unwrap(); + tx = tx.delete(u32_to_key(*key)).unwrap(); } tx.commit().unwrap(); @@ -212,7 +211,7 @@ fn it_should_delete_elements_from_subtree() { let db = Database::memory().unwrap(); let mut tx = db.begin_write().unwrap(); for key in initial_set { - tx.insert(u32_to_key(key), vec![0]).unwrap(); + tx = tx.insert(u32_to_key(key), vec![0]).unwrap(); } tx.commit().unwrap(); @@ -221,7 +220,7 @@ fn it_should_delete_elements_from_subtree() { // Add all elements that we wish to delete as well let mut tx = db.begin_write().unwrap(); for key in &keys_to_delete { - tx.insert(u32_to_key(*key), vec![0]).unwrap(); + tx = tx.insert(u32_to_key(*key), vec![0]).unwrap(); } tx.commit().unwrap(); @@ -254,7 +253,7 @@ fn it_should_store_metadata() { for i in 0..100 { let key = Sha256Hasher::hash(format!("key{}", i).as_bytes()); let value = format!("data{}", i).as_bytes().to_vec(); - tx.insert(key, value.clone()).unwrap(); + tx = tx.insert(key, value.clone()).unwrap(); } tx.metadata("snapshot 1".as_bytes().to_vec()).unwrap(); tx.commit().unwrap(); @@ -281,7 +280,7 @@ fn it_should_rollback() -> spacedb::Result<()> { for snapshot_index in 0..snapshots_len { let mut tx = db.begin_write()?; for entry in 0..items_per_snapshot { - tx.insert(u32_to_key((snapshot_index * entry) as u32), entry.to_be_bytes().to_vec())?; + tx = tx.insert(u32_to_key((snapshot_index * entry) as u32), entry.to_be_bytes().to_vec())?; } tx.commit()?; }