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::{ByteOrder, BigEndian, 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, }) } @@ -90,8 +94,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 { @@ -125,6 +148,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))?;