From fffee377ddd523aa7e301006fc0a272134e62ef0 Mon Sep 17 00:00:00 2001
From: Antioch Peverell <30642645+antiochp@users.noreply.github.com>
Date: Wed, 6 Jun 2018 22:46:21 +0100
Subject: [PATCH] Fix 2nd compact (#1143)

* add failing test to cover case where we compact
an already compacted data file

* fix the logic for pruning the data file
---
 store/src/pmmr.rs   |  7 +++--
 store/tests/pmmr.rs | 74 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/store/src/pmmr.rs b/store/src/pmmr.rs
index ce70ab38f..baf37e646 100644
--- a/store/src/pmmr.rs
+++ b/store/src/pmmr.rs
@@ -335,8 +335,9 @@ where
 			let record_len = T::len() as u64;
 
 			let off_to_rm = map_vec!(leaf_pos_to_rm, |pos| {
-				let shift = self.pruned_nodes.get_leaf_shift(*pos);
-				(pmmr::n_leaves(pos - shift.unwrap()) - 1) * record_len
+				let flat_pos = pmmr::n_leaves(*pos);
+				let shift = self.pruned_nodes.get_leaf_shift(*pos).unwrap();
+				(flat_pos - 1 - shift) * record_len
 			});
 
 			self.data_file.save_prune(
@@ -352,6 +353,8 @@ where
 			for &pos in &rm_pre_cutoff {
 				self.pruned_nodes.add(pos);
 			}
+			// TODO - we can get rid of leaves in the prunelist here (and things still work)
+			// self.pruned_nodes.pruned_nodes.retain(|&x| !pmmr::is_leaf(x));
 
 			write_vec(
 				format!("{}/{}", self.data_dir, PMMR_PRUNED_FILE),
diff --git a/store/tests/pmmr.rs b/store/tests/pmmr.rs
index 5e470444c..c3b0ebe27 100644
--- a/store/tests/pmmr.rs
+++ b/store/tests/pmmr.rs
@@ -19,8 +19,8 @@ extern crate time;
 
 use std::fs;
 
-use core::ser::*;
 use core::core::pmmr::{Backend, PMMR};
+use core::ser::*;
 use store::types::prune_noop;
 
 #[test]
@@ -566,6 +566,78 @@ fn pmmr_compact_horizon() {
 	teardown(data_dir);
 }
 
+#[test]
+fn compact_twice() {
+	let (data_dir, elems) = setup("compact_twice");
+
+	// setup the mmr store with all elements
+	let mut backend = store::pmmr::PMMRBackend::new(data_dir.to_string()).unwrap();
+	let mmr_size = load(0, &elems[..], &mut backend);
+	backend.sync().unwrap();
+
+	// save the root
+	let root = {
+		let pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
+		pmmr.root()
+	};
+
+	// pruning some choice nodes
+	{
+		let mut pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
+		pmmr.prune(1, 1).unwrap();
+		pmmr.prune(2, 1).unwrap();
+		pmmr.prune(4, 1).unwrap();
+	}
+	backend.sync().unwrap();
+
+	// check the root and stored data
+	{
+		let pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
+		assert_eq!(root, pmmr.root());
+		assert_eq!(pmmr.get_data(5).unwrap(), TestElem(4));
+		assert_eq!(pmmr.get_data(11).unwrap(), TestElem(7));
+	}
+
+	// compact
+	backend.check_compact(2, 2, &prune_noop).unwrap();
+
+	// recheck the root and stored data
+	{
+		let pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
+		assert_eq!(root, pmmr.root());
+		assert_eq!(pmmr.get_data(5).unwrap(), TestElem(4));
+		assert_eq!(pmmr.get_data(11).unwrap(), TestElem(7));
+	}
+
+	// now prune some more nodes
+	{
+		let mut pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
+		pmmr.prune(5, 2).unwrap();
+		pmmr.prune(8, 2).unwrap();
+		pmmr.prune(9, 2).unwrap();
+	}
+	backend.sync().unwrap();
+
+	// recheck the root and stored data
+	{
+		let pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
+		assert_eq!(root, pmmr.root());
+		assert_eq!(pmmr.get_data(11).unwrap(), TestElem(7));
+	}
+
+	// compact
+	backend.check_compact(2, 3, &prune_noop).unwrap();
+
+	// recheck the root and stored data
+	{
+		let pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
+		assert_eq!(root, pmmr.root());
+		assert_eq!(pmmr.get_data(11).unwrap(), TestElem(7));
+	}
+
+	teardown(data_dir);
+}
+
 fn setup(tag: &str) -> (String, Vec<TestElem>) {
 	let _ = env_logger::init();
 	let t = time::get_time();