diff --git a/storage/src/tree_state.rs b/storage/src/tree_state.rs index 67703b63..c88a526f 100644 --- a/storage/src/tree_state.rs +++ b/storage/src/tree_state.rs @@ -189,6 +189,7 @@ pub struct TreeState { left: Option, right: Option, parents: Vec>, + is_empty: bool, } impl TreeState { @@ -198,6 +199,7 @@ impl TreeState { left: None, right: None, parents: vec![None; D::HEIGHT - 1], + is_empty: true, } } @@ -229,10 +231,16 @@ impl TreeState { return Err("Appending to full tree"); } + + self.is_empty = false; Ok(()) } pub fn root(&self) -> H256 { + if self.is_empty { + return Self::empty_root(); + } + let left = self.left.as_ref().unwrap_or(&H::empty()[0]); let right = self.right.as_ref().unwrap_or(&H::empty()[0]); @@ -273,6 +281,10 @@ impl serialization::Deserializable for TreeState { tree_state.right = reader.read()?; tree_state.parents = reader.read_list()?; + tree_state.is_empty = tree_state.left.is_none() + && tree_state.right.is_none() + && tree_state.parents.iter().all(Option::is_none); + Ok(tree_state) } } @@ -508,6 +520,7 @@ mod tests { tree.append(TEST_COMMITMENTS[i].clone()).expect(&format!("Failed to add commitment #{}", i)); } + assert!(!tree.is_empty); assert_eq!(tree.root(), H256::from("0bf622cb9f901b7532433ea2e7c1b7632f5935899b62dcf897a71551997dc8cc")); let mut stream = serialization::Stream::new(); @@ -518,9 +531,24 @@ mod tests { let mut reader = serialization::Reader::new(&bytes[..]); let deserialized_tree: TestSproutTreeState = reader.read().expect("Failed to deserialize"); + assert!(!tree.is_empty); assert_eq!(deserialized_tree.root(), H256::from("0bf622cb9f901b7532433ea2e7c1b7632f5935899b62dcf897a71551997dc8cc")); } + #[test] + fn serde_empty() { + let tree = TestSproutTreeState::new(); + let mut stream = serialization::Stream::new(); + assert!(tree.is_empty); + stream.append(&tree); + + let bytes = stream.out(); + + let mut reader = serialization::Reader::new(&bytes[..]); + let deserialized_tree: TestSproutTreeState = reader.read().expect("Failed to deserialize"); + assert!(deserialized_tree.is_empty); + } + #[test] fn sapling_empty_root() { let expected_root = H256::from_reversed_str("3e49b5f954aa9d3545bc6c37744661eea48d7c34e3000d82b7f0010c30f4c2fb");