fix various issues with rewinding data file (#983)

* fix issue with rewinding data file
also fix fragility around an empty MMR (and the memory map)

* cleanup

core/src/core/pmmr.rs

@@ -460,7 +460,7 @@ where
 		Ok(true)
 	}
 
-	/// Get a hash at provided position in the MMR.
+	/// Get the hash at provided position in the MMR.
 	pub fn get_hash(&self, pos: u64) -> Option<Hash> {
 		if pos > self.last_pos {
 			None
@@ -469,7 +469,7 @@ where
 		}
 	}
 
-	/// Get the data element at provided in the MMR.
+	/// Get the data element at provided position in the MMR.
 	pub fn get_data(&self, pos: u64) -> Option<T> {
 		if pos > self.last_pos {
 			None

store/src/pmmr.rs

@@ -189,14 +189,16 @@ where
 			.rewind(index)
 			.map_err(|e| format!("Could not truncate remove log: {}", e))?;
 
+		// Hash file
 		let shift = self.pruned_nodes.get_shift(position).unwrap_or(0);
 		let record_len = 32;
 		let file_pos = (position - shift) * (record_len as u64);
 		self.hash_file.rewind(file_pos);
 
 		// Data file
-		let flatfile_pos = pmmr::n_leaves(position) - 1;
-		let file_pos = (flatfile_pos as usize + 1) * T::len();
+		let leaf_shift = self.pruned_nodes.get_leaf_shift(position).unwrap_or(0);
+		let flatfile_pos = pmmr::n_leaves(position);
+		let file_pos = (flatfile_pos - leaf_shift) * T::len() as u64;
 		self.data_file.rewind(file_pos as u64);
 		Ok(())
 	}

store/src/types.rs

@@ -102,11 +102,20 @@ impl AppendOnlyFile {
 			self.truncate(self.buffer_start)?;
 			self.buffer_start_bak = 0;
 		}
+
 		self.buffer_start += self.buffer.len();
 		self.file.write(&self.buffer[..])?;
 		self.file.sync_all()?;
+
 		self.buffer = vec![];
+
+		// Note: file must be non-empty to memory map it
+		if self.file.metadata()?.len() == 0 {
+			self.mmap = None;
+		} else {
 			self.mmap = Some(unsafe { memmap::Mmap::map(&self.file)? });
+		}
+
 		Ok(())
 	}
 
@@ -126,10 +135,13 @@ impl AppendOnlyFile {
 		self.buffer = vec![];
 	}
 
-	/// Read length bytes of data at offset from the file. Leverages the memory
-	/// map.
+	/// Read length bytes of data at offset from the file.
+	/// Leverages the memory map.
 	pub fn read(&self, offset: usize, length: usize) -> Vec<u8> {
 		if offset >= self.buffer_start {
+			if self.buffer.is_empty() {
+				return vec![];
+			}
 			let offset = offset - self.buffer_start;
 			return self.buffer[offset..(offset + length)].to_vec();
 		}

store/tests/pmmr.rs

@@ -21,7 +21,6 @@ use std::fs;
 
 use core::ser::*;
 use core::core::pmmr::{Backend, PMMR};
-use core::core::hash::Hash;
 use store::types::prune_noop;
 
 #[test]
@@ -292,30 +291,33 @@ fn pmmr_rewind() {
 	// adding elements and keeping the corresponding root
 	let mut mmr_size = load(0, &elems[0..4], &mut backend);
 	backend.sync().unwrap();
-	let root1: Hash;
-	{
+	let root1 = {
 		let pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
-		root1 = pmmr.root();
-	}
+		pmmr.root()
+	};
 
 	mmr_size = load(mmr_size, &elems[4..6], &mut backend);
 	backend.sync().unwrap();
-	let root2: Hash;
-	{
+	let root2 = {
 		let pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, mmr_size);
-		root2 = pmmr.root();
-	}
+		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
+	// prune the first 4 elements (leaves at pos 1, 2, 4, 5)
 	{
 		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();
+		pmmr.prune(5, 1).unwrap();
 	}
-	backend.check_compact(1, 2, &prune_noop).unwrap();
+	backend.sync().unwrap();
+
+	// and compact the MMR to remove the 2 pruned elements
+	backend.check_compact(2, 2, &prune_noop).unwrap();
 	backend.sync().unwrap();
 
 	// rewind and check the roots still match
@@ -330,6 +332,17 @@ fn pmmr_rewind() {
 		assert_eq!(pmmr.root(), root2);
 	}
 
+	// also check the data file looks correct
+	// everything up to and including pos 7 should be pruned from the data file
+	// except the data at pos 8 and 9 (corresponding to elements 5 and 6)
+	for pos in 1..8 {
+		assert_eq!(backend.get_data(pos), None);
+	}
+	assert_eq!(backend.get_data(8), Some(elems[4]));
+	assert_eq!(backend.get_data(9), Some(elems[5]));
+
+	assert_eq!(backend.data_size().unwrap(), 2);
+
 	{
 		let mut pmmr: PMMR<TestElem, _> = PMMR::at(&mut backend, 10);
 		pmmr.rewind(5, 3).unwrap();
@@ -341,6 +354,17 @@ fn pmmr_rewind() {
 		assert_eq!(pmmr.root(), root1);
 	}
 
+	// also check the data file looks correct
+	// everything up to and including pos 7 should be pruned from the data file
+	// but we have rewound to pos 5 so everything after that should be None
+	for pos in 1..10 {
+		assert_eq!(backend.get_data(pos), None);
+	}
+
+	// check we have no data in the backend after
+	// pruning, compacting and rewinding
+	assert_eq!(backend.data_size().unwrap(), 0);
+
 	teardown(data_dir);
 }