diff rust/hg-cpython/src/revlog.rs @ 51233:9b06e7f32bc5

rust-index: add support for `find_snapshots`
author Rapha?l Gom?s <rgomes@octobus.net>
date Thu, 03 Aug 2023 15:01:34 +0200
parents b8c89957a6b7
children 62e39bef36ca
line wrap: on
line diff
--- 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,