rust-discovery: implementing and exposing stats()

This time, it's simple enough that we can do it in all layers in
one shot.

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

rust/hg-core/src/discovery.rs

@@ -1,196 +1,209 b''
 // discovery.rs
 //
 // Copyright 2019 Georges Racinet <[e-mail]>
 //
 // This software may be used and distributed according to the terms of the
 // GNU General Public License version 2 or any later version.
 
 //! Discovery operations
 //!
 //! This is a Rust counterpart to the `partialdiscovery` class of
 //! `mercurial.setdiscovery`
 
 use super::{Graph, GraphError, Revision};
 use crate::ancestors::MissingAncestors;
 use crate::dagops;
 use std::collections::HashSet;
 
 pub struct PartialDiscovery<G: Graph + Clone> {
     target_heads: Option<Vec<Revision>>,
     graph: G, // plays the role of self._repo
     common: MissingAncestors<G>,
     undecided: Option<HashSet<Revision>>,
     missing: HashSet<Revision>,
 }
 
+pub struct DiscoveryStats {
+    pub undecided: Option<usize>,
+}
+
 impl<G: Graph + Clone> PartialDiscovery<G> {
     /// Create a PartialDiscovery object, with the intent
     /// of comparing our `::<target_heads>` revset to the contents of another
     /// repo.
     ///
     /// For now `target_heads` is passed as a vector, and will be used
     /// at the first call to `ensure_undecided()`.
     ///
     /// If we want to make the signature more flexible,
     /// we'll have to make it a type argument of `PartialDiscovery` or a trait
     /// object since we'll keep it in the meanwhile
     pub fn new(graph: G, target_heads: Vec<Revision>) -> Self {
         PartialDiscovery {
             undecided: None,
             target_heads: Some(target_heads),
             graph: graph.clone(),
             common: MissingAncestors::new(graph, vec![]),
             missing: HashSet::new(),
         }
     }
 
     /// Register revisions known as being common
     pub fn add_common_revisions(
         &mut self,
         common: impl IntoIterator<Item = Revision>,
     ) -> Result<(), GraphError> {
         self.common.add_bases(common);
         if let Some(ref mut undecided) = self.undecided {
             self.common.remove_ancestors_from(undecided)?;
         }
         Ok(())
     }
 
     /// Register revisions known as being missing
     pub fn add_missing_revisions(
         &mut self,
         missing: impl IntoIterator<Item = Revision>,
     ) -> Result<(), GraphError> {
         self.ensure_undecided()?;
         let range = dagops::range(
             &self.graph,
             missing,
             self.undecided.as_ref().unwrap().iter().cloned(),
         )?;
         let undecided_mut = self.undecided.as_mut().unwrap();
         for missrev in range {
             self.missing.insert(missrev);
             undecided_mut.remove(&missrev);
         }
         Ok(())
     }
 
     /// Do we have any information about the peer?
     pub fn has_info(&self) -> bool {
         self.common.has_bases()
     }
 
     /// Did we acquire full knowledge of our Revisions that the peer has?
     pub fn is_complete(&self) -> bool {
         self.undecided.as_ref().map_or(false, |s| s.is_empty())
     }
 
     /// Return the heads of the currently known common set of revisions.
     ///
     /// If the discovery process is not complete (see `is_complete()`), the
     /// caller must be aware that this is an intermediate state.
     ///
     /// On the other hand, if it is complete, then this is currently
     /// the only way to retrieve the end results of the discovery process.
     ///
     /// We may introduce in the future an `into_common_heads` call that
     /// would be more appropriate for normal Rust callers, dropping `self`
     /// if it is complete.
     pub fn common_heads(&self) -> Result<HashSet<Revision>, GraphError> {
         self.common.bases_heads()
     }
 
     /// Force first computation of `self.undecided`
     ///
     /// After this, `self.undecided.as_ref()` and `.as_mut()` can be
     /// unwrapped to get workable immutable or mutable references without
     /// any panic.
     ///
     /// This is an imperative call instead of an access with added lazyness
     /// to reduce easily the scope of mutable borrow for the caller,
     /// compared to undecided(&'a mut self) -> &'a… that would keep it
     /// as long as the resulting immutable one.
     fn ensure_undecided(&mut self) -> Result<(), GraphError> {
         if self.undecided.is_some() {
             return Ok(());
         }
         let tgt = self.target_heads.take().unwrap();
         self.undecided =
             Some(self.common.missing_ancestors(tgt)?.into_iter().collect());
         Ok(())
     }
+
+    /// Provide statistics about the current state of the discovery process
+    pub fn stats(&self) -> DiscoveryStats {
+        DiscoveryStats {
+            undecided: self.undecided.as_ref().map(|s| s.len()),
+        }
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::testing::SampleGraph;
 
     /// A PartialDiscovery as for pushing all the heads of `SampleGraph`
     fn full_disco() -> PartialDiscovery<SampleGraph> {
         PartialDiscovery::new(SampleGraph, vec![10, 11, 12, 13])
     }
 
     fn sorted_undecided(
         disco: &PartialDiscovery<SampleGraph>,
     ) -> Vec<Revision> {
         let mut as_vec: Vec<Revision> =
             disco.undecided.as_ref().unwrap().iter().cloned().collect();
         as_vec.sort();
         as_vec
     }
 
     fn sorted_missing(disco: &PartialDiscovery<SampleGraph>) -> Vec<Revision> {
         let mut as_vec: Vec<Revision> =
             disco.missing.iter().cloned().collect();
         as_vec.sort();
         as_vec
     }
 
     fn sorted_common_heads(
         disco: &PartialDiscovery<SampleGraph>,
     ) -> Result<Vec<Revision>, GraphError> {
         let mut as_vec: Vec<Revision> =
             disco.common_heads()?.iter().cloned().collect();
         as_vec.sort();
         Ok(as_vec)
     }
 
     #[test]
     fn test_add_common_get_undecided() -> Result<(), GraphError> {
         let mut disco = full_disco();
         assert_eq!(disco.undecided, None);
         assert!(!disco.has_info());
+        assert_eq!(disco.stats().undecided, None);
 
         disco.add_common_revisions(vec![11, 12])?;
         assert!(disco.has_info());
         assert!(!disco.is_complete());
         assert!(disco.missing.is_empty());
 
         // add_common_revisions did not trigger a premature computation
         // of `undecided`, let's check that and ask for them
         assert_eq!(disco.undecided, None);
         disco.ensure_undecided()?;
         assert_eq!(sorted_undecided(&disco), vec![5, 8, 10, 13]);
+        assert_eq!(disco.stats().undecided, Some(4));
         Ok(())
     }
 
     /// in this test, we pretend that our peer misses exactly (8+10)::
     /// and we're comparing all our repo to it (as in a bare push)
     #[test]
     fn test_discovery() -> Result<(), GraphError> {
         let mut disco = full_disco();
         disco.add_common_revisions(vec![11, 12])?;
         disco.add_missing_revisions(vec![8, 10])?;
         assert_eq!(sorted_undecided(&disco), vec![5]);
         assert_eq!(sorted_missing(&disco), vec![8, 10, 13]);
         assert!(!disco.is_complete());
 
         disco.add_common_revisions(vec![5])?;
         assert_eq!(sorted_undecided(&disco), vec![]);
         assert_eq!(sorted_missing(&disco), vec![8, 10, 13]);
         assert!(disco.is_complete());
         assert_eq!(sorted_common_heads(&disco)?, vec![5, 11, 12]);
         Ok(())
     }
 }

rust/hg-cpython/src/discovery.rs

@@ -1,114 +1,125 b''
 // discovery.rs
 //
 // Copyright 2018 Georges Racinet <[e-mail]>
 //
 // This software may be used and distributed according to the terms of the
 // GNU General Public License version 2 or any later version.
 
 //! Bindings for the `hg::discovery` module provided by the
 //! `hg-core` crate. From Python, this will be seen as `rustext.discovery`
 //!
 //! # Classes visible from Python:
 //! - [`PartialDiscover`] is the Rust implementation of
 //!   `mercurial.setdiscovery.partialdiscovery`.
 
 use crate::conversion::{py_set, rev_pyiter_collect};
 use cindex::Index;
-use cpython::{ObjectProtocol, PyDict, PyModule, PyObject, PyResult, Python};
+use cpython::{
+    ObjectProtocol, PyDict, PyModule, PyObject, PyResult, Python, ToPyObject,
+};
 use exceptions::GraphError;
 use hg::discovery::PartialDiscovery as CorePartialDiscovery;
 use hg::Revision;
 
 use std::cell::RefCell;
 
 py_class!(pub class PartialDiscovery |py| {
     data inner: RefCell<Box<CorePartialDiscovery<Index>>>;
 
     def __new__(
         _cls,
         index: PyObject,
         targetheads: PyObject
     ) -> PyResult<PartialDiscovery> {
         Self::create_instance(
             py,
             RefCell::new(Box::new(CorePartialDiscovery::new(
                 Index::new(py, index)?,
                 rev_pyiter_collect(py, &targetheads)?,
             )))
         )
     }
 
     def addcommons(&self, commons: PyObject) -> PyResult<PyObject> {
         let mut inner = self.inner(py).borrow_mut();
         let commons_vec: Vec<Revision> = rev_pyiter_collect(py, &commons)?;
         inner.add_common_revisions(commons_vec)
             .map_err(|e| GraphError::pynew(py, e))?;
         Ok(py.None())
     }
 
     def addmissings(&self, missings: PyObject) -> PyResult<PyObject> {
         let mut inner = self.inner(py).borrow_mut();
         let missings_vec: Vec<Revision> = rev_pyiter_collect(py, &missings)?;
         inner.add_missing_revisions(missings_vec)
             .map_err(|e| GraphError::pynew(py, e))?;
         Ok(py.None())
     }
 
     def addinfo(&self, sample: PyObject) -> PyResult<PyObject> {
         let mut missing: Vec<Revision> = Vec::new();
         let mut common: Vec<Revision> = Vec::new();
         for info in sample.iter(py)? { // info is a pair (Revision, bool)
             let mut revknown = info?.iter(py)?;
             let rev: Revision = revknown.next().unwrap()?.extract(py)?;
             let known: bool = revknown.next().unwrap()?.extract(py)?;
             if known {
                 common.push(rev);
             } else {
                 missing.push(rev);
             }
         }
         let mut inner = self.inner(py).borrow_mut();
         inner.add_common_revisions(common)
             .map_err(|e| GraphError::pynew(py, e))?;
         inner.add_missing_revisions(missing)
             .map_err(|e| GraphError::pynew(py, e))?;
         Ok(py.None())
     }
 
     def hasinfo(&self) -> PyResult<bool> {
         Ok(self.inner(py).borrow().has_info())
     }
 
     def iscomplete(&self) -> PyResult<bool> {
         Ok(self.inner(py).borrow().is_complete())
     }
 
+    def stats(&self) -> PyResult<PyDict> {
+        let stats = self.inner(py).borrow().stats();
+        let as_dict: PyDict = PyDict::new(py);
+        as_dict.set_item(py, "undecided",
+                         stats.undecided.map(|l| l.to_py_object(py))
+                              .unwrap_or_else(|| py.None()))?;
+        Ok(as_dict)
+    }
+
     def commonheads(&self) -> PyResult<PyObject> {
         py_set(
             py,
             &self.inner(py).borrow().common_heads()
                 .map_err(|e| GraphError::pynew(py, e))?
         )
     }
 });
 
 /// Create the module, with __package__ given from parent
 pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
     let dotted_name = &format!("{}.discovery", package);
     let m = PyModule::new(py, dotted_name)?;
     m.add(py, "__package__", package)?;
     m.add(
         py,
         "__doc__",
         "Discovery of common node sets - Rust implementation",
     )?;
     m.add_class::<PartialDiscovery>(py)?;
 
     let sys = PyModule::import(py, "sys")?;
     let sys_modules: PyDict = sys.get(py, "modules")?.extract(py)?;
     sys_modules.set_item(py, dotted_name, &m)?;
     // Example C code (see pyexpat.c and import.c) will "give away the
     // reference", but we won't because it will be consumed once the
     // Rust PyObject is dropped.
     Ok(m)
 }

