diff --git a/rust/treedirstate/src/dirstate.rs b/rust/treedirstate/src/dirstate.rs --- a/rust/treedirstate/src/dirstate.rs +++ b/rust/treedirstate/src/dirstate.rs @@ -36,6 +36,13 @@ Backend::File(ref file) => file, } } + + pub fn cache(&mut self) -> Result<()> { + match *self { + Backend::Empty(ref _null) => Ok(()), + Backend::File(ref mut file) => file.cache(), + } + } } /// A dirstate object. This contains the state of all files in the dirstate, stored in tree @@ -181,6 +188,7 @@ /// Get the name and state of the first file in the tracked tree. pub fn get_first_tracked<'a>(&'a mut self) -> Result> { + self.store.cache()?; self.tracked.get_first(self.store.store_view()) } diff --git a/rust/treedirstate/src/filestore.rs b/rust/treedirstate/src/filestore.rs --- a/rust/treedirstate/src/filestore.rs +++ b/rust/treedirstate/src/filestore.rs @@ -1,7 +1,7 @@ // Copyright Facebook, Inc. 2017 //! Implementation of a store using file I/O. -use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; +use byteorder::{BigEndian, ByteOrder, ReadBytesExt, WriteBytesExt}; use errors::*; use std::borrow::Cow; use std::cell::RefCell; @@ -40,6 +40,9 @@ /// True if the file is read-only. read_only: bool, + + /// Cache of data loaded from disk. Used when iterating over the whole dirstate. + cache: Option>, } impl FileStore { @@ -58,6 +61,7 @@ position: HEADER_LEN as u64, at_end: RefCell::new(true), read_only: false, + cache: None, }) } @@ -96,8 +100,27 @@ position, at_end: RefCell::new(true), read_only, + cache: None, }) } + + pub fn cache(&mut self) -> Result<()> { + if self.cache.is_none() { + let file = self.file.get_mut(); + file.flush()?; + file.seek(SeekFrom::Start(0))?; + let mut buffer = Vec::with_capacity(self.position as usize); + unsafe { + // This is safe as we've just allocated the buffer and are about to read into it. + buffer.set_len(self.position as usize); + } + file.get_mut().read_exact(buffer.as_mut_slice())?; + file.seek(SeekFrom::Start(self.position))?; + *self.at_end.get_mut() = true; + self.cache = Some(buffer); + } + Ok(()) + } } impl Store for FileStore { @@ -131,6 +154,23 @@ bail!(ErrorKind::InvalidStoreId(id)); } + if let Some(ref cache) = self.cache { + if (id as u64) < cache.len() as u64 { + if (id as u64) > cache.len() as u64 - 4 { + // The ID falls in the last 3 bytes of the cache. This is invalid. + bail!(ErrorKind::InvalidStoreId(id)); + } + let start = id as usize + 4; + let size = BigEndian::read_u32(&cache[id as usize..start]) as usize; + if size as u64 > cache.len() as u64 - start as u64 { + // The stored size of this block exceeds the number of bytes left in the + // cache. We must have been given an invalid ID. + bail!(ErrorKind::InvalidStoreId(id)); + } + return Ok(Cow::from(&cache[start..start + size])); + } + } + // Get mutable access to the file, and seek to the right location. let mut file = self.file.borrow_mut(); file.seek(SeekFrom::Start(id as u64))?; @@ -273,6 +313,32 @@ } #[test] + fn cache() { + let dir = TempDir::new("filestore_test").expect("create temp dir"); + let p = dir.path().join("store"); + let mut s = FileStore::create(&p).expect("create store"); + let id1 = s.append("data block 1".as_bytes()).expect("write block 1"); + let id2 = s.append("data block two".as_bytes()) + .expect("write block 2"); + s.flush().expect("flush"); + assert_eq!(s.read(id1).expect("read 1"), "data block 1".as_bytes()); + assert_eq!(s.read(id2).expect("read 2"), "data block two".as_bytes()); + drop(s); + let mut s = FileStore::open(&p).expect("open store"); + s.cache().expect("can cache"); + assert_eq!(s.read(id1).expect("read 1"), "data block 1".as_bytes()); + assert_eq!(s.read(id2).expect("read 2"), "data block two".as_bytes()); + let id3 = s.append("third data block".as_bytes()) + .expect("write block 3"); + s.flush().expect("flush"); + assert_eq!(s.read(id3).expect("read 3"), "third data block".as_bytes()); + assert_eq!( + s.read(id3 - 2).unwrap_err().to_string(), + format!("invalid store id: {}", id3 - 2) + ); + } + + #[test] fn invalid_store_ids() { let dir = TempDir::new("filestore_test").expect("create temp dir"); let p = dir.path().join("store");