Skip to content

Commit

Permalink
Consume self in write tx
Browse files Browse the repository at this point in the history 
Necessary because insert/delete consume the in memory state and in case of error tx should be re-created
  • Loading branch information
@buffrr
buffrr committed Aug 1, 2024
1 parent ddd9849 commit 4ad7092
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 37 deletions.
46 changes: 22 additions & 24 deletions src/tx.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -373,28 +373,26 @@ impl<'db, H: NodeHasher> WriteTransaction<'db, H> {
Ok(())
}

pub fn insert(&mut self, key: Hash, value: Vec<u8>) -> Result<()> {
pub fn insert(mut self, key: Hash, value: Vec<u8>) -> Result<Self> {
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<Option<Vec<u8>>> {
pub fn delete(mut self, key: Hash) -> Result<Self> {
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(
Expand Down Expand Up @@ -507,17 +505,17 @@ impl<'db, H: NodeHasher> WriteTransaction<'db, H> {
node: Node,
key: Path<Hash>,
depth: usize,
) -> Result<(Option<Node>, Option<Vec<u8>>)> {
) -> Result<Option<Node>> {
let inner = self.read_inner(node)?;
return match inner {
NodeInner::Leaf {
key: node_key, value
} => {
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,
Expand All @@ -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))
))
}
};
}
Expand Down Expand Up @@ -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();
Expand Down
25 changes: 12 additions & 13 deletions tests/integration_test.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand Down Expand Up @@ -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();
}

Expand Down Expand Up @@ -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());
}

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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();

Expand All @@ -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();

Expand All @@ -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();

Expand Down Expand Up @@ -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();

Expand All @@ -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();

Expand Down Expand Up @@ -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();
Expand All @@ -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()?;
}
Expand Down

0 comments on commit 4ad7092

Please sign in to comment.