tests/test-rust-discovery.py

@@ -1,103 +1,111 b''
 from __future__ import absolute_import
 import unittest
 
 try:
     from mercurial import rustext
     rustext.__name__  # trigger immediate actual import
 except ImportError:
     rustext = None
 else:
     # this would fail already without appropriate ancestor.__package__
     from mercurial.rustext.discovery import (
         PartialDiscovery,
     )
 
 try:
     from mercurial.cext import parsers as cparsers
 except ImportError:
     cparsers = None
 
 # picked from test-parse-index2, copied rather than imported
 # so that it stays stable even if test-parse-index2 changes or disappears.
 data_non_inlined = (
     b'\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01D\x19'
     b'\x00\x07e\x12\x00\x00\x00\x00\x00\x00\x00\x00\xff\xff\xff\xff'
     b'\xff\xff\xff\xff\xd1\xf4\xbb\xb0\xbe\xfc\x13\xbd\x8c\xd3\x9d'
     b'\x0f\xcd\xd9;\x8c\x07\x8cJ/\x00\x00\x00\x00\x00\x00\x00\x00\x00'
     b'\x00\x00\x00\x00\x00\x00\x01D\x19\x00\x00\x00\x00\x00\xdf\x00'
     b'\x00\x01q\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\xff'
     b'\xff\xff\xff\xc1\x12\xb9\x04\x96\xa4Z1t\x91\xdfsJ\x90\xf0\x9bh'
     b'\x07l&\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
     b'\x00\x01D\xf8\x00\x00\x00\x00\x01\x1b\x00\x00\x01\xb8\x00\x00'
     b'\x00\x01\x00\x00\x00\x02\x00\x00\x00\x01\xff\xff\xff\xff\x02\n'
     b'\x0e\xc6&\xa1\x92\xae6\x0b\x02i\xfe-\xe5\xbao\x05\xd1\xe7\x00'
     b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01F'
     b'\x13\x00\x00\x00\x00\x01\xec\x00\x00\x03\x06\x00\x00\x00\x01'
     b'\x00\x00\x00\x03\x00\x00\x00\x02\xff\xff\xff\xff\x12\xcb\xeby1'
     b'\xb6\r\x98B\xcb\x07\xbd`\x8f\x92\xd9\xc4\x84\xbdK\x00\x00\x00'
     b'\x00\x00\x00\x00\x00\x00\x00\x00\x00'
     )
 
 
 @unittest.skipIf(rustext is None or cparsers is None,
                  "rustext or the C Extension parsers module "
                  "discovery relies on is not available")
 class rustdiscoverytest(unittest.TestCase):
     """Test the correctness of binding to Rust code.
 
     This test is merely for the binding to Rust itself: extraction of
     Python variable, giving back the results etc.
 
     It is not meant to test the algorithmic correctness of the provided
     methods. Hence the very simple embedded index data is good enough.
 
     Algorithmic correctness is asserted by the Rust unit tests.
     """
 
     def parseindex(self):
         return cparsers.parse_index2(data_non_inlined, False)[0]
 
     def testindex(self):
         idx = self.parseindex()
         # checking our assumptions about the index binary data:
         self.assertEqual({i: (r[5], r[6]) for i, r in enumerate(idx)},
                          {0: (-1, -1),
                           1: (0, -1),
                           2: (1, -1),
                           3: (2, -1)})
 
     def testaddcommonsmissings(self):
         idx = self.parseindex()
         disco = PartialDiscovery(idx, [3])
         self.assertFalse(disco.hasinfo())
         self.assertFalse(disco.iscomplete())
 
         disco.addcommons([1])
         self.assertTrue(disco.hasinfo())
         self.assertFalse(disco.iscomplete())
 
         disco.addmissings([2])
         self.assertTrue(disco.hasinfo())
         self.assertTrue(disco.iscomplete())
 
         self.assertEqual(disco.commonheads(), {1})
 
+    def testaddmissingsstats(self):
+        idx = self.parseindex()
+        disco = PartialDiscovery(idx, [3])
+        self.assertIsNone(disco.stats()['undecided'], None)
+
+        disco.addmissings([2])
+        self.assertEqual(disco.stats()['undecided'], 2)
+
     def testaddinfocommonfirst(self):
         idx = self.parseindex()
         disco = PartialDiscovery(idx, [3])
         disco.addinfo([(1, True), (2, False)])
         self.assertTrue(disco.hasinfo())
         self.assertTrue(disco.iscomplete())
         self.assertEqual(disco.commonheads(), {1})
 
     def testaddinfomissingfirst(self):
         idx = self.parseindex()
         disco = PartialDiscovery(idx, [3])
         disco.addinfo([(2, False), (1, True)])
         self.assertTrue(disco.hasinfo())
         self.assertTrue(disco.iscomplete())
         self.assertEqual(disco.commonheads(), {1})
 
 if __name__ == '__main__':
     import silenttestrunner
     silenttestrunner.main(__name__)