rust: Add Repo::dirstate_map and use it in `rhg status`

This moves low-level dirstate wrangling out of the status command and into
a more reusable location.

The open dirstate map is lazily initialized and kept on the Repo object,
for reuse by sub-sequent calls.

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

rust/hg-core/src/dirstate.rs

@@ -19,7 +19,7 b' pub mod dirstate_map;'
 pub mod parsers;
 pub mod status;
 
-#[derive(Debug, PartialEq, Clone, BytesCast)]
+#[derive(Debug, PartialEq, Copy, Clone, BytesCast)]
 #[repr(C)]
 pub struct DirstateParents {
     pub p1: Node,

rust/hg-core/src/repo.rs

@@ -1,10 +1,16 b''
 use crate::config::{Config, ConfigError, ConfigParseError};
+use crate::dirstate::DirstateParents;
+use crate::dirstate_tree::dirstate_map::DirstateMap;
+use crate::dirstate_tree::owning::OwningDirstateMap;
 use crate::errors::HgError;
+use crate::errors::HgResultExt;
 use crate::exit_codes;
 use crate::requirements;
 use crate::utils::files::get_path_from_bytes;
 use crate::utils::SliceExt;
 use crate::vfs::{is_dir, is_file, Vfs};
+use crate::DirstateError;
+use std::cell::{Cell, Ref, RefCell, RefMut};
 use std::collections::HashSet;
 use std::path::{Path, PathBuf};
 
@@ -15,6 +21,9 b' pub struct Repo {'
     store: PathBuf,
     requirements: HashSet<String>,
     config: Config,
+    // None means not known/initialized yet
+    dirstate_parents: Cell<Option<DirstateParents>>,
+    dirstate_map: RefCell<Option<OwningDirstateMap>>,
 }
 
 #[derive(Debug, derive_more::From)]
@@ -186,6 +195,8 b' impl Repo {'
             store: store_path,
             dot_hg,
             config: repo_config,
+            dirstate_parents: Cell::new(None),
+            dirstate_map: RefCell::new(None),
         };
 
         requirements::check(&repo)?;
@@ -228,19 +239,101 b' impl Repo {'
             .contains(requirements::DIRSTATE_V2_REQUIREMENT)
     }
 
-    pub fn dirstate_parents(
-        &self,
-    ) -> Result<crate::dirstate::DirstateParents, HgError> {
-        let dirstate = self.hg_vfs().mmap_open("dirstate")?;
-        if dirstate.is_empty() {
-            return Ok(crate::dirstate::DirstateParents::NULL);
+    fn dirstate_file_contents(&self) -> Result<Vec<u8>, HgError> {
+        Ok(self
+            .hg_vfs()
+            .read("dirstate")
+            .io_not_found_as_none()?
+            .unwrap_or(Vec::new()))
+    }
+
+    pub fn dirstate_parents(&self) -> Result<DirstateParents, HgError> {
+        if let Some(parents) = self.dirstate_parents.get() {
+            return Ok(parents);
         }
-        let parents = if self.has_dirstate_v2() {
+        let dirstate = self.dirstate_file_contents()?;
+        let parents = if dirstate.is_empty() {
+            DirstateParents::NULL
+        } else if self.has_dirstate_v2() {
             crate::dirstate_tree::on_disk::read_docket(&dirstate)?.parents()
         } else {
             crate::dirstate::parsers::parse_dirstate_parents(&dirstate)?
                 .clone()
         };
+        self.dirstate_parents.set(Some(parents));
         Ok(parents)
     }
+
+    fn new_dirstate_map(&self) -> Result<OwningDirstateMap, DirstateError> {
+        let dirstate_file_contents = self.dirstate_file_contents()?;
+        if dirstate_file_contents.is_empty() {
+            self.dirstate_parents.set(Some(DirstateParents::NULL));
+            Ok(OwningDirstateMap::new_empty(Vec::new()))
+        } else if self.has_dirstate_v2() {
+            let docket = crate::dirstate_tree::on_disk::read_docket(
+                &dirstate_file_contents,
+            )?;
+            self.dirstate_parents.set(Some(docket.parents()));
+            let data_size = docket.data_size();
+            let metadata = docket.tree_metadata();
+            let mut map = if let Some(data_mmap) = self
+                .hg_vfs()
+                .mmap_open(docket.data_filename())
+                .io_not_found_as_none()?
+            {
+                OwningDirstateMap::new_empty(MmapWrapper(data_mmap))
+            } else {
+                OwningDirstateMap::new_empty(Vec::new())
+            };
+            let (on_disk, placeholder) = map.get_mut_pair();
+            *placeholder = DirstateMap::new_v2(on_disk, data_size, metadata)?;
+            Ok(map)
+        } else {
+            let mut map = OwningDirstateMap::new_empty(dirstate_file_contents);
+            let (on_disk, placeholder) = map.get_mut_pair();
+            let (inner, parents) = DirstateMap::new_v1(on_disk)?;
+            self.dirstate_parents
+                .set(Some(parents.unwrap_or(DirstateParents::NULL)));
+            *placeholder = inner;
+            Ok(map)
+        }
+    }
+
+    pub fn dirstate_map(
+        &self,
+    ) -> Result<Ref<OwningDirstateMap>, DirstateError> {
+        let mut borrowed = self.dirstate_map.borrow();
+        if borrowed.is_none() {
+            drop(borrowed);
+            // Only use `borrow_mut` if it is really needed to avoid panic in
+            // case there is another outstanding borrow but mutation is not
+            // needed.
+            *self.dirstate_map.borrow_mut() = Some(self.new_dirstate_map()?);
+            borrowed = self.dirstate_map.borrow()
+        }
+        Ok(Ref::map(borrowed, |option| option.as_ref().unwrap()))
+    }
+
+    pub fn dirstate_map_mut(
+        &self,
+    ) -> Result<RefMut<OwningDirstateMap>, DirstateError> {
+        let mut borrowed = self.dirstate_map.borrow_mut();
+        if borrowed.is_none() {
+            *borrowed = Some(self.new_dirstate_map()?);
+        }
+        Ok(RefMut::map(borrowed, |option| option.as_mut().unwrap()))
+    }
 }
+
+// TODO: remove this when https://github.com/RazrFalcon/memmap2-rs/pull/22 is on crates.io
+struct MmapWrapper(memmap2::Mmap);
+
+impl std::ops::Deref for MmapWrapper {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        self.0.deref()
+    }
+}
+
+unsafe impl stable_deref_trait::StableDeref for MmapWrapper {}

rust/rhg/src/commands/status.rs

@@ -9,9 +9,7 b' use crate::error::CommandError;'
 use crate::ui::Ui;
 use clap::{Arg, SubCommand};
 use hg;
-use hg::dirstate_tree::dirstate_map::DirstateMap;
-use hg::dirstate_tree::on_disk;
-use hg::errors::HgResultExt;
+use hg::dirstate_tree::dispatch::DirstateMapMethods;
 use hg::errors::IoResultExt;
 use hg::matchers::AlwaysMatcher;
 use hg::operations::cat;
@@ -166,40 +164,7 b' pub fn run(invocation: &crate::CliInvoca'
     };
 
     let repo = invocation.repo?;
