# HG changeset patch
# User Raphaël Gomès <rgomes@octobus.net>
# Date 1691067694 -7200
# Node ID 9b06e7f32bc5e94aaf0bfb8a59538fea518496be
# Parent  b8c89957a6b7709f69e215c0cb5796ad33360ef3
rust-index: add support for `find_snapshots`

diff -r b8c89957a6b7 -r 9b06e7f32bc5 rust/hg-core/src/revlog/index.rs
--- a/rust/hg-core/src/revlog/index.rs	Thu Aug 03 12:05:32 2023 +0200
+++ b/rust/hg-core/src/revlog/index.rs	Thu Aug 03 15:01:34 2023 +0200
@@ -1,3 +1,4 @@
+use std::collections::HashSet;
 use std::fmt::Debug;
 use std::ops::Deref;
 use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
@@ -10,7 +11,10 @@
 use crate::node::{NODE_BYTES_LENGTH, NULL_NODE, STORED_NODE_ID_BYTES};
 use crate::revlog::node::Node;
 use crate::revlog::{Revision, NULL_REVISION};
-use crate::{Graph, GraphError, RevlogError, RevlogIndex, UncheckedRevision};
+use crate::{
+    BaseRevision, FastHashMap, Graph, GraphError, RevlogError, RevlogIndex,
+    UncheckedRevision,
+};
 
 pub const INDEX_ENTRY_SIZE: usize = 64;
 pub const COMPRESSION_MODE_INLINE: u8 = 2;
@@ -283,6 +287,35 @@
     }
 }
 
+/// A cache suitable for find_snapshots
+///
+/// Logically equivalent to a mapping whose keys are [`BaseRevision`] and
+/// values sets of [`BaseRevision`]
+///
+/// TODO the dubious part is insisting that errors must be RevlogError
+/// we would probably need to sprinkle some magic here, such as an associated
+/// type that would be Into<RevlogError> but even that would not be
+/// satisfactory, as errors potentially have nothing to do with the revlog.
+pub trait SnapshotsCache {
+    fn insert_for(
+        &mut self,
+        rev: BaseRevision,
+        value: BaseRevision,
+    ) -> Result<(), RevlogError>;
+}
+
+impl SnapshotsCache for FastHashMap<BaseRevision, HashSet<BaseRevision>> {
+    fn insert_for(
+        &mut self,
+        rev: BaseRevision,
+        value: BaseRevision,
+    ) -> Result<(), RevlogError> {
+        let all_values = self.entry(rev).or_insert_with(HashSet::new);
+        all_values.insert(value);
+        Ok(())
+    }
+}
+
 impl Index {
     /// Create an index from bytes.
     /// Calculate the start of each entry when is_inline is true.
@@ -479,6 +512,38 @@
         }
     }
 
+    pub fn find_snapshots(
+        &self,
+        start_rev: UncheckedRevision,
+        end_rev: UncheckedRevision,
+        cache: &mut impl SnapshotsCache,
+    ) -> Result<(), RevlogError> {
+        let mut start_rev = start_rev.0;
+        let mut end_rev = end_rev.0;
+        end_rev += 1;
+        let len = self.len().try_into().unwrap();
+        if end_rev > len {
+            end_rev = len;
+        }
+        if start_rev < 0 {
+            start_rev = 0;
+        }
+        for rev in start_rev..end_rev {
+            if !self.is_snapshot_unchecked(Revision(rev))? {
+                continue;
+            }
+            let mut base = self
+                .get_entry(Revision(rev))
+                .unwrap()
+                .base_revision_or_base_of_delta_chain();
+            if base.0 == rev {
+                base = NULL_REVISION.into();
+            }
+            cache.insert_for(base.0, rev)?;
+        }
+        Ok(())
+    }
+
     /// TODO move this to the trait probably, along with other things
     pub fn append(
         &mut self,
diff -r b8c89957a6b7 -r 9b06e7f32bc5 rust/hg-cpython/src/revlog.rs
--- a/rust/hg-cpython/src/revlog.rs	Thu Aug 03 12:05:32 2023 +0200
+++ b/rust/hg-cpython/src/revlog.rs	Thu Aug 03 15:01:34 2023 +0200
@@ -14,12 +14,14 @@
     buffer::{Element, PyBuffer},
     exc::{IndexError, ValueError},
     ObjectProtocol, PyBool, PyBytes, PyClone, PyDict, PyErr, PyInt, PyModule,
-    PyObject, PyResult, PyString, PyTuple, Python, PythonObject, ToPyObject,
+    PyObject, PyResult, PySet, PyString, PyTuple, Python, PythonObject,
+    ToPyObject,
 };
 use hg::{
-    index::{IndexHeader, RevisionDataParams},
+    errors::HgError,
+    index::{IndexHeader, RevisionDataParams, SnapshotsCache},
     nodemap::{Block, NodeMapError, NodeTree},
-    revlog::{nodemap::NodeMap, NodePrefix, RevlogIndex},
+    revlog::{nodemap::NodeMap, NodePrefix, RevlogError, RevlogIndex},
     BaseRevision, Revision, UncheckedRevision, NULL_REVISION,
 };
 use std::cell::RefCell;
