dirstatemap: move `set_tracked` out of common methods and plug in Rust

We now have a Rust-specific implementation of this method, it is no longer
shared between both implementations.

Differential Revision: https://phab.mercurial-scm.org/D12496

mercurial/dirstatemap.py

@@ -1,730 +1,733 b''
 # dirstatemap.py
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 
 import errno
 
 from .i18n import _
 
 from . import (
     error,
     pathutil,
     policy,
     txnutil,
     util,
 )
 
 from .dirstateutils import (
     docket as docketmod,
     v2,
 )
 
 parsers = policy.importmod('parsers')
 rustmod = policy.importrust('dirstate')
 
 propertycache = util.propertycache
 
 if rustmod is None:
     DirstateItem = parsers.DirstateItem
 else:
     DirstateItem = rustmod.DirstateItem
 
 rangemask = 0x7FFFFFFF
 
 
 class _dirstatemapcommon:
     """
     Methods that are identical for both implementations of the dirstatemap
     class, with and without Rust extensions enabled.
     """
 
     # please pytype
 
     _map = None
     copymap = None
 
     def __init__(self, ui, opener, root, nodeconstants, use_dirstate_v2):
         self._use_dirstate_v2 = use_dirstate_v2
         self._nodeconstants = nodeconstants
         self._ui = ui
         self._opener = opener
         self._root = root
         self._filename = b'dirstate'
         self._nodelen = 20  # Also update Rust code when changing this!
         self._parents = None
         self._dirtyparents = False
         self._docket = None
 
         # for consistent view between _pl() and _read() invocations
         self._pendingmode = None
 
     def preload(self):
         """Loads the underlying data, if it's not already loaded"""
         self._map
 
     def get(self, key, default=None):
         return self._map.get(key, default)
 
     def __len__(self):
         return len(self._map)
 
     def __iter__(self):
         return iter(self._map)
 
     def __contains__(self, key):
         return key in self._map
 
     def __getitem__(self, item):
         return self._map[item]
 
     ### sub-class utility method
     #
     # Use to allow for generic implementation of some method while still coping
     # with minor difference between implementation.
 
     def _dirs_incr(self, filename, old_entry=None):
         """increment the dirstate counter if applicable
 
         This might be a no-op for some subclasses who deal with directory
         tracking in a different way.
         """
 
     def _dirs_decr(self, filename, old_entry=None, remove_variant=False):
         """decrement the dirstate counter if applicable
 
         This might be a no-op for some subclasses who deal with directory
         tracking in a different way.
         """
 
     def _refresh_entry(self, f, entry):
         """record updated state of an entry"""
 
     def _insert_entry(self, f, entry):
         """add a new dirstate entry (or replace an unrelated one)
 
         The fact it is actually new is the responsability of the caller
         """
 
     def _drop_entry(self, f):
         """remove any entry for file f
 
         This should also drop associated copy information
 
         The fact we actually need to drop it is the responsability of the caller"""
 
     ### method to manipulate the entries
 
     def set_possibly_dirty(self, filename):
         """record that the current state of the file on disk is unknown"""
         entry = self[filename]
         entry.set_possibly_dirty()
         self._refresh_entry(filename, entry)
 
     def set_clean(self, filename, mode, size, mtime):
         """mark a file as back to a clean state"""
         entry = self[filename]
         size = size & rangemask
         entry.set_clean(mode, size, mtime)
         self._refresh_entry(filename, entry)
         self.copymap.pop(filename, None)
 