-    let dirstate_data_mmap;
-    let (mut dmap, parents) = if repo.has_dirstate_v2() {
-        let docket_data =
-            repo.hg_vfs().read("dirstate").io_not_found_as_none()?;
-        let parents;
-        let dirstate_data;
-        let data_size;
-        let docket;
-        let tree_metadata;
-        if let Some(docket_data) = &docket_data {
-            docket = on_disk::read_docket(docket_data)?;
-            tree_metadata = docket.tree_metadata();
-            parents = Some(docket.parents());
-            data_size = docket.data_size();
-            dirstate_data_mmap = repo
-                .hg_vfs()
-                .mmap_open(docket.data_filename())
-                .io_not_found_as_none()?;
-            dirstate_data = dirstate_data_mmap.as_deref().unwrap_or(b"");
-        } else {
-            parents = None;
-            tree_metadata = b"";
-            data_size = 0;
-            dirstate_data = b"";
-        }
-        let dmap =
-            DirstateMap::new_v2(dirstate_data, data_size, tree_metadata)?;
-        (dmap, parents)
-    } else {
-        dirstate_data_mmap =
-            repo.hg_vfs().mmap_open("dirstate").io_not_found_as_none()?;
-        let dirstate_data = dirstate_data_mmap.as_deref().unwrap_or(b"");
-        DirstateMap::new_v1(dirstate_data)?
-    };
+    let mut dmap = repo.dirstate_map_mut()?;
 
     let options = StatusOptions {
         // TODO should be provided by the dirstate parsing and
@@ -216,8 +181,7 b' pub fn run(invocation: &crate::CliInvoca'
         collect_traversed_dirs: false,
     };
     let ignore_file = repo.working_directory_vfs().join(".hgignore"); // TODO hardcoded
-    let (mut ds_status, pattern_warnings) = hg::dirstate_tree::status::status(
-        &mut dmap,
+    let (mut ds_status, pattern_warnings) = dmap.status(
         &AlwaysMatcher,
         repo.working_directory_path().to_owned(),
         vec![ignore_file],
@@ -239,13 +203,7 b' pub fn run(invocation: &crate::CliInvoca'
     if !ds_status.unsure.is_empty()
         && (display_states.modified || display_states.clean)
     {
-        let p1: Node = parents
-            .expect(
-                "Dirstate with no parents should not list any file to
-            be rechecked for modifications",
-            )
-            .p1
-            .into();
+        let p1: Node = repo.dirstate_parents()?.p1.into();
         let p1_hex = format!("{:x}", p1);
         for to_check in ds_status.unsure {
             if cat_file_is_modified(repo, &to_check, &p1_hex)? {