--- a/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Wed Feb 05 11:04:19 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Thu Jan 30 12:24:41 2025 +0100
@@ -8,35 +8,47 @@
 //! Bindings for the `hg::dirstate::dirstate_map` file provided by the
 //! `hg-core` package.
 
-use pyo3::exceptions::PyKeyError;
+use pyo3::exceptions::{PyKeyError, PyOSError};
 use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyBytesMethods, PyTuple};
+use pyo3::types::{
+    PyBytes, PyBytesMethods, PyDict, PyDictMethods, PyList, PyTuple,
+};
 use pyo3_sharedref::{py_shared_iterator, PyShareable};
 
-use std::sync::RwLockReadGuard;
+use std::sync::{RwLockReadGuard, RwLockWriteGuard};
 
 use hg::{
     dirstate::{
         dirstate_map::{
-            DirstateIdentity as CoreDirstateIdentity, DirstateMapWriteMode,
+            DirstateEntryReset, DirstateIdentity as CoreDirstateIdentity,
+            DirstateMapWriteMode,
         },
-        entry::DirstateEntry,
+        entry::{DirstateEntry, ParentFileData, TruncatedTimestamp},
         on_disk::DirstateV2ParseError,
         owning::OwningDirstateMap,
         StateMapIter,
     },
-    utils::hg_path::HgPath,
+    utils::{files::normalize_case, hg_path::HgPath},
     DirstateParents,
 };
 
 use super::item::DirstateItem;
 use crate::{
-    exceptions::{dirstate_error, dirstate_v2_error, map_try_lock_error},
+    exceptions::{
+        dirstate_error, dirstate_v2_error, map_try_lock_error,
+        to_string_value_error,
+    },
     node::{node_from_py_bytes, PyNode},
-    path::PyHgPathRef,
+    path::{PyHgPathBuf, PyHgPathDirstateV2Result, PyHgPathRef},
     utils::PyBytesDeref,
 };
 
+/// Type alias to satisfy Clippy in `DirstateMap::reset_state)`
+///
+/// It is *not* the same as [`super::item::UncheckedTruncatedTimeStamp`] and
+/// this is worth reviewing.
+type UncheckedTruncatedTimeStamp = Option<(i64, u32, bool)>;
+
 #[pyclass(mapping)]
 pub struct DirstateMap {
     inner: PyShareable<OwningDirstateMap>,
@@ -97,6 +109,13 @@
         })
     }
 
+    fn clear(slf: &Bound<'_, Self>) -> PyResult<()> {
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            inner.clear();
+            Ok(())
+        })
+    }
+
     #[pyo3(signature = (key, default=None))]
     fn get(
         slf: &Bound<'_, Self>,
@@ -115,6 +134,141 @@
         })
     }
 
+    fn set_tracked(
+        slf: &Bound<'_, Self>,
+        f: &Bound<'_, PyBytes>,
+    ) -> PyResult<bool> {
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            inner
+                .set_tracked(HgPath::new(f.as_bytes()))
+                .map_err(dirstate_v2_error)
+        })
+    }
+
+    fn set_untracked(
+        slf: &Bound<'_, Self>,
+        f: &Bound<'_, PyBytes>,
+    ) -> PyResult<bool> {
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            // here it would be more straightforward to use dirstate_v2_error,
+            // but that raises ValueError instead of OSError
+            inner
+                .set_untracked(HgPath::new(f.as_bytes()))
+                .map_err(|_| PyOSError::new_err("Dirstate error"))
+        })
+    }
+
+    fn set_clean(
+        slf: &Bound<'_, Self>,
+        f: &Bound<'_, PyBytes>,
+        mode: u32,
+        size: u32,
+        mtime: (i64, u32, bool),
+    ) -> PyResult<()> {
+        let (mtime_s, mtime_ns, second_ambiguous) = mtime;
+        let timestamp = TruncatedTimestamp::new_truncate(
+            mtime_s,
+            mtime_ns,
+            second_ambiguous,
+        );
+
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            inner
+                .set_clean(HgPath::new(f.as_bytes()), mode, size, timestamp)
+                .map_err(dirstate_error)
+        })
+    }
+
+    fn set_possibly_dirty(
+        slf: &Bound<'_, Self>,
+        f: &Bound<'_, PyBytes>,
+    ) -> PyResult<()> {
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            inner
+                .set_possibly_dirty(HgPath::new(f.as_bytes()))
+                .map_err(dirstate_error)
+        })
+    }
+
+    #[pyo3(signature = (f,
+                        wc_tracked=false,
+                        p1_tracked=false,
+                        p2_info=false,
+                        has_meaningful_mtime=true,
+                        parentfiledata=None))]
+    fn reset_state(
+        slf: &Bound<'_, Self>,
+        f: &Bound<'_, PyBytes>,
+        wc_tracked: bool,
+        p1_tracked: bool,
+        p2_info: bool,
+        has_meaningful_mtime: bool,
+        parentfiledata: Option<(u32, u32, UncheckedTruncatedTimeStamp)>,
+    ) -> PyResult<()> {
+        let mut has_meaningful_mtime = has_meaningful_mtime;
+        let parent_file_data = match parentfiledata {
+            None => {
+                has_meaningful_mtime = false;
+                None
+            }
+            Some(data) => {
+                let (mode, size, mtime_info) = data;
+                let mtime = if let Some(mtime_info) = mtime_info {
+                    let (mtime_s, mtime_ns, second_ambiguous) = mtime_info;
+                    let timestamp = TruncatedTimestamp::new_truncate(
+                        mtime_s,
+                        mtime_ns,
+                        second_ambiguous,
+                    );
+                    Some(timestamp)
+                } else {
+                    has_meaningful_mtime = false;
+                    None
+                };
+                Some(ParentFileData {
+                    mode_size: Some((mode, size)),
+                    mtime,
+                })
+            }
+        };
+
+        let reset = DirstateEntryReset {
+            filename: HgPath::new(f.as_bytes()),
+            wc_tracked,
+            p1_tracked,
+            p2_info,
+            has_meaningful_mtime,
+            parent_file_data_opt: parent_file_data,
+            from_empty: false,
+        };
+
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            inner.reset_state(reset).map_err(dirstate_error)
+        })
+    }
+
+    fn hastrackeddir(
+        slf: &Bound<'_, Self>,
+        d: &Bound<'_, PyBytes>,
+    ) -> PyResult<bool> {
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            inner
+                .has_tracked_dir(HgPath::new(d.as_bytes()))
+                .map_err(to_string_value_error)
+        })
+    }
+
+    fn hasdir(
+        slf: &Bound<'_, Self>,
+        d: &Bound<'_, PyBytes>,
+    ) -> PyResult<bool> {
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            inner
+                .has_dir(HgPath::new(d.as_bytes()))
+                .map_err(to_string_value_error)
+        })
+    }
+
     /// Returns suitable data for writing on disk in v1 format
     ///
     /// Despite the name, this is not a mutation of the object.
