# HG changeset patch
# User Raphaël Gomès <rgomes@octobus.net>
# Date 1695129703 -7200
# Node ID a7bba7df9189cdacb2c3f3069f9906cc8944001f
# Parent  050098d60c30130e450eb879ffc9eba28dc32167
rust-index: implement headrevs

diff -r 050098d60c30 -r a7bba7df9189 rust/hg-core/src/revlog/index.rs
--- a/rust/hg-core/src/revlog/index.rs	Sat Sep 30 16:52:40 2023 +0200
+++ b/rust/hg-core/src/revlog/index.rs	Tue Sep 19 15:21:43 2023 +0200
@@ -1,3 +1,4 @@
+use std::collections::hash_map::RandomState;
 use std::collections::HashSet;
 use std::fmt::Debug;
 use std::ops::Deref;
@@ -12,8 +13,8 @@
 use crate::revlog::node::Node;
 use crate::revlog::{Revision, NULL_REVISION};
 use crate::{
-    BaseRevision, FastHashMap, Graph, GraphError, RevlogError, RevlogIndex,
-    UncheckedRevision,
+    dagops, BaseRevision, FastHashMap, Graph, GraphError, RevlogError,
+    RevlogIndex, UncheckedRevision,
 };
 
 pub const INDEX_ENTRY_SIZE: usize = 64;
@@ -259,6 +260,9 @@
     offsets: RwLock<Option<Vec<usize>>>,
     uses_generaldelta: bool,
     is_inline: bool,
+    /// Cache of the head revisions in this index, kept in sync. Should
+    /// be accessed via the [`Self::head_revs`] method.
+    head_revs: Vec<Revision>,
 }
 
 impl Debug for Index {
@@ -358,6 +362,7 @@
                     offsets: RwLock::new(Some(offsets)),
                     uses_generaldelta,
                     is_inline: true,
+                    head_revs: vec![],
                 })
             } else {
                 Err(HgError::corrupted("unexpected inline revlog length"))
@@ -368,6 +373,7 @@
                 offsets: RwLock::new(None),
                 uses_generaldelta,
                 is_inline: false,
+                head_revs: vec![],
             })
         }
     }
@@ -512,6 +518,26 @@
         }
     }
 
+    /// Return the head revisions of this index
+    pub fn head_revs(&mut self) -> Result<Vec<Revision>, GraphError> {
+        if !self.head_revs.is_empty() {
+            return Ok(self.head_revs.to_owned());
+        }
+        let mut revs: HashSet<Revision, RandomState> = (0..self.len())
+            .into_iter()
+            .map(|i| Revision(i as BaseRevision))
+            .collect();
+        dagops::retain_heads(self, &mut revs)?;
+        if self.is_empty() {
+            revs.insert(NULL_REVISION);
+        }
+        let mut as_vec: Vec<Revision> =
+            revs.into_iter().map(Into::into).collect();
+        as_vec.sort_unstable();
+        self.head_revs = as_vec.to_owned();
+        Ok(as_vec)
+    }
+
     /// Obtain the delta chain for a revision.
     ///
     /// `stop_rev` specifies a revision to stop at. If not specified, we
@@ -599,6 +625,7 @@
             offsets.push(new_offset)
         }
         self.bytes.added.extend(revision_data.into_v1().as_bytes());
+        self.head_revs.clear();
         Ok(())
     }
 
@@ -612,6 +639,7 @@
         if let Some(offsets) = &mut *self.get_offsets_mut() {
             offsets.truncate(rev.0 as usize)
         }
+        self.head_revs.clear();
         Ok(())
     }
 
@@ -620,6 +648,7 @@
         // instead of offsets to determine whether we're inline since we might
         // clear caches. This implies re-populating the offsets on-demand.
         self.offsets = RwLock::new(None);
+        self.head_revs.clear();
     }
 
     /// Unchecked version of `is_snapshot`.
diff -r 050098d60c30 -r a7bba7df9189 rust/hg-cpython/src/revlog.rs
--- a/rust/hg-cpython/src/revlog.rs	Sat Sep 30 16:52:40 2023 +0200
+++ b/rust/hg-cpython/src/revlog.rs	Tue Sep 19 15:21:43 2023 +0200
@@ -13,9 +13,9 @@
 use cpython::{
     buffer::{Element, PyBuffer},
     exc::{IndexError, ValueError},
-    ObjectProtocol, PyBool, PyBytes, PyClone, PyDict, PyErr, PyInt, PyModule,
-    PyObject, PyResult, PySet, PyString, PyTuple, Python, PythonObject,
-    ToPyObject,
+    ObjectProtocol, PyBool, PyBytes, PyClone, PyDict, PyErr, PyInt, PyList,
+    PyModule, PyObject, PyResult, PySet, PyString, PyTuple, Python,
+    PythonObject, ToPyObject,
 };
 use hg::{
     errors::HgError,
@@ -250,7 +250,11 @@
 
     /// get head revisions
     def headrevs(&self, *args, **kw) -> PyResult<PyObject> {
-        self.call_cindex(py, "headrevs", args, kw)
+        let rust_res = self.inner_headrevs(py)?;
+
+        let c_res = self.call_cindex(py, "headrevs", args, kw)?;
+        assert_py_eq(py, "headrevs", &rust_res, &c_res)?;
+        Ok(rust_res)
     }
 
     /// get filtered head revisions
@@ -782,6 +786,17 @@
             ),
         })
     }
+
+    fn inner_headrevs(&self, py: Python) -> PyResult<PyObject> {
+        let index = &mut *self.index(py).borrow_mut();
+        let as_vec: Vec<PyObject> = index
+            .head_revs()
+            .map_err(|e| graph_error(py, e))?
+            .iter()
+            .map(|r| PyRevision::from(*r).into_py_object(py).into_object())
+            .collect();
+        Ok(PyList::new(py, &as_vec).into_object())
+    }
 }
 
 fn revlog_error(py: Python) -> PyErr {
@@ -797,6 +812,12 @@
     }
 }
 
+fn graph_error(py: Python, _err: hg::GraphError) -> PyErr {
+    // ParentOutOfRange is currently the only alternative
+    // in `hg::GraphError`. The C index always raises this simple ValueError.
+    PyErr::new::<ValueError, _>(py, "parent out of range")
+}
+
 fn nodemap_rev_not_in_index(py: Python, rev: UncheckedRevision) -> PyErr {
     PyErr::new::<ValueError, _>(
         py,