diff --git a/rust/treedirstate/Cargo.lock b/rust/treedirstate/Cargo.lock --- a/rust/treedirstate/Cargo.lock +++ b/rust/treedirstate/Cargo.lock @@ -50,6 +50,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] +name = "cpython" +version = "0.1.0" +source = "git+https://github.com/dgrunwald/rust-cpython.git#b35031e2670d7571f03c313cde8fd91105bd5322" +dependencies = [ + "libc 0.2.33 (registry+https://github.com/rust-lang/crates.io-index)", + "num-traits 0.1.40 (registry+https://github.com/rust-lang/crates.io-index)", + "python27-sys 0.1.2 (git+https://github.com/dgrunwald/rust-cpython.git)", +] + +[[package]] name = "dbghelp-sys" version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -132,6 +142,20 @@ ] [[package]] +name = "num-traits" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" + +[[package]] +name = "python27-sys" +version = "0.1.2" +source = "git+https://github.com/dgrunwald/rust-cpython.git#b35031e2670d7571f03c313cde8fd91105bd5322" +dependencies = [ + "libc 0.2.33 (registry+https://github.com/rust-lang/crates.io-index)", + "regex 0.1.80 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] name = "quickcheck" version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -177,6 +201,7 @@ version = "0.1.0" dependencies = [ "byteorder 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)", + "cpython 0.1.0 (git+https://github.com/dgrunwald/rust-cpython.git)", "error-chain 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)", "itertools 0.7.2 (registry+https://github.com/rust-lang/crates.io-index)", "quickcheck 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)", @@ -231,6 +256,7 @@ "checksum byteorder 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "ff81738b726f5d099632ceaffe7fb65b90212e8dce59d518729e7e8634032d3d" "checksum cc 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "a9b13a57efd6b30ecd6598ebdb302cca617930b5470647570468a65d12ef9719" "checksum cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "d4c819a1287eb618df47cc647173c5c4c66ba19d888a6e50d605672aed3140de" +"checksum cpython 0.1.0 (git+https://github.com/dgrunwald/rust-cpython.git)" = "" "checksum dbghelp-sys 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "97590ba53bcb8ac28279161ca943a924d1fd4a8fb3fa63302591647c4fc5b850" "checksum either 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)" = "740178ddf48b1a9e878e6d6509a1442a2d42fd2928aae8e7a6f8a36fb01981b3" "checksum env_logger 0.3.5 (registry+https://github.com/rust-lang/crates.io-index)" = "15abd780e45b3ea4f76b4e9a26ff4843258dd8a3eed2775a0e7368c2e7936c2f" @@ -242,6 +268,8 @@ "checksum libc 0.2.33 (registry+https://github.com/rust-lang/crates.io-index)" = "5ba3df4dcb460b9dfbd070d41c94c19209620c191b0340b929ce748a2bcd42d2" "checksum log 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)" = "880f77541efa6e5cc74e76910c9884d9859683118839d6a1dc3b11e63512565b" "checksum memchr 0.1.11 (registry+https://github.com/rust-lang/crates.io-index)" = "d8b629fb514376c675b98c1421e80b151d3817ac42d7c667717d282761418d20" +"checksum num-traits 0.1.40 (registry+https://github.com/rust-lang/crates.io-index)" = "99843c856d68d8b4313b03a17e33c4bb42ae8f6610ea81b28abe076ac721b9b0" +"checksum python27-sys 0.1.2 (git+https://github.com/dgrunwald/rust-cpython.git)" = "" "checksum quickcheck 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)" = "02c2411d418cea2364325b18a205664f9ef8252e06b2e911db97c0b0d98b1406" "checksum rand 0.3.18 (registry+https://github.com/rust-lang/crates.io-index)" = "6475140dfd8655aeb72e1fd4b7a1cc1c202be65d71669476e392fe62532b9edd" "checksum regex 0.1.80 (registry+https://github.com/rust-lang/crates.io-index)" = "4fd4ace6a8cf7860714a2c2280d6c1f7e6a413486c13298bbc86fd3da019402f" diff --git a/rust/treedirstate/Cargo.toml b/rust/treedirstate/Cargo.toml --- a/rust/treedirstate/Cargo.toml +++ b/rust/treedirstate/Cargo.toml @@ -18,3 +18,9 @@ itertools = "0.7.2" quickcheck = "*" tempdir = "*" + +[dependencies.cpython] +version = "0.1" +default-features = false +features = ["python27-sys", "extension-module-2-7"] +git = "https://github.com/dgrunwald/rust-cpython.git" diff --git a/rust/treedirstate/src/lib.rs b/rust/treedirstate/src/lib.rs --- a/rust/treedirstate/src/lib.rs +++ b/rust/treedirstate/src/lib.rs @@ -14,6 +14,10 @@ extern crate byteorder; +#[cfg(not(test))] +#[macro_use] +extern crate cpython; + #[macro_use] extern crate error_chain; @@ -31,6 +35,8 @@ pub mod errors; pub mod filestate; pub mod filestore; +#[cfg(not(test))] +pub mod python; pub mod store; pub mod tree; pub mod vecmap; diff --git a/rust/treedirstate/src/python.rs b/rust/treedirstate/src/python.rs new file mode 100644 --- /dev/null +++ b/rust/treedirstate/src/python.rs @@ -0,0 +1,258 @@ +// Copyright Facebook, Inc. 2017 +//! Python bindings for treedirstate. + +use cpython::*; +use cpython::exc; +use dirstate::Dirstate; +use filestate::FileState; +use std::cell::RefCell; +use std::path::PathBuf; +use store::BlockId; + +py_module_initializer!( + rusttreedirstate, + initrusttreedirstate, + PyInit_rusttreedirstate, + |py, m| { + m.add_class::(py)?; + Ok(()) + } +); + +py_class!(class RustDirstateMap |py| { + data repodir: PathBuf; + data dirstate: RefCell>; + + def __new__( + _cls, + _ui: &PyObject, + opener: &PyObject + ) -> PyResult { + let repodir = opener.getattr(py, "base")?.extract::(py)?; + let dirstate = Dirstate::new(); + RustDirstateMap::create_instance( + py, + repodir.into(), + RefCell::new(dirstate)) + } + + def clear(&self) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + dirstate.clear(); + Ok(py.None()) + } + + // Read a dirstate file. + def read(&self, name: &str, root_id: u64) -> PyResult> { + let mut dirstate = self.dirstate(py).borrow_mut(); + dirstate + .open(self.repodir(py).join(name), BlockId(root_id)) + .map_err(|e| PyErr::new::(py, e.description()))?; + Ok(None) + } + + // Import another map of dirstate tuples into a treedirstate file. + def importmap(&self, old_map: PyObject) -> PyResult> { + let mut dirstate = self.dirstate(py).borrow_mut(); + let iter = PyIterator::from_object( + py, + old_map.call_method(py, "iteritems", NoArgs, None)?)?; + for item in iter { + let item_tuple = item?.extract::(py)?; + let filename = item_tuple.get_item(py, 0).extract::(py)?; + let data = item_tuple.get_item(py, 1).extract::(py)?; + let state = *data.get_item(py, 0)?.extract::(py)?.data(py).first().unwrap(); + let mode = data.get_item(py, 1)?.extract::(py)?; + let size = data.get_item(py, 2)?.extract::(py)?; + let mtime = data.get_item(py, 3)?.extract::(py)?; + if state == b'r' { + dirstate + .remove_file(filename.data(py), &FileState::new(b'r', 0, size, 0)) + .map_err(|e| PyErr::new::(py, e.description()))?; + } else { + dirstate + .add_file(filename.data(py), + &FileState::new(state, mode, size, mtime)) + .map_err(|e| PyErr::new::(py, e.description()))?; + } + + } + Ok(None) + } + + def write(&self, name: &str) -> PyResult> { + let mut dirstate = self.dirstate(py).borrow_mut(); + dirstate + .write_full(self.repodir(py).join(name)) + .map_err(|e| PyErr::new::(py, e.description()))?; + Ok(None) + } + + def writedelta(&self) -> PyResult> { + let mut dirstate = self.dirstate(py).borrow_mut(); + dirstate.write_delta().map_err(|e| PyErr::new::(py, e.description()))?; + Ok(None) + } + + def filecount(&self) -> PyResult { + let dirstate = self.dirstate(py).borrow(); + Ok((dirstate.tracked_count() + dirstate.removed_count()) as usize) + } + + def hastrackedfile(&self, filename: PyBytes) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let value = dirstate + .get_tracked(filename.data(py)) + .map_err(|e| PyErr::new::(py, e.description()))?; + + Ok(value.is_some()) + } + + def hasremovedfile(&self, filename: PyBytes) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let value = dirstate + .get_removed(filename.data(py)) + .map_err(|e| PyErr::new::(py, e.description()))?; + Ok(value.is_some()) + } + + def gettracked(&self, filename: PyObject, default: PyObject) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let res = if let Ok(filename) = filename.extract::(py) { + let value = dirstate + .get_tracked(filename.data(py)) + .map_err(|e| PyErr::new::(py, e.description()))?; + match value { + Some(ref file) => + PyTuple::new(py, &[ + PyBytes::new(py, &[file.state; 1]).to_py_object(py).into_object(), + file.mode.to_py_object(py).into_object(), + file.size.to_py_object(py).into_object(), + file.mtime.to_py_object(py).into_object()]).into_object(), + None => default, + } + } else { + default + }; + Ok(res) + } + + def getremoved(&self, filename: PyObject, default: PyObject) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let res = if let Ok(filename) = filename.extract::(py) { + let value = dirstate + .get_removed(filename.data(py)) + .map_err(|e| PyErr::new::(py, e.description()))?; + match value { + Some(ref file) => + PyTuple::new(py, &[ + PyBytes::new(py, &[file.state; 1]).to_py_object(py).into_object(), + file.mode.to_py_object(py).into_object(), + file.size.to_py_object(py).into_object(), + file.mtime.to_py_object(py).into_object()]).into_object(), + None => default, + } + } else { + default + }; + Ok(res) + } + + def hastrackeddir(&self, dirname: PyBytes) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + dirstate + .has_tracked_dir(dirname.data(py)) + .map_err(|e| PyErr::new::(py, e.description())) + } + + def hasremoveddir(&self, dirname: PyBytes) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + dirstate + .has_removed_dir(dirname.data(py)) + .map_err(|e| PyErr::new::(py, e.description())) + } + + // 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. + def getnext(&self, filename: Option, removed: bool) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let next = if removed { + match filename { + Some(filename) => { + dirstate + .get_next_removed(filename.data(py)) + .map_err(|e| PyErr::new::(py, e.description()))? + } + None => { + dirstate + .get_first_removed() + .map_err(|e| PyErr::new::(py, e.description()))? + } + } + } else { + match filename { + Some(filename) => { + dirstate + .get_next_tracked(filename.data(py)) + .map_err(|e| PyErr::new::(py, e.description()))? + } + None => { + dirstate.get_first_tracked() + .map_err(|e| PyErr::new::(py, e.description()))? + } + } + }; + let res = match next { + Some((ref f, ref s)) => + PyTuple::new(py, &[ + PyBytes::new(py, &f).into_object(), + PyTuple::new(py, &[ + PyBytes::new(py, &[s.state; 1]).to_py_object(py).into_object(), + s.mode.to_py_object(py).into_object(), + s.size.to_py_object(py).into_object(), + s.mtime.to_py_object(py).into_object()]).into_object() + ]).into_object(), + None => py.None(), + }; + Ok(res) + } + + def addfile( + &self, + filename: PyBytes, + _old_state: PyBytes, + state: PyBytes, + mode: u32, + size: i32, + mtime: i32 + ) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let state = *state.data(py).first().unwrap_or(&b'?'); + dirstate + .add_file(filename.data(py), &FileState::new(state, mode, size, mtime)) + .map_err(|e| PyErr::new::(py, e.description()))?; + Ok(py.None()) + } + + def removefile(&self, filename: PyBytes, _old_state: PyBytes, size: i32) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + dirstate + .remove_file(filename.data(py), &FileState::new(b'r', 0, size, 0)) + .map_err(|e| PyErr::new::(py, e.description()))?; + Ok(py.None()) + } + + def dropfile(&self, filename: PyBytes) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + dirstate + .drop_file(filename.data(py)) + .map_err(|e| PyErr::new::(py, e.description())) + } + + // Returns the ID of the root node. + def rootid(&self) -> PyResult> { + Ok(self.dirstate(py).borrow().root_id().map(|id| id.0)) + } + +}); diff --git a/treedirstate/__init__.py b/treedirstate/__init__.py new file mode 100644 --- /dev/null +++ b/treedirstate/__init__.py @@ -0,0 +1,451 @@ +# Copyright Facebook, Inc. 2017 +# +# This software may be used and distributed according to the terms of the +# GNU General Public License version 2 or any later version. +"""tree-based dirstate implementation""" + +from __future__ import absolute_import +from mercurial import ( + dirstate, + error, + extensions, + localrepo, + node, + pycompat, + registrar, + txnutil, + util, +) +from mercurial.i18n import _ +import errno +import itertools +import struct + +from .rusttreedirstate import RustDirstateMap + +dirstateheader = b'########################treedirstate####' +treedirstateversion = 1 +useinnewrepos = True + +class _reader(object): + def __init__(self, data, offset): + self.data = data + self.offset = offset + + def readuint(self): + v = struct.unpack(">L", self.data[self.offset:self.offset + 4]) + self.offset += 4 + return v[0] + + def readstr(self): + l = self.readuint() + v = self.data[self.offset:self.offset + l] + self.offset += l + return v + +class _writer(object): + def __init__(self): + self.buffer = pycompat.stringio() + + def writeuint(self, v): + self.buffer.write(struct.pack(">L", v)) + + def writestr(self, v): + self.writeuint(len(v)) + self.buffer.write(v) + +# The treedirstatemap iterator uses the getnext method on the dirstatemap +# to find the next item on each call. This involves searching down the +# tree each time. A future improvement is to keep the state between each +# call to avoid these extra searches. +class treedirstatemapiterator(object): + def __init__(self, map_, removed=False, items=False): + self._rmap = map_ + self._removed = removed + self._items = items + self._at = None + + def __iter__(self): + return self + + def __next__(self): + nextitem = self._rmap.getnext(self._at, self._removed) + if nextitem is None: + raise StopIteration + self._at = nextitem[0] + return nextitem if self._items else nextitem[0] + + def next(self): + return self.__next__() + +class treedirstatemap(object): + def __init__(self, ui, opener, root, importmap=None): + self._ui = ui + self._opener = opener + self._root = root + self.copymap = {} + + self._filename = 'dirstate' + self._rmap = RustDirstateMap(ui, opener) + self._treeid = None + self._parents = None + self._dirtyparents = False + self._nonnormalset = set() + self._otherparentset = set() + + if importmap is not None: + self._rmap.importmap(importmap) + self._parents = importmap._parents + self._nonnormalset = importmap.nonnormalset + self._otherparentset = importmap.otherparentset + self.copymap = importmap.copymap + else: + self.read() + + def preload(self): + pass + + def clear(self): + self._rmap.clear() + self.copymap.clear() + self._nonnormalset.clear() + self._otherparentset.clear() + self.setparents(node.nullid, node.nullid) + + def __len__(self): + """Returns the number of files, including removed files.""" + return self._rmap.filecount() + + def itertrackeditems(self): + """Returns an iterator over (filename, (state, mode, size, mtime)).""" + return treedirstatemapiterator(self._rmap, removed=False, items=True) + + 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) + + def iteritems(self): + return itertools.chain(self.itertrackeditems(), + self.iterremoveditems()) + + def gettracked(self, filename, default=None): + """Returns (state, mode, size, mtime) for the tracked file.""" + return self._rmap.gettracked(filename, default) + + def getremoved(self, filename, default=None): + """Returns (state, mode, size, mtime) for the removed file.""" + return self._rmap.getremoved(filename, default) + + def get(self, filename, default=None): + return (self._rmap.gettracked(filename, None) or + self._rmap.getremoved(filename, default)) + + def __getitem__(self, filename): + item = (self._rmap.gettracked(filename, None) or + self._rmap.getremoved(filename, None)) + if item is None: + raise KeyError(filename) + return item + + def hastrackedfile(self, filename): + """Returns true if the file is tracked in the dirstate.""" + return self._rmap.hastrackedfile(filename) + + def hasremovedfile(self, filename): + """Returns true if the file is recorded as removed in the dirstate.""" + return self._rmap.hasremovedfile(filename) + + def __contains__(self, filename): + 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) + + def removedfiles(self): + """Returns an iterator of all removed files in the dirstate.""" + return treedirstatemapiterator(self._rmap, removed=True, items=False) + + def __iter__(self): + """Returns an iterator of all files in the dirstate.""" + return itertools.chain(self.trackedfiles(), self.removedfiles()) + + def keys(self): + return list(iter(self)) + + def hastrackeddir(self, dirname): + """ + Returns True if the dirstate includes a directory. + """ + return self._rmap.hastrackeddir(dirname + '/') + + def hasremoveddir(self, dirname): + """ + Returns True if the directories containing files marked for removal + includes a directory. + """ + return self._rmap.hasremoveddir(dirname + '/') + + def hasdir(self, dirname): + """ + Returns True if the directory exists in the dirstate for either + tracked or removed files. + """ + return self.hastrackeddir(dirname) or self.hasremoveddir(dirname) + + def addfile(self, f, oldstate, state, mode, size, mtime): + self._rmap.addfile(f, oldstate, state, mode, size, mtime) + if state != 'n' or mtime == -1: + self._nonnormalset.add(f) + else: + self._nonnormalset.discard(f) + if size == -2: + self._otherparentset.add(f) + else: + self._otherparentset.discard(f) + + def removefile(self, f, oldstate, size): + self._rmap.removefile(f, oldstate, size) + self._nonnormalset.add(f) + if size == -2: + self._otherparentset.add(f) + + def dropfile(self, f, oldstate): + """ + Drops a file from the dirstate. Returns True if it was previously + recorded. + """ + self._nonnormalset.discard(f) + self._otherparentset.discard(f) + return self._rmap.dropfile(f) + + def clearambiguoustimes(self, files, now): + ### TODO + pass + + def parents(self): + """ + Returns the parents of the dirstate. + """ + return self._parents + + def setparents(self, p1, p2): + """ + Sets the dirstate parents. + """ + self._parents = (p1, p2) + self._dirtyparents = True + + @property + def nonnormalset(self): + return self._nonnormalset + + @property + def otherparentset(self): + return self._otherparentset + + def getfilefoldmap(self): + """Returns a dictionary mapping normalized case paths to their + non-normalized versions. + """ + raise NotImplementedError() + + def getdirfoldmap(self): + """ + Returns a dictionary mapping normalized case paths to their + non-normalized versions for directories. + """ + raise NotImplementedError() + + def identity(self): + if self._identity is None: + self.read() + return self._identity + + def _opendirstatefile(self): + fp, _mode = txnutil.trypending(self._root, self._opener, self._filename) + return fp + + def read(self): + # ignore HG_PENDING because identity is used only for writing + self._identity = util.filestat.frompath( + self._opener.join(self._filename)) + + try: + data = self._opendirstatefile().read() + except IOError as err: + if err.errno != errno.ENOENT: + raise + # File doesn't exist so current state is empty. + if not self._dirtyparents: + self._parents = (node.nullid, node.nullid) + return + + if data[40:80] != dirstateheader: + raise error.Abort(_('dirstate is not a valid treedirstate')) + + if not self._dirtyparents: + self._parents = data[:20], data[20:40] + + r = _reader(data, 80) + + version = r.readuint() + if version != treedirstateversion: + raise error.Abort(_('unsupported treedirstate version: %s') + % version) + + self._treeid = r.readstr() + rootid = r.readuint() + self._rmap.read('dirstate.tree.000', rootid) + clen = r.readuint() + copymap = {} + for _i in range(clen): + k = r.readstr() + v = r.readstr() + copymap[k] = v + + def readset(): + slen = r.readuint() + s = set() + for _i in range(slen): + s.add(r.readstr()) + return s + + nonnormalset = readset() + otherparentset = readset() + + self.copymap = copymap + self._nonnormalset = nonnormalset + self._otherparentset = otherparentset + + def startwrite(self, tr): + # TODO: register map store offset with 'tr' + pass + + def write(self, st, now): + """Write the dirstate to the filehandle st.""" + if self._treeid is None: + self._treeid = '000' + self._rmap.write('dirstate.tree.000') + else: + self._rmap.writedelta() + st.write(self._genrootdata()) + st.close() + self._dirtyparents = False + + def writeflat(self): + with self._opener("dirstate", "w", + atomictemp=True, checkambig=True) as st: + newdmap = {} + for k, v in self.iteritems(): + newdmap[k] = dirstate.dirstatetuple(*v) + + st.write(dirstate.parsers.pack_dirstate( + newdmap, self.copymap, self._parents, + dirstate._getfsnow(self._opener))) + + def _genrootdata(self): + w = _writer() + if self._parents: + w.buffer.write(self._parents[0]) + w.buffer.write(self._parents[1]) + else: + w.buffer.write(node.nullid) + w.buffer.write(node.nullid) + w.buffer.write(dirstateheader) + w.writeuint(treedirstateversion) + w.writestr(self._treeid) + w.writeuint(self._rmap.rootid()) + w.writeuint(len(self.copymap)) + for k, v in self.copymap.iteritems(): + w.writestr(k) + w.writestr(v) + w.writeuint(len(self._nonnormalset)) + for v in self._nonnormalset: + w.writestr(v) + w.writeuint(len(self._otherparentset)) + for v in self._otherparentset: + w.writestr(v) + + return w.buffer.getvalue() + +def istreedirstate(repo): + return (util.safehasattr(repo, 'requirements') and + 'treedirstate' in repo.requirements) + +def upgrade(ui, repo): + if istreedirstate(repo): + raise error.Abort('repo already has treedirstate') + with repo.wlock(): + newmap = treedirstatemap(ui, repo.dirstate._opener, repo.root, + importmap=repo.dirstate._map) + f = repo.dirstate._opener('dirstate', 'w') + newmap.write(f, dirstate._getfsnow(repo.dirstate._opener)) + repo.requirements.add('treedirstate') + repo._writerequirements() + del repo.dirstate + +def downgrade(ui, repo): + if not istreedirstate(repo): + raise error.Abort('repo doesn\'t have treedirstate') + with repo.wlock(): + repo.dirstate._map.writeflat() + repo.requirements.remove('treedirstate') + repo._writerequirements() + del repo.dirstate + +def wrapdirstate(orig, self): + ds = orig(self) + if istreedirstate(self): + ds._mapcls = treedirstatemap + return ds + +def wrapnewreporequirements(orig, repo): + reqs = orig(repo) + if useinnewrepos: + reqs.add('treedirstate') + return reqs + +def featuresetup(ui, supported): + supported |= {'treedirstate'} + +def extsetup(ui): + # Check this version of Mercurial has the extension points we need + if not util.safehasattr(dirstate.dirstatemap, "hasdir"): + ui.warn(_("this version of Mercurial doesn't support treedirstate\n")) + return + + if util.safehasattr(localrepo, 'newreporequirements'): + extensions.wrapfunction(localrepo, 'newreporequirements', + wrapnewreporequirements) + + localrepo.localrepository.featuresetupfuncs.add(featuresetup) + extensions.wrapfunction(localrepo.localrepository.dirstate, 'func', + wrapdirstate) + +def reposetup(ui, repo): + ui.setconfig('treedirstate', 'enabled', istreedirstate(repo)) + +# debug commands +cmdtable = {} +command = registrar.command(cmdtable) + +@command('debugtreedirstate', [], 'hg debugtreedirstate [on|off|status]') +def debugtreedirstate(ui, repo, cmd, **opts): + """migrate to treedirstate""" + if cmd == "on": + upgrade(ui, repo) + elif cmd == "off": + downgrade(ui, repo) + elif cmd == "status": + if istreedirstate(repo): + ui.status(_("treedirstate enabled " + + "(using dirstate.tree.%s, %s files tracked)") + % (repo.dirstate._map._treeid, len(repo.dirstate._map))) + else: + ui.status(_("treedirstate not enabled")) + else: + raise error.Abort("unrecognised command: %s" % cmd)