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 @@ -186,6 +186,15 @@ self.tracked.get(self.store.store_view(), name) } + /// Visit all tracked files with a visitor. + pub fn visit_tracked(&mut self, visitor: &mut F) -> Result<()> + where + F: FnMut(&Vec, &T) -> Result<()>, + { + self.store.cache()?; + self.tracked.visit(self.store.store_view(), visitor) + } + /// 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()?; @@ -206,6 +215,14 @@ self.removed.get(self.store.store_view(), name) } + /// Visit all removed files with a visitor. + pub fn visit_removed(&mut self, visitor: &mut F) -> Result<()> + where + F: FnMut(&Vec, &T) -> Result<()>, + { + self.removed.visit(self.store.store_view(), visitor) + } + /// Get the name and state of the first file in the removed tree. pub fn get_first_removed<'a>(&'a mut self) -> Result> { self.removed.get_first(self.store.store_view()) diff --git a/rust/treedirstate/src/errors.rs b/rust/treedirstate/src/errors.rs --- a/rust/treedirstate/src/errors.rs +++ b/rust/treedirstate/src/errors.rs @@ -15,6 +15,10 @@ description("invalid store id"), display("invalid store id: {}", id), } + InternalError { + description("internal error"), + display("internal error"), + } } foreign_links { Io(::std::io::Error); 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::{ByteOrder, BigEndian, ReadBytesExt, WriteBytesExt}; +use byteorder::{BigEndian, ByteOrder, ReadBytesExt, WriteBytesExt}; use errors::*; use std::borrow::Cow; use std::cell::RefCell; @@ -161,7 +161,7 @@ // cache. We must have been given an invalid ID. bail!(ErrorKind::InvalidStoreId(id)); } - return Ok(Cow::from(&cache[start..start + size])) + return Ok(Cow::from(&cache[start..start + size])); } } diff --git a/rust/treedirstate/src/python.rs b/rust/treedirstate/src/python.rs --- a/rust/treedirstate/src/python.rs +++ b/rust/treedirstate/src/python.rs @@ -4,10 +4,12 @@ use cpython::*; use cpython::exc; use dirstate::Dirstate; +use errors::ErrorKind; use filestate::FileState; use std::cell::RefCell; use std::path::PathBuf; use store::BlockId; +use tree::KeyRef; py_module_initializer!( rusttreedirstate, @@ -169,6 +171,28 @@ map_pyerr!(py, dirstate.has_removed_dir(dirname.data(py)), exc::IOError) } + def visittrackedfiles(&self, target: PyObject) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let mut visitor = |filepath: &Vec, _state: &FileState| { + let filename = PyBytes::new(py, &filepath.concat()).into_object(); + target.call(py, (filename,), None).map_err(|_e| ErrorKind::InternalError)?; + Ok(()) + }; + map_pyerr!(py, dirstate.visit_tracked(&mut visitor), exc::IOError)?; + Ok(py.None()) + } + + def visitremovedfiles(&self, target: PyObject) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let mut visitor = |filepath: &Vec, _state: &FileState| { + let filename = PyBytes::new(py, &filepath.concat()).into_object(); + target.call(py, (filename,), None).map_err(|_e| ErrorKind::InternalError)?; + Ok(()) + }; + map_pyerr!(py, dirstate.visit_removed(&mut visitor), exc::IOError)?; + Ok(py.None()) + } + // Get the next dirstate object after the provided filename. If the filename is None, // returns the first file in the tree. If the provided filename is the last file, returns // None. diff --git a/rust/treedirstate/src/tree.rs b/rust/treedirstate/src/tree.rs --- a/rust/treedirstate/src/tree.rs +++ b/rust/treedirstate/src/tree.rs @@ -219,6 +219,31 @@ } } + // Visit all of the files in under this node, by calling the visitor function on each one. + fn visit<'a, F>( + &'a mut self, + store: &StoreView, + path: &mut Vec>, + visitor: &mut F, + ) -> Result<()> + where + F: FnMut(&Vec, &T) -> Result<()>, + { + for (name, entry) in self.load_entries(store)?.iter_mut() { + path.push(name); + match entry { + &mut NodeEntry::Directory(ref mut node) => { + node.visit(store, path, visitor)?; + } + &mut NodeEntry::File(ref file) => { + visitor(path, file)?; + } + } + path.pop(); + } + Ok(()) + } + /// Get the first file in the subtree under this node. If the subtree is not empty, returns a /// pair containing the path to the file as a reversed vector of key references for each path /// element, and a reference to the file. @@ -455,6 +480,14 @@ Ok(self.root.get(store, name)?) } + pub fn visit(&mut self, store: &StoreView, visitor: &mut F) -> Result<()> + where + F: FnMut(&Vec, &T) -> Result<()>, + { + let mut path = Vec::new(); + self.root.visit(store, &mut path, visitor) + } + pub fn get_first<'a>(&'a mut self, store: &StoreView) -> Result> { Ok(self.root.get_first(store)?.map(|(mut path, file)| { path.reverse(); @@ -499,7 +532,7 @@ use store::NullStore; use store::tests::MapStore; - use tree::Tree; + use tree::{KeyRef, Tree}; use filestate::FileState; #[test] @@ -562,5 +595,15 @@ t.write_full(&mut ms2, &ms).expect("write to new store"); t.remove(&ms, b"bar/baz/foo").expect("remove"); t.write_delta(&mut ms2).expect("write delta"); + + let mut files = Vec::new(); + { + let mut v = |path: &Vec, _fs: &FileState| { + files.push(path.concat()); + Ok(()) + }; + t.visit(&mut ms2, &mut v).expect("visit should succeed"); + } + assert_eq!(files, vec![b"foo/bar/baz", b"foo/bar/xyz", b"foo/baz/foo"]); } } diff --git a/treedirstate/__init__.py b/treedirstate/__init__.py --- a/treedirstate/__init__.py +++ b/treedirstate/__init__.py @@ -52,10 +52,9 @@ self.buffer.write(v) class treedirstatemapiterator(object): - def __init__(self, map_, removed=False, items=False): + def __init__(self, map_, removed=False): self._rmap = map_ self._removed = removed - self._items = items self._at = None def __iter__(self): @@ -66,7 +65,7 @@ if nextitem is None: raise StopIteration self._at = nextitem[0] - return nextitem if self._items else nextitem[0] + return nextitem def next(self): return self.__next__() @@ -111,14 +110,14 @@ def itertrackeditems(self): """Returns an iterator over (filename, (state, mode, size, mtime)).""" - return treedirstatemapiterator(self._rmap, removed=False, items=True) + return treedirstatemapiterator(self._rmap, removed=False) def iterremoveditems(self): """ Returns an iterator over (filename, (state, mode, size, mtime)) for files that have been marked as removed. """ - return treedirstatemapiterator(self._rmap, removed=True, items=True) + return treedirstatemapiterator(self._rmap, removed=True) def iteritems(self): return itertools.chain(self.itertrackeditems(), @@ -155,16 +154,25 @@ return self.hastrackedfile(filename) or self.hasremovedfile(filename) def trackedfiles(self): - """Returns an iterator of all filenames tracked by the dirstate.""" - return treedirstatemapiterator(self._rmap, removed=False, items=False) + """Returns a list of all filenames tracked by the dirstate.""" + trackedfiles = [] + self._rmap.visittrackedfiles(trackedfiles.append) + return iter(trackedfiles) def removedfiles(self): - """Returns an iterator of all removed files in the dirstate.""" - return treedirstatemapiterator(self._rmap, removed=True, items=False) + """Returns a list of all removed files in the dirstate.""" + removedfiles = [] + self._rmap.visitremovedfiles(removedfiles.append) + return removedfiles def __iter__(self): """Returns an iterator of all files in the dirstate.""" - return itertools.chain(self.trackedfiles(), self.removedfiles()) + trackedfiles = self.trackedfiles() + removedfiles = self.removedfiles() + if removedfiles: + return heapq.merge(iter(trackedfiles), iter(removedfiles)) + else: + return iter(trackedfiles) def keys(self): return list(iter(self))