changeset 51220:c817d9f626d3

rust-index: add support for `computephasesmapsets`

Exposition in `hg-cpython` done in the regular `impl` block to enjoy
rustfmt and clearer compilartion errors.

--- a/rust/hg-core/src/revlog/index.rs	Sat Sep 30 15:59:03 2023 +0200
+++ b/rust/hg-core/src/revlog/index.rs	Mon Oct 30 11:54:42 2023 +0100
@@ -1,5 +1,5 @@
 use std::collections::hash_map::RandomState;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
 use std::ops::Deref;
 use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
@@ -910,7 +910,128 @@
         }
         &revs[start..end]
     }
+
+    /// Computes the set of revisions for each non-public phase from `roots`,
+    /// which are the last known roots for each non-public phase.
+    pub fn compute_phases_map_sets(
+        &self,
+        roots: HashMap<Phase, Vec<Revision>>,
+    ) -> Result<(usize, RootsPerPhase), GraphError> {
+        let mut phases = HashMap::new();
+        let mut min_phase_rev = NULL_REVISION;
+
+        for phase in Phase::non_public_phases() {
+            if let Some(phase_roots) = roots.get(phase) {
+                let min_rev =
+                    self.add_roots_get_min(phase_roots, &mut phases, *phase);
+                if min_rev != NULL_REVISION
+                    && (min_phase_rev == NULL_REVISION
+                        || min_rev < min_phase_rev)
+                {
+                    min_phase_rev = min_rev;
+                }
+            } else {
+                continue;
+            };
+        }
+        let mut phase_sets: RootsPerPhase = Default::default();
+
+        if min_phase_rev == NULL_REVISION {
+            min_phase_rev = Revision(self.len() as BaseRevision);
+        }
+
+        for rev in min_phase_rev.0..self.len() as BaseRevision {
+            let rev = Revision(rev);
+            let [p1, p2] = self.parents(rev)?;
+
+            const DEFAULT_PHASE: &Phase = &Phase::Public;
+            if p1.0 >= 0
+                && phases.get(&p1).unwrap_or(DEFAULT_PHASE)
+                    > phases.get(&rev).unwrap_or(DEFAULT_PHASE)
+            {
+                phases.insert(rev, phases[&p1]);
+            }
+            if p2.0 >= 0
+                && phases.get(&p2).unwrap_or(DEFAULT_PHASE)
+                    > phases.get(&rev).unwrap_or(DEFAULT_PHASE)
+            {
+                phases.insert(rev, phases[&p2]);
+            }
+            let set = match phases.get(&rev).unwrap_or(DEFAULT_PHASE) {
+                Phase::Public => continue,
+                phase => &mut phase_sets[*phase as usize - 1],
+            };
+            set.insert(rev);
+        }
+
+        Ok((self.len(), phase_sets))
+    }
+
+    fn add_roots_get_min(
+        &self,
+        phase_roots: &[Revision],
+        phases: &mut HashMap<Revision, Phase>,
+        phase: Phase,
+    ) -> Revision {
+        let mut min_rev = NULL_REVISION;
+
+        for root in phase_roots {
+            phases.insert(*root, phase);
+            if min_rev == NULL_REVISION || min_rev > *root {
+                min_rev = *root;
+            }
+        }
+        min_rev
+    }
 }
+
+/// Set of roots of all non-public phases
+pub type RootsPerPhase = [HashSet<Revision>; Phase::non_public_phases().len()];
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash)]
+pub enum Phase {
+    Public = 0,
+    Draft = 1,
+    Secret = 2,
+    Archived = 3,
+    Internal = 4,
+}
+
+impl TryFrom<usize> for Phase {
+    type Error = RevlogError;
+
+    fn try_from(value: usize) -> Result<Self, Self::Error> {
+        Ok(match value {
+            0 => Self::Public,
+            1 => Self::Draft,
+            2 => Self::Secret,
+            32 => Self::Archived,
+            96 => Self::Internal,
+            v => {
+                return Err(RevlogError::corrupted(format!(
+                    "invalid phase value {}",
+                    v
+                )))
+            }
+        })
+    }
+}
+
+impl Phase {
+    pub const fn all_phases() -> &'static [Self] {
+        &[
+            Self::Public,
+            Self::Draft,
+            Self::Secret,
+            Self::Archived,
+            Self::Internal,
+        ]
+    }
+    pub const fn non_public_phases() -> &'static [Self] {
+        &[Self::Draft, Self::Secret, Self::Archived, Self::Internal]
+    }
+}
+
 fn inline_scan(bytes: &[u8]) -> (usize, Vec<usize>) {
     let mut offset: usize = 0;
     let mut offsets = Vec::new();

--- a/rust/hg-cpython/src/revlog.rs	Sat Sep 30 15:59:03 2023 +0200
+++ b/rust/hg-cpython/src/revlog.rs	Mon Oct 30 11:54:42 2023 +0100
@@ -20,12 +20,12 @@
 };
 use hg::{
     errors::HgError,
-    index::{IndexHeader, RevisionDataParams, SnapshotsCache},
+    index::{IndexHeader, Phase, RevisionDataParams, SnapshotsCache},
     nodemap::{Block, NodeMapError, NodeTree},
     revlog::{nodemap::NodeMap, NodePrefix, RevlogError, RevlogIndex},
     BaseRevision, Revision, UncheckedRevision, NULL_REVISION,
 };