@@ -271,7 +273,39 @@
 
     /// Gather snapshot data in a cache dict
     def findsnapshots(&self, *args, **kw) -> PyResult<PyObject> {
-        self.call_cindex(py, "findsnapshots", args, kw)
+        let index = self.index(py).borrow();
+        let cache: PyDict = args.get_item(py, 0).extract(py)?;
+        // this methods operates by setting new values in the cache,
+        // hence we will compare results by letting the C implementation
+        // operate over a deepcopy of the cache, and finally compare both
+        // caches.
+        let c_cache = PyDict::new(py);
+        for (k, v) in cache.items(py) {
+            c_cache.set_item(py, k, PySet::new(py, v)?)?;
+        }
+
+        let start_rev = UncheckedRevision(args.get_item(py, 1).extract(py)?);
+        let end_rev = UncheckedRevision(args.get_item(py, 2).extract(py)?);
+        let mut cache_wrapper = PySnapshotsCache{ py, dict: cache };
+        index.find_snapshots(
+            start_rev,
+            end_rev,
+            &mut cache_wrapper,
+        ).map_err(|_| revlog_error(py))?;
+
+        let c_args = PyTuple::new(
+            py,
+            &[
+                c_cache.clone_ref(py).into_object(),
+                args.get_item(py, 1),
+                args.get_item(py, 2)
+            ]
+        );
+        self.call_cindex(py, "findsnapshots", &c_args, kw)?;
+        assert_py_eq(py, "findsnapshots cache",
+                     &cache_wrapper.into_object(),
+                     &c_cache.into_object())?;
+        Ok(py.None())
     }
 
     /// determine revisions with deltas to reconstruct fulltext
@@ -487,6 +521,39 @@
     )
 }
 
+struct PySnapshotsCache<'p> {
+    py: Python<'p>,
+    dict: PyDict,
+}
+
+impl<'p> PySnapshotsCache<'p> {
+    fn into_object(self) -> PyObject {
+        self.dict.into_object()
+    }
+}
+
+impl<'p> SnapshotsCache for PySnapshotsCache<'p> {
+    fn insert_for(
+        &mut self,
+        rev: BaseRevision,
+        value: BaseRevision,
+    ) -> Result<(), RevlogError> {
+        let pyvalue = value.into_py_object(self.py).into_object();
+        match self.dict.get_item(self.py, rev) {
+            Some(obj) => obj
+                .extract::<PySet>(self.py)
+                .and_then(|set| set.add(self.py, pyvalue)),
+            None => PySet::new(self.py, vec![pyvalue])
+                .and_then(|set| self.dict.set_item(self.py, rev, set)),
+        }
+        .map_err(|_| {
+            RevlogError::Other(HgError::unsupported(
+                "Error in Python caches handling",
+            ))
+        })
+    }
+}
+
 impl MixedIndex {
     fn new(
         py: Python,