@@ -168,6 +322,24 @@
         })
     }
 
+    fn filefoldmapasdict(
+        slf: &Bound<'_, Self>,
+        py: Python,
+    ) -> PyResult<Py<PyDict>> {
+        let dict = PyDict::new(py);
+        Self::with_inner_read(slf, |_self_ref, inner| {
+            for item in inner.iter() {
+                let (path, entry) = item.map_err(dirstate_v2_error)?;
+                if !entry.removed() {
+                    let key = normalize_case(path);
+                    dict.set_item(PyHgPathBuf(key), PyHgPathRef(path))?;
+                }
+            }
+            Ok(())
+        })?;
+        Ok(dict.unbind())
+    }
+
     fn __len__(slf: &Bound<'_, Self>) -> PyResult<usize> {
         Self::with_inner_read(slf, |_self_ref, inner| Ok(inner.len()))
     }
@@ -211,6 +383,39 @@
         Self::keys(slf)
     }
 
+    fn tracked_dirs(
+        slf: &Bound<'_, Self>,
+        py: Python,
+    ) -> PyResult<Py<PyList>> {
+        // core iterator is not exact sized, we cannot use `PyList::new`
+        let dirs = PyList::empty(py);
+        Self::with_inner_write(slf, |_self_ref, mut inner| {
+            for path in inner.iter_tracked_dirs().map_err(dirstate_error)? {
+                dirs.append(PyHgPathDirstateV2Result(path))?;
+            }
+            Ok(())
+        })?;
+        Ok(dirs.unbind())
+    }
+
+    fn setparents_fixup(
+        slf: &Bound<'_, Self>,
+        py: Python,
+    ) -> PyResult<Py<PyDict>> {
+        let dict = PyDict::new(py);
+        let copies = Self::with_inner_write(slf, |_self_ref, mut inner| {
+            inner.setparents_fixup().map_err(dirstate_v2_error)
+        })?;
+
+        // it might be interesting to try and use the `IntoPyDict` trait,
+        // but it does about the same thing
+        // but that would require performing the inner `as_bytes()` as well
+        for (key, value) in copies {
+            dict.set_item(PyHgPathBuf(key), PyHgPathBuf(value))?;
+        }
+        Ok(dict.unbind())
+    }
+
     fn debug_iter(
         slf: &Bound<'_, Self>,
         py: Python,
@@ -287,6 +492,21 @@
         let guard = shareable_ref.try_read().map_err(map_try_lock_error)?;
         f(&self_ref, guard)
     }
+
+    fn with_inner_write<'py, T>(
+        slf: &Bound<'py, Self>,
+        f: impl FnOnce(
+            &PyRef<'py, Self>,
+            RwLockWriteGuard<OwningDirstateMap>,
+        ) -> PyResult<T>,
+    ) -> PyResult<T> {
+        let self_ref = slf.borrow();
+        // Safety: the owner is the right one. We will anyway
+        // not actually `share` it.
+        let shareable_ref = unsafe { self_ref.inner.borrow_with_owner(slf) };
+        let guard = shareable_ref.try_write().map_err(map_try_lock_error)?;
+        f(&self_ref, guard)
+    }
 }
 
 #[pyclass]

--- a/rust/hg-pyo3/src/exceptions.rs	Wed Feb 05 11:04:19 2025 +0100
+++ b/rust/hg-pyo3/src/exceptions.rs	Thu Jan 30 12:24:41 2025 +0100
@@ -52,7 +52,6 @@
     PyRuntimeError::new_err(format!("In Rust PyO3 bindings: {e}"))
 }
 
-#[allow(unused)]
 pub fn to_string_value_error<T: Display>(e: T) -> PyErr {
     PyValueError::new_err(e.to_string())
 }