-    def set_tracked(self, filename):
-        new = False
-        entry = self.get(filename)
-        if entry is None:
-            self._dirs_incr(filename)
-            entry = DirstateItem(
-                wc_tracked=True,
-            )
-
-            self._insert_entry(filename, entry)
-            new = True
-        elif not entry.tracked:
-            self._dirs_incr(filename, entry)
-            entry.set_tracked()
-            self._refresh_entry(filename, entry)
-            new = True
-        else:
-            # XXX This is probably overkill for more case, but we need this to
-            # fully replace the `normallookup` call with `set_tracked` one.
-            # Consider smoothing this in the future.
-            entry.set_possibly_dirty()
-            self._refresh_entry(filename, entry)
-        return new
-
     def set_untracked(self, f):
         """Mark a file as no longer tracked in the dirstate map"""
         entry = self.get(f)
         if entry is None:
             return False
         else:
             self._dirs_decr(f, old_entry=entry, remove_variant=not entry.added)
             if not entry.p2_info:
                 self.copymap.pop(f, None)
             entry.set_untracked()
             self._refresh_entry(f, entry)
             return True
 
     def reset_state(
         self,
         filename,
         wc_tracked=False,
         p1_tracked=False,
         p2_info=False,
         has_meaningful_mtime=True,
         has_meaningful_data=True,
         parentfiledata=None,
     ):
         """Set a entry to a given state, diregarding all previous state
 
         This is to be used by the part of the dirstate API dedicated to
         adjusting the dirstate after a update/merge.
 
         note: calling this might result to no entry existing at all if the
         dirstate map does not see any point at having one for this file
         anymore.
         """
         # copy information are now outdated
         # (maybe new information should be in directly passed to this function)
         self.copymap.pop(filename, None)
 
         if not (p1_tracked or p2_info or wc_tracked):
             old_entry = self._map.get(filename)
             self._drop_entry(filename)
             self._dirs_decr(filename, old_entry=old_entry)
             return
 
         old_entry = self._map.get(filename)
         self._dirs_incr(filename, old_entry)
         entry = DirstateItem(
             wc_tracked=wc_tracked,
             p1_tracked=p1_tracked,
             p2_info=p2_info,
             has_meaningful_mtime=has_meaningful_mtime,
             parentfiledata=parentfiledata,
         )
         self._insert_entry(filename, entry)
 
     ### disk interaction
 
     def _opendirstatefile(self):
         fp, mode = txnutil.trypending(self._root, self._opener, self._filename)
         if self._pendingmode is not None and self._pendingmode != mode:
             fp.close()
             raise error.Abort(
                 _(b'working directory state may be changed parallelly')
             )
         self._pendingmode = mode
         return fp
 
     def _readdirstatefile(self, size=-1):
         try:
             with self._opendirstatefile() as fp:
                 return fp.read(size)
         except IOError as err:
             if err.errno != errno.ENOENT:
                 raise
             # File doesn't exist, so the current state is empty
             return b''
 
     @property
     def docket(self):
         if not self._docket:
             if not self._use_dirstate_v2:
                 raise error.ProgrammingError(
                     b'dirstate only has a docket in v2 format'
                 )
             self._docket = docketmod.DirstateDocket.parse(
                 self._readdirstatefile(), self._nodeconstants
             )
         return self._docket
 
     def write_v2_no_append(self, tr, st, meta, packed):
         old_docket = self.docket
         new_docket = docketmod.DirstateDocket.with_new_uuid(
             self.parents(), len(packed), meta
         )
         data_filename = new_docket.data_filename()
         if tr:
             tr.add(data_filename, 0)
         self._opener.write(data_filename, packed)
         # Write the new docket after the new data file has been
         # written. Because `st` was opened with `atomictemp=True`,
         # the actual `.hg/dirstate` file is only affected on close.
         st.write(new_docket.serialize())
         st.close()
         # Remove the old data file after the new docket pointing to
         # the new data file was written.
         if old_docket.uuid:
             data_filename = old_docket.data_filename()
             unlink = lambda _tr=None: self._opener.unlink(data_filename)
             if tr:
                 category = b"dirstate-v2-clean-" + old_docket.uuid
                 tr.addpostclose(category, unlink)
             else:
                 unlink()
         self._docket = new_docket
 
     ### reading/setting parents
 
     def parents(self):
         if not self._parents:
             if self._use_dirstate_v2:
                 self._parents = self.docket.parents
             else:
                 read_len = self._nodelen * 2
                 st = self._readdirstatefile(read_len)
                 l = len(st)
                 if l == read_len:
                     self._parents = (
                         st[: self._nodelen],
                         st[self._nodelen : 2 * self._nodelen],
                     )
                 elif l == 0:
                     self._parents = (
                         self._nodeconstants.nullid,
                         self._nodeconstants.nullid,
                     )
                 else:
                     raise error.Abort(
                         _(b'working directory state appears damaged!')
                     )
 
         return self._parents
 
 
 class dirstatemap(_dirstatemapcommon):
     """Map encapsulating the dirstate's contents.
 
     The dirstate contains the following state:
 
     - `identity` is the identity of the dirstate file, which can be used to
       detect when changes have occurred to the dirstate file.
 
     - `parents` is a pair containing the parents of the working copy. The
       parents are updated by calling `setparents`.
 
     - the state map maps filenames to tuples of (state, mode, size, mtime),
       where state is a single character representing 'normal', 'added',
       'removed', or 'merged'. It is read by treating the dirstate as a
       dict.  File state is updated by calling various methods (see each
       documentation for details):
 
       - `reset_state`,
       - `set_tracked`
       - `set_untracked`
       - `set_clean`
       - `set_possibly_dirty`
 
     - `copymap` maps destination filenames to their source filename.
 
     The dirstate also provides the following views onto the state:
 
     - `filefoldmap` is a dict mapping normalized filenames to the denormalized
       form that they appear as in the dirstate.
 
     - `dirfoldmap` is a dict mapping normalized directory names to the
       denormalized form that they appear as in the dirstate.
     """
 
     ### Core data storage and access
 
     @propertycache
     def _map(self):
         self._map = {}
         self.read()
         return self._map
 
     @propertycache
     def copymap(self):
         self.copymap = {}
         self._map
         return self.copymap
 
     def clear(self):
         self._map.clear()
         self.copymap.clear()
         self.setparents(self._nodeconstants.nullid, self._nodeconstants.nullid)
         util.clearcachedproperty(self, b"_dirs")
         util.clearcachedproperty(self, b"_alldirs")
         util.clearcachedproperty(self, b"filefoldmap")
         util.clearcachedproperty(self, b"dirfoldmap")
 
     def items(self):
         return self._map.items()
 
     # forward for python2,3 compat
     iteritems = items
 
     def debug_iter(self, all):
         """
         Return an iterator of (filename, state, mode, size, mtime) tuples
 
         `all` is unused when Rust is not enabled
         """
         for (filename, item) in self.items():
             yield (filename, item.state, item.mode, item.size, item.mtime)
 
     def keys(self):
         return self._map.keys()
 
     ### reading/setting parents
 
     def setparents(self, p1, p2, fold_p2=False):
         self._parents = (p1, p2)
         self._dirtyparents = True
         copies = {}
         if fold_p2:
             for f, s in self._map.items():
                 # Discard "merged" markers when moving away from a merge state
                 if s.p2_info:
                     source = self.copymap.pop(f, None)
                     if source:
                         copies[f] = source
                     s.drop_merge_data()
         return copies
 
     ### disk interaction
 
     def read(self):
         # ignore HG_PENDING because identity is used only for writing
         self.identity = util.filestat.frompath(
             self._opener.join(self._filename)
         )
 
         if self._use_dirstate_v2:
             if not self.docket.uuid:
                 return
             st = self._opener.read(self.docket.data_filename())
         else:
             st = self._readdirstatefile()
 
         if not st:
             return
 
         # TODO: adjust this estimate for dirstate-v2
         if util.safehasattr(parsers, b'dict_new_presized'):
             # Make an estimate of the number of files in the dirstate based on
             # its size. This trades wasting some memory for avoiding costly
             # resizes. Each entry have a prefix of 17 bytes followed by one or
             # two path names. Studies on various large-scale real-world repositories
             # found 54 bytes a reasonable upper limit for the average path names.
             # Copy entries are ignored for the sake of this estimate.
             self._map = parsers.dict_new_presized(len(st) // 71)
 
         # Python's garbage collector triggers a GC each time a certain number
         # of container objects (the number being defined by
         # gc.get_threshold()) are allocated. parse_dirstate creates a tuple
         # for each file in the dirstate. The C version then immediately marks
         # them as not to be tracked by the collector. However, this has no
         # effect on when GCs are triggered, only on what objects the GC looks
         # into. This means that O(number of files) GCs are unavoidable.
         # Depending on when in the process's lifetime the dirstate is parsed,
         # this can get very expensive. As a workaround, disable GC while
         # parsing the dirstate.
         #
         # (we cannot decorate the function directly since it is in a C module)
         if self._use_dirstate_v2:
             p = self.docket.parents
             meta = self.docket.tree_metadata
             parse_dirstate = util.nogc(v2.parse_dirstate)
             parse_dirstate(self._map, self.copymap, st, meta)
         else:
             parse_dirstate = util.nogc(parsers.parse_dirstate)
             p = parse_dirstate(self._map, self.copymap, st)
         if not self._dirtyparents:
             self.setparents(*p)
 
         # Avoid excess attribute lookups by fast pathing certain checks
         self.__contains__ = self._map.__contains__
         self.__getitem__ = self._map.__getitem__
         self.get = self._map.get
 
     def write(self, tr, st):
         if self._use_dirstate_v2:
             packed, meta = v2.pack_dirstate(self._map, self.copymap)
             self.write_v2_no_append(tr, st, meta, packed)
         else:
             packed = parsers.pack_dirstate(
                 self._map, self.copymap, self.parents()
             )
             st.write(packed)
             st.close()
         self._dirtyparents = False
 
     @propertycache
     def identity(self):
         self._map
         return self.identity
 
     ### code related to maintaining and accessing "extra" property
     # (e.g. "has_dir")
 
     def _dirs_incr(self, filename, old_entry=None):
         """incremente the dirstate counter if applicable"""
         if (
             old_entry is None or old_entry.removed
         ) and "_dirs" in self.__dict__:
             self._dirs.addpath(filename)
         if old_entry is None and "_alldirs" in self.__dict__:
             self._alldirs.addpath(filename)
 
     def _dirs_decr(self, filename, old_entry=None, remove_variant=False):
         """decremente the dirstate counter if applicable"""
         if old_entry is not None:
             if "_dirs" in self.__dict__ and not old_entry.removed:
                 self._dirs.delpath(filename)
             if "_alldirs" in self.__dict__ and not remove_variant:
                 self._alldirs.delpath(filename)
         elif remove_variant and "_alldirs" in self.__dict__:
             self._alldirs.addpath(filename)
         if "filefoldmap" in self.__dict__:
             normed = util.normcase(filename)
             self.filefoldmap.pop(normed, None)
 
     @propertycache
     def filefoldmap(self):
         """Returns a dictionary mapping normalized case paths to their
         non-normalized versions.
         """
         try:
             makefilefoldmap = parsers.make_file_foldmap
         except AttributeError:
             pass
         else:
             return makefilefoldmap(
                 self._map, util.normcasespec, util.normcasefallback
             )
 
         f = {}
         normcase = util.normcase
         for name, s in self._map.items():
             if not s.removed:
                 f[normcase(name)] = name
         f[b'.'] = b'.'  # prevents useless util.fspath() invocation
         return f
 
     @propertycache
     def dirfoldmap(self):
         f = {}
         normcase = util.normcase
         for name in self._dirs:
             f[normcase(name)] = name
         return f
 
     def hastrackeddir(self, d):
         """
         Returns True if the dirstate contains a tracked (not removed) file
         in this directory.
         """
         return d in self._dirs
 
     def hasdir(self, d):
         """
         Returns True if the dirstate contains a file (tracked or removed)
         in this directory.
         """
         return d in self._alldirs
 
     @propertycache
     def _dirs(self):
         return pathutil.dirs(self._map, only_tracked=True)
 
     @propertycache
     def _alldirs(self):
         return pathutil.dirs(self._map)
 
     ### code related to manipulation of entries and copy-sources
 
+    def set_tracked(self, filename):
+        new = False
+        entry = self.get(filename)
+        if entry is None:
+            self._dirs_incr(filename)
+            entry = DirstateItem(
+                wc_tracked=True,
+            )
+
+            self._insert_entry(filename, entry)
+            new = True
+        elif not entry.tracked:
+            self._dirs_incr(filename, entry)
+            entry.set_tracked()
+            self._refresh_entry(filename, entry)
+            new = True
+        else:
+            # XXX This is probably overkill for more case, but we need this to
+            # fully replace the `normallookup` call with `set_tracked` one.
+            # Consider smoothing this in the future.
+            entry.set_possibly_dirty()
+            self._refresh_entry(filename, entry)
+        return new
+
     def _refresh_entry(self, f, entry):
         if not entry.any_tracked:
             self._map.pop(f, None)
 
     def _insert_entry(self, f, entry):
         self._map[f] = entry
 
     def _drop_entry(self, f):
         self._map.pop(f, None)
         self.copymap.pop(f, None)
 
 
 if rustmod is not None:
 
     class dirstatemap(_dirstatemapcommon):
 
         ### Core data storage and access
 
         @propertycache
         def _map(self):
             """
             Fills the Dirstatemap when called.
             """
             # ignore HG_PENDING because identity is used only for writing
             self.identity = util.filestat.frompath(
                 self._opener.join(self._filename)
             )
 
             if self._use_dirstate_v2:
                 if self.docket.uuid:
                     # TODO: use mmap when possible
                     data = self._opener.read(self.docket.data_filename())
                 else:
                     data = b''
                 self._map = rustmod.DirstateMap.new_v2(
                     data, self.docket.data_size, self.docket.tree_metadata
                 )
                 parents = self.docket.parents
             else:
                 self._map, parents = rustmod.DirstateMap.new_v1(
                     self._readdirstatefile()
                 )
 
             if parents and not self._dirtyparents:
                 self.setparents(*parents)
 
             self.__contains__ = self._map.__contains__
             self.__getitem__ = self._map.__getitem__
             self.get = self._map.get
             return self._map
 
         @property
         def copymap(self):
             return self._map.copymap()
 
         def debug_iter(self, all):
             """
             Return an iterator of (filename, state, mode, size, mtime) tuples
 
             `all`: also include with `state == b' '` dirstate tree nodes that
             don't have an associated `DirstateItem`.
 
             """
             return self._map.debug_iter(all)
 
         def clear(self):
             self._map.clear()
             self.setparents(
                 self._nodeconstants.nullid, self._nodeconstants.nullid
             )
             util.clearcachedproperty(self, b"_dirs")
             util.clearcachedproperty(self, b"_alldirs")
             util.clearcachedproperty(self, b"dirfoldmap")
 
         def items(self):
             return self._map.items()
 
         # forward for python2,3 compat
         iteritems = items
 
         def keys(self):
             return iter(self._map)
 
         ### reading/setting parents
 
         def setparents(self, p1, p2, fold_p2=False):
             self._parents = (p1, p2)
             self._dirtyparents = True
             copies = {}
             if fold_p2:
                 # Collect into an intermediate list to avoid a `RuntimeError`
                 # exception due to mutation during iteration.
                 # TODO: move this the whole loop to Rust where `iter_mut`
                 # enables in-place mutation of elements of a collection while
                 # iterating it, without mutating the collection itself.
                 files_with_p2_info = [
                     f for f, s in self._map.items() if s.p2_info
                 ]
                 rust_map = self._map
                 for f in files_with_p2_info:
                     e = rust_map.get(f)
                     source = self.copymap.pop(f, None)
                     if source:
                         copies[f] = source
                     e.drop_merge_data()
                     rust_map.set_dirstate_item(f, e)
             return copies
 
         ### disk interaction
 
         @propertycache
         def identity(self):
             self._map
             return self.identity
 
         def write(self, tr, st):
             if not self._use_dirstate_v2:
                 p1, p2 = self.parents()
                 packed = self._map.write_v1(p1, p2)
                 st.write(packed)
                 st.close()
                 self._dirtyparents = False
                 return
 
             # We can only append to an existing data file if there is one
             can_append = self.docket.uuid is not None
             packed, meta, append = self._map.write_v2(can_append)
             if append:
                 docket = self.docket
                 data_filename = docket.data_filename()
                 if tr:
                     tr.add(data_filename, docket.data_size)
                 with self._opener(data_filename, b'r+b') as fp:
                     fp.seek(docket.data_size)
                     assert fp.tell() == docket.data_size
                     written = fp.write(packed)
                     if written is not None:  # py2 may return None
                         assert written == len(packed), (written, len(packed))
                 docket.data_size += len(packed)
                 docket.parents = self.parents()
                 docket.tree_metadata = meta
                 st.write(docket.serialize())
                 st.close()
             else:
                 self.write_v2_no_append(tr, st, meta, packed)
             # Reload from the newly-written file
             util.clearcachedproperty(self, b"_map")
             self._dirtyparents = False
 
         ### code related to maintaining and accessing "extra" property
         # (e.g. "has_dir")
 
         @propertycache
         def filefoldmap(self):
             """Returns a dictionary mapping normalized case paths to their
             non-normalized versions.
             """
             return self._map.filefoldmapasdict()
 
         def hastrackeddir(self, d):
             return self._map.hastrackeddir(d)
 
         def hasdir(self, d):
             return self._map.hasdir(d)
 
         @propertycache
         def dirfoldmap(self):
             f = {}
             normcase = util.normcase
             for name in self._map.tracked_dirs():
                 f[normcase(name)] = name
             return f
 
         ### code related to manipulation of entries and copy-sources
 
         def _refresh_entry(self, f, entry):
             if not entry.any_tracked:
                 self._map.drop_item_and_copy_source(f)
             else:
                 self._map.addfile(f, entry)
 
         def _insert_entry(self, f, entry):
             self._map.addfile(f, entry)
 
+        def set_tracked(self, f):
+            return self._map.set_tracked(f)
+
         def _drop_entry(self, f):
             self._map.drop_item_and_copy_source(f)
 
         def __setitem__(self, key, value):
             assert isinstance(value, DirstateItem)
             self._map.set_dirstate_item(key, value)

rust/hg-cpython/src/dirstate/dirstate_map.rs

@@ -1,490 +1,500 b''
 // dirstate_map.rs
 //
 // Copyright 2019 Raphaël Gomès <[e-mail]>
 //
 // This software may be used and distributed according to the terms of the
 // GNU General Public License version 2 or any later version.
 
 //! Bindings for the `hg::dirstate::dirstate_map` file provided by the
 //! `hg-core` package.
 
 use std::cell::{RefCell, RefMut};
 use std::convert::TryInto;
 
 use cpython::{
     exc, PyBool, PyBytes, PyClone, PyDict, PyErr, PyList, PyNone, PyObject,
     PyResult, Python, PythonObject, ToPyObject, UnsafePyLeaked,
 };
 
 use crate::{
     dirstate::copymap::{CopyMap, CopyMapItemsIterator, CopyMapKeysIterator},
     dirstate::item::DirstateItem,
     pybytes_deref::PyBytesDeref,
 };
 use hg::{
     dirstate::StateMapIter,
     dirstate_tree::on_disk::DirstateV2ParseError,
     dirstate_tree::owning::OwningDirstateMap,
     revlog::Node,
     utils::files::normalize_case,
     utils::hg_path::{HgPath, HgPathBuf},
     DirstateEntry, DirstateError, DirstateParents, EntryState,
 };
 
 // TODO
 //     This object needs to share references to multiple members of its Rust
 //     inner struct, namely `copy_map`, `dirs` and `all_dirs`.
 //     Right now `CopyMap` is done, but it needs to have an explicit reference
 //     to `RustDirstateMap` which itself needs to have an encapsulation for
 //     every method in `CopyMap` (copymapcopy, etc.).
 //     This is ugly and hard to maintain.
 //     The same logic applies to `dirs` and `all_dirs`, however the `Dirs`
 //     `py_class!` is already implemented and does not mention
 //     `RustDirstateMap`, rightfully so.
 //     All attributes also have to have a separate refcount data attribute for
 //     leaks, with all methods that go along for reference sharing.
 py_class!(pub class DirstateMap |py| {
     @shared data inner: OwningDirstateMap;
 
     /// Returns a `(dirstate_map, parents)` tuple
     @staticmethod
     def new_v1(
         on_disk: PyBytes,
     ) -> PyResult<PyObject> {
         let on_disk = PyBytesDeref::new(py, on_disk);
         let (map, parents) = OwningDirstateMap::new_v1(on_disk)
             .map_err(|e| dirstate_error(py, e))?;
         let map = Self::create_instance(py, map)?;
         let p1 = PyBytes::new(py, parents.p1.as_bytes());
         let p2 = PyBytes::new(py, parents.p2.as_bytes());
         let parents = (p1, p2);
         Ok((map, parents).to_py_object(py).into_object())
     }
 
     /// Returns a DirstateMap
     @staticmethod
     def new_v2(
         on_disk: PyBytes,
         data_size: usize,
         tree_metadata: PyBytes,
     ) -> PyResult<PyObject> {
         let dirstate_error = |e: DirstateError| {
             PyErr::new::<exc::OSError, _>(py, format!("Dirstate error: {:?}", e))
         };
         let on_disk = PyBytesDeref::new(py, on_disk);
         let map = OwningDirstateMap::new_v2(
             on_disk, data_size, tree_metadata.data(py),
         ).map_err(dirstate_error)?;
         let map = Self::create_instance(py, map)?;
         Ok(map.into_object())
     }
 
     def clear(&self) -> PyResult<PyObject> {
         self.inner(py).borrow_mut().clear();
         Ok(py.None())
     }
 
     def get(
         &self,
         key: PyObject,
         default: Option<PyObject> = None
     ) -> PyResult<Option<PyObject>> {
         let key = key.extract::<PyBytes>(py)?;
         match self
             .inner(py)
             .borrow()
             .get(HgPath::new(key.data(py)))
             .map_err(|e| v2_error(py, e))?
         {
             Some(entry) => {
                 Ok(Some(DirstateItem::new_as_pyobject(py, entry)?))
             },
             None => Ok(default)
         }
     }
 
     def set_dirstate_item(
         &self,
         path: PyObject,
         item: DirstateItem
     ) -> PyResult<PyObject> {
         let f = path.extract::<PyBytes>(py)?;
         let filename = HgPath::new(f.data(py));
         self.inner(py)
             .borrow_mut()
             .set_entry(filename, item.get_entry(py))
             .map_err(|e| v2_error(py, e))?;
         Ok(py.None())
     }
 
     def addfile(
         &self,
         f: PyBytes,
         item: DirstateItem,
     ) -> PyResult<PyNone> {
         let filename = HgPath::new(f.data(py));
         let entry = item.get_entry(py);
         self.inner(py)
             .borrow_mut()
             .add_file(filename, entry)
             .map_err(|e |dirstate_error(py, e))?;
         Ok(PyNone)
     }
 
+    def set_tracked(&self, f: PyObject) -> PyResult<PyBool> {
+        let bytes = f.extract::<PyBytes>(py)?;
+        let path = HgPath::new(bytes.data(py));
+        let res = self.inner(py).borrow_mut().set_tracked(path);
+        let was_tracked = res.or_else(|_| {
+            Err(PyErr::new::<exc::OSError, _>(py, "Dirstate error".to_string()))
+        })?;
+        Ok(was_tracked.to_py_object(py))
+    }
+
     def removefile(
         &self,
         f: PyObject,
         in_merge: PyObject
     ) -> PyResult<PyObject> {
         self.inner(py).borrow_mut()
             .remove_file(
                 HgPath::new(f.extract::<PyBytes>(py)?.data(py)),
                 in_merge.extract::<PyBool>(py)?.is_true(),
             )
             .or_else(|_| {
                 Err(PyErr::new::<exc::OSError, _>(
                     py,
                     "Dirstate error".to_string(),
                 ))
             })?;
         Ok(py.None())
     }
 
     def drop_item_and_copy_source(
         &self,
         f: PyBytes,
     ) -> PyResult<PyNone> {
         self.inner(py)
             .borrow_mut()
             .drop_entry_and_copy_source(HgPath::new(f.data(py)))
             .map_err(|e |dirstate_error(py, e))?;
         Ok(PyNone)
     }
 
     def hastrackeddir(&self, d: PyObject) -> PyResult<PyBool> {
         let d = d.extract::<PyBytes>(py)?;
         Ok(self.inner(py).borrow_mut()
             .has_tracked_dir(HgPath::new(d.data(py)))
             .map_err(|e| {
                 PyErr::new::<exc::ValueError, _>(py, e.to_string())
             })?
             .to_py_object(py))
     }
 
     def hasdir(&self, d: PyObject) -> PyResult<PyBool> {
         let d = d.extract::<PyBytes>(py)?;
         Ok(self.inner(py).borrow_mut()
             .has_dir(HgPath::new(d.data(py)))
             .map_err(|e| {
                 PyErr::new::<exc::ValueError, _>(py, e.to_string())
             })?
             .to_py_object(py))
     }
 
     def write_v1(
         &self,
         p1: PyObject,
         p2: PyObject,
     ) -> PyResult<PyBytes> {
         let inner = self.inner(py).borrow();
         let parents = DirstateParents {
             p1: extract_node_id(py, &p1)?,
             p2: extract_node_id(py, &p2)?,
         };
         let result = inner.pack_v1(parents);
         match result {
             Ok(packed) => Ok(PyBytes::new(py, &packed)),
             Err(_) => Err(PyErr::new::<exc::OSError, _>(
                 py,
                 "Dirstate error".to_string(),
             )),
         }
     }
 
     /// Returns new data together with whether that data should be appended to
     /// the existing data file whose content is at `self.on_disk` (True),
     /// instead of written to a new data file (False).
     def write_v2(
         &self,
         can_append: bool,
     ) -> PyResult<PyObject> {
         let inner = self.inner(py).borrow();
         let result = inner.pack_v2(can_append);
         match result {
             Ok((packed, tree_metadata, append)) => {
                 let packed = PyBytes::new(py, &packed);
                 let tree_metadata = PyBytes::new(py, tree_metadata.as_bytes());
                 let tuple = (packed, tree_metadata, append);
                 Ok(tuple.to_py_object(py).into_object())
             },
             Err(_) => Err(PyErr::new::<exc::OSError, _>(
                 py,
                 "Dirstate error".to_string(),
             )),
         }
     }
 
     def filefoldmapasdict(&self) -> PyResult<PyDict> {
         let dict = PyDict::new(py);
         for item in self.inner(py).borrow_mut().iter() {
             let (path, entry) = item.map_err(|e| v2_error(py, e))?;
             if entry.state() != EntryState::Removed {
                 let key = normalize_case(path);
                 let value = path;
                 dict.set_item(
                     py,
                     PyBytes::new(py, key.as_bytes()).into_object(),
                     PyBytes::new(py, value.as_bytes()).into_object(),
                 )?;
             }
         }
         Ok(dict)
     }
 
     def __len__(&self) -> PyResult<usize> {
         Ok(self.inner(py).borrow().len())
     }
 
     def __contains__(&self, key: PyObject) -> PyResult<bool> {
         let key = key.extract::<PyBytes>(py)?;
         self.inner(py)
             .borrow()
             .contains_key(HgPath::new(key.data(py)))
             .map_err(|e| v2_error(py, e))
     }
 
     def __getitem__(&self, key: PyObject) -> PyResult<PyObject> {
         let key = key.extract::<PyBytes>(py)?;
         let key = HgPath::new(key.data(py));
         match self
             .inner(py)
             .borrow()
             .get(key)
             .map_err(|e| v2_error(py, e))?
         {
             Some(entry) => {
                 Ok(DirstateItem::new_as_pyobject(py, entry)?)
             },
             None => Err(PyErr::new::<exc::KeyError, _>(
                 py,
                 String::from_utf8_lossy(key.as_bytes()),
             )),
         }
     }
 
     def keys(&self) -> PyResult<DirstateMapKeysIterator> {
         let leaked_ref = self.inner(py).leak_immutable();
         DirstateMapKeysIterator::from_inner(
             py,
             unsafe { leaked_ref.map(py, |o| o.iter()) },
         )
     }
 
     def items(&self) -> PyResult<DirstateMapItemsIterator> {
         let leaked_ref = self.inner(py).leak_immutable();
         DirstateMapItemsIterator::from_inner(
             py,
             unsafe { leaked_ref.map(py, |o| o.iter()) },
         )
     }
 
     def __iter__(&self) -> PyResult<DirstateMapKeysIterator> {
         let leaked_ref = self.inner(py).leak_immutable();
         DirstateMapKeysIterator::from_inner(
             py,
             unsafe { leaked_ref.map(py, |o| o.iter()) },
         )
     }
 
     // TODO all copymap* methods, see docstring above
     def copymapcopy(&self) -> PyResult<PyDict> {
         let dict = PyDict::new(py);
         for item in self.inner(py).borrow().copy_map_iter() {
             let (key, value) = item.map_err(|e| v2_error(py, e))?;
             dict.set_item(
                 py,
                 PyBytes::new(py, key.as_bytes()),
                 PyBytes::new(py, value.as_bytes()),
             )?;
         }
         Ok(dict)
     }
 
     def copymapgetitem(&self, key: PyObject) -> PyResult<PyBytes> {
         let key = key.extract::<PyBytes>(py)?;
         match self
             .inner(py)
             .borrow()
             .copy_map_get(HgPath::new(key.data(py)))
             .map_err(|e| v2_error(py, e))?
         {
             Some(copy) => Ok(PyBytes::new(py, copy.as_bytes())),
             None => Err(PyErr::new::<exc::KeyError, _>(
                 py,
                 String::from_utf8_lossy(key.data(py)),
             )),
         }
     }
     def copymap(&self) -> PyResult<CopyMap> {
         CopyMap::from_inner(py, self.clone_ref(py))
     }
 
     def copymaplen(&self) -> PyResult<usize> {
         Ok(self.inner(py).borrow().copy_map_len())
     }
     def copymapcontains(&self, key: PyObject) -> PyResult<bool> {
         let key = key.extract::<PyBytes>(py)?;
         self.inner(py)
             .borrow()
             .copy_map_contains_key(HgPath::new(key.data(py)))
             .map_err(|e| v2_error(py, e))
     }
     def copymapget(
         &self,
         key: PyObject,
         default: Option<PyObject>
     ) -> PyResult<Option<PyObject>> {
         let key = key.extract::<PyBytes>(py)?;
         match self
             .inner(py)
             .borrow()
             .copy_map_get(HgPath::new(key.data(py)))
             .map_err(|e| v2_error(py, e))?
         {
             Some(copy) => Ok(Some(
                 PyBytes::new(py, copy.as_bytes()).into_object(),
             )),
             None => Ok(default),
         }
     }
     def copymapsetitem(
         &self,
         key: PyObject,
         value: PyObject
     ) -> PyResult<PyObject> {
         let key = key.extract::<PyBytes>(py)?;
         let value = value.extract::<PyBytes>(py)?;
         self.inner(py)
             .borrow_mut()
             .copy_map_insert(
                 HgPathBuf::from_bytes(key.data(py)),
                 HgPathBuf::from_bytes(value.data(py)),
             )
             .map_err(|e| v2_error(py, e))?;
         Ok(py.None())
     }
     def copymappop(
         &self,
         key: PyObject,
         default: Option<PyObject>
     ) -> PyResult<Option<PyObject>> {
         let key = key.extract::<PyBytes>(py)?;
         match self
             .inner(py)
             .borrow_mut()
             .copy_map_remove(HgPath::new(key.data(py)))
             .map_err(|e| v2_error(py, e))?
         {
             Some(copy) => Ok(Some(
                 PyBytes::new(py, copy.as_bytes()).into_object(),
             )),
             None => Ok(default),
         }
     }
 
     def copymapiter(&self) -> PyResult<CopyMapKeysIterator> {
         let leaked_ref = self.inner(py).leak_immutable();
         CopyMapKeysIterator::from_inner(
             py,
             unsafe { leaked_ref.map(py, |o| o.copy_map_iter()) },
         )
     }
 
     def copymapitemsiter(&self) -> PyResult<CopyMapItemsIterator> {
         let leaked_ref = self.inner(py).leak_immutable();
         CopyMapItemsIterator::from_inner(
             py,
             unsafe { leaked_ref.map(py, |o| o.copy_map_iter()) },
         )
     }
 
     def tracked_dirs(&self) -> PyResult<PyList> {
         let dirs = PyList::new(py, &[]);
         for path in self.inner(py).borrow_mut().iter_tracked_dirs()
             .map_err(|e |dirstate_error(py, e))?
         {
             let path = path.map_err(|e| v2_error(py, e))?;
             let path = PyBytes::new(py, path.as_bytes());
             dirs.append(py, path.into_object())
         }
         Ok(dirs)
     }
 
     def debug_iter(&self, all: bool) -> PyResult<PyList> {
         let dirs = PyList::new(py, &[]);
         for item in self.inner(py).borrow().debug_iter(all) {
             let (path, (state, mode, size, mtime)) =
                 item.map_err(|e| v2_error(py, e))?;
             let path = PyBytes::new(py, path.as_bytes());
             let item = (path, state, mode, size, mtime);
             dirs.append(py, item.to_py_object(py).into_object())
         }
         Ok(dirs)
     }
 });
 
 impl DirstateMap {
     pub fn get_inner_mut<'a>(
         &'a self,
         py: Python<'a>,
     ) -> RefMut<'a, OwningDirstateMap> {
         self.inner(py).borrow_mut()
     }
     fn translate_key(
         py: Python,
         res: Result<(&HgPath, DirstateEntry), DirstateV2ParseError>,
     ) -> PyResult<Option<PyBytes>> {
         let (f, _entry) = res.map_err(|e| v2_error(py, e))?;
         Ok(Some(PyBytes::new(py, f.as_bytes())))
     }
     fn translate_key_value(
         py: Python,
         res: Result<(&HgPath, DirstateEntry), DirstateV2ParseError>,
     ) -> PyResult<Option<(PyBytes, PyObject)>> {
         let (f, entry) = res.map_err(|e| v2_error(py, e))?;
         Ok(Some((
             PyBytes::new(py, f.as_bytes()),
             DirstateItem::new_as_pyobject(py, entry)?,
         )))
     }
 }
 
 py_shared_iterator!(
     DirstateMapKeysIterator,
     UnsafePyLeaked<StateMapIter<'static>>,
     DirstateMap::translate_key,
     Option<PyBytes>
 );
 
 py_shared_iterator!(
     DirstateMapItemsIterator,
     UnsafePyLeaked<StateMapIter<'static>>,
     DirstateMap::translate_key_value,
     Option<(PyBytes, PyObject)>
 );
 
 fn extract_node_id(py: Python, obj: &PyObject) -> PyResult<Node> {
     let bytes = obj.extract::<PyBytes>(py)?;
     match bytes.data(py).try_into() {
         Ok(s) => Ok(s),
         Err(e) => Err(PyErr::new::<exc::ValueError, _>(py, e.to_string())),
     }
 }
 
 pub(super) fn v2_error(py: Python<'_>, _: DirstateV2ParseError) -> PyErr {
     PyErr::new::<exc::ValueError, _>(py, "corrupted dirstate-v2")
 }
 
 fn dirstate_error(py: Python<'_>, e: DirstateError) -> PyErr {
     PyErr::new::<exc::OSError, _>(py, format!("Dirstate error: {:?}", e))
 }