-use std::cell::RefCell;
+use std::{cell::RefCell, collections::HashMap};
 
 /// Return a Struct implementing the Graph trait
 pub(crate) fn pyindex_to_graph(
@@ -241,7 +241,12 @@
 
     /// compute phases
     def computephasesmapsets(&self, *args, **kw) -> PyResult<PyObject> {
-        self.call_cindex(py, "computephasesmapsets", args, kw)
+        let py_roots = args.get_item(py, 0).extract::<PyDict>(py)?;
+        let rust_res = self.inner_computephasesmapsets(py, py_roots)?;
+
+        let c_res = self.call_cindex(py, "computephasesmapsets", args, kw)?;
+        assert_py_eq(py, "computephasesmapsets", &rust_res, &c_res)?;
+        Ok(rust_res)
     }
 
     /// reachableroots
@@ -829,6 +834,71 @@
         Ok(PyList::new(py, &as_vec).into_object())
     }
 
+    fn inner_computephasesmapsets(
+        &self,
+        py: Python,
+        py_roots: PyDict,
+    ) -> PyResult<PyObject> {
+        let index = &*self.index(py).borrow();
+        let opt = self.get_nodetree(py)?.borrow();
+        let nt = opt.as_ref().unwrap();
+        let roots: Result<HashMap<Phase, Vec<Revision>>, PyErr> = py_roots
+            .items_list(py)
+            .iter(py)
+            .map(|r| {
+                let phase = r.get_item(py, 0)?;
+                let nodes = r.get_item(py, 1)?;
+                // Transform the nodes from Python to revs here since we
+                // have access to the nodemap
+                let revs: Result<_, _> = nodes
+                    .iter(py)?
+                    .map(|node| match node?.extract::<PyBytes>(py) {
+                        Ok(py_bytes) => {
+                            let node = node_from_py_bytes(py, &py_bytes)?;
+                            nt.find_bin(index, node.into())
+                                .map_err(|e| nodemap_error(py, e))?
+                                .ok_or_else(|| revlog_error(py))
+                        }
+                        Err(e) => Err(e),
+                    })
+                    .collect();
+                let phase = Phase::try_from(phase.extract::<usize>(py)?)
+                    .map_err(|_| revlog_error(py));
+                Ok((phase?, revs?))
+            })
+            .collect();
+        let (len, phase_maps) = index
+            .compute_phases_map_sets(roots?)
+            .map_err(|e| graph_error(py, e))?;
+
+        // Ugly hack, but temporary
+        const IDX_TO_PHASE_NUM: [usize; 4] = [1, 2, 32, 96];
+        let py_phase_maps = PyDict::new(py);
+        for (idx, roots) in phase_maps.iter().enumerate() {
+            let phase_num = IDX_TO_PHASE_NUM[idx].into_py_object(py);
+            // OPTIM too bad we have to collect here. At least, we could
+            // reuse the same Vec and allocate it with capacity at
+            // max(len(phase_maps)
+            let roots_vec: Vec<PyInt> = roots
+                .iter()
+                .map(|r| PyRevision::from(*r).into_py_object(py))
+                .collect();
+            py_phase_maps.set_item(
+                py,
+                phase_num,
+                PySet::new(py, roots_vec)?,
+            )?;
+        }
+        Ok(PyTuple::new(
+            py,
+            &[
+                len.into_py_object(py).into_object(),
+                py_phase_maps.into_object(),
+            ],
+        )
+        .into_object())
+    }
+
     fn inner_slicechunktodensity(
         &self,
         py: Python,