Store test for sumtree rewind

core/src/core/pmmr.rs

@@ -284,7 +284,7 @@ where
 		Ok(elmt_pos)
 	}
 
-	/// Rewind the PMMR to a previous position, as is all push operations after
+	/// Rewind the PMMR to a previous position, as if all push operations after
 	/// that had been canceled. Expects a position in the PMMR to rewind to as
 	/// well as the consumer-provided index of when the change occurred.
 	pub fn rewind(&mut self, position: u64, index: u32) -> Result<(), String> {
@@ -302,9 +302,8 @@ where
 	}
 
 	/// Prune an element from the tree given its position. Note that to be able
-	/// to
+	/// to provide that position and prune, consumers of this API are expected
-	/// provide that position and prune, consumers of this API are expected to
+	/// to keep an index of elements to positions in the tree. Prunes parent
-	/// keep an index of elements to positions in the tree. Prunes parent
 	/// nodes as well when they become childless.
 	pub fn prune(&mut self, position: u64, index: u32) -> Result<bool, String> {
 		if let None = self.backend.get(position) {

store/src/sumtree.rs

@@ -442,8 +442,7 @@ where
 	/// otherwise the RM_LOG_MAX_NODES default value is used.
 	///
 	/// TODO whatever is calling this should also clean up the commit to
-	/// position
+	/// position index in db
-	/// index in db
 	pub fn check_compact(&mut self, max_len: usize) -> io::Result<()> {
 		if !(max_len > 0 && self.remove_log.len() > max_len ||
 			     max_len == 0 && self.remove_log.len() > RM_LOG_MAX_NODES)

store/tests/sumtree.rs

@@ -144,6 +144,64 @@ fn sumtree_reload() {
 	}
 }
 
+#[test]
+fn sumtree_rewind() {
+	let (data_dir, elems) = setup();
+	let mut backend = store::sumtree::PMMRBackend::new(data_dir).unwrap();
+
+	// adding elements and keeping the corresponding root
+	let mut mmr_size = load(0, &elems[0..4], &mut backend);
+	backend.sync().unwrap();
+	let root1: HashSum<TestElem>;
+	{
+		let pmmr = PMMR::at(&mut backend, mmr_size);
+		root1 = pmmr.root();
+	}
+
+	mmr_size = load(mmr_size, &elems[4..6], &mut backend);
+	backend.sync().unwrap();
+	let root2: HashSum<TestElem>;
+	{
+		let pmmr = PMMR::at(&mut backend, mmr_size);
+		root2 = pmmr.root();
+	}
+
+	mmr_size = load(mmr_size, &elems[6..9], &mut backend);
+	backend.sync().unwrap();
+
+	// prune and compact the 2 first elements to spice things up
+	{
+		let mut pmmr = PMMR::at(&mut backend, mmr_size);
+		pmmr.prune(1, 1).unwrap();
+		pmmr.prune(2, 1).unwrap();
+	}
+	backend.check_compact(1).unwrap();
+	backend.sync().unwrap();
+	
+	// rewind and check the roots still match
+	{
+		let mut pmmr = PMMR::at(&mut backend, mmr_size);
+		pmmr.rewind(9, 3).unwrap();
+		assert_eq!(pmmr.root(), root2);
+	}
+	backend.sync().unwrap();
+	{
+		let pmmr = PMMR::at(&mut backend, 10);
+		assert_eq!(pmmr.root(), root2);
+	}
+
+	{
+		let mut pmmr = PMMR::at(&mut backend, 10);
+		pmmr.rewind(5, 3).unwrap();
+		assert_eq!(pmmr.root(), root1);
+	}
+	backend.sync().unwrap();
+	{
+		let pmmr = PMMR::at(&mut backend, 7);
+		assert_eq!(pmmr.root(), root1);
+	}
+}
+
 fn setup() -> (String, Vec<TestElem>) {
 	let _ = env_logger::init();
 	let t = time::get_time();