rust-dirstate: improve API of `DirsMultiset`

    - Use opaque `Iterator` type instead of implementation-specific one from
      `HashMap`
    - Make `DirsMultiset` behave like a set both in Rust and from Python

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

rust/hg-core/src/dirstate/dirs_multiset.rs

@@ -1,341 +1,341 b''
 // dirs_multiset.rs
 //
 // Copyright 2019 Raphaël Gomès <[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.
 
 //! A multiset of directory names.
 //!
 //! Used to counts the references to directories in a manifest or dirstate.
 use crate::{
     dirstate::EntryState, utils::files, DirsIterable, DirstateEntry,
     DirstateMapError,
 };
-use std::collections::hash_map::{Entry, Iter};
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 
 #[derive(PartialEq, Debug)]
 pub struct DirsMultiset {
     inner: HashMap<Vec<u8>, u32>,
 }
 
 impl DirsMultiset {
     /// Initializes the multiset from a dirstate or a manifest.
     ///
     /// If `skip_state` is provided, skips dirstate entries with equal state.
     pub fn new(
         iterable: DirsIterable,
         skip_state: Option<EntryState>,
     ) -> Self {
         let mut multiset = DirsMultiset {
             inner: HashMap::new(),
         };
 
         match iterable {
             DirsIterable::Dirstate(vec) => {
                 for (filename, DirstateEntry { state, .. }) in vec {
                     // This `if` is optimized out of the loop
                     if let Some(skip) = skip_state {
                         if skip != *state {
                             multiset.add_path(filename);
                         }
                     } else {
                         multiset.add_path(filename);
                     }
                 }
             }
             DirsIterable::Manifest(vec) => {
                 for filename in vec {
                     multiset.add_path(filename);
                 }
             }
         }
 
         multiset
     }
 
     /// Increases the count of deepest directory contained in the path.
     ///
     /// If the directory is not yet in the map, adds its parents.
     pub fn add_path(&mut self, path: &[u8]) {
         for subpath in files::find_dirs(path) {
             if let Some(val) = self.inner.get_mut(subpath) {
                 *val += 1;
                 break;
             }
             self.inner.insert(subpath.to_owned(), 1);
         }
     }
 
     /// Decreases the count of deepest directory contained in the path.
     ///
     /// If it is the only reference, decreases all parents until one is
     /// removed.
     /// If the directory is not in the map, something horrible has happened.
     pub fn delete_path(
         &mut self,
         path: &[u8],
     ) -> Result<(), DirstateMapError> {
         for subpath in files::find_dirs(path) {
             match self.inner.entry(subpath.to_owned()) {
                 Entry::Occupied(mut entry) => {
                     let val = entry.get().clone();
                     if val > 1 {
                         entry.insert(val - 1);
                         break;
                     }
                     entry.remove();
                 }
                 Entry::Vacant(_) => {
                     return Err(DirstateMapError::PathNotFound(
                         path.to_owned(),
                     ))
                 }
             };
         }
 
         Ok(())
     }
 
-    pub fn contains_key(&self, key: &[u8]) -> bool {
+    pub fn contains(&self, key: &[u8]) -> bool {
         self.inner.contains_key(key)
     }
 
-    pub fn iter(&self) -> Iter<Vec<u8>, u32> {
-        self.inner.iter()
+    pub fn iter(&self) -> impl Iterator<Item = &Vec<u8>> {
+        self.inner.keys()
     }
 
     pub fn len(&self) -> usize {
         self.inner.len()
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use std::collections::HashMap;
 
     #[test]
     fn test_delete_path_path_not_found() {
         let mut map = DirsMultiset::new(DirsIterable::Manifest(&vec![]), None);
         let path = b"doesnotexist/";
         assert_eq!(
             Err(DirstateMapError::PathNotFound(path.to_vec())),
             map.delete_path(path)
         );
     }
 
     #[test]
     fn test_delete_path_empty_path() {
         let mut map =
             DirsMultiset::new(DirsIterable::Manifest(&vec![vec![]]), None);
         let path = b"";
         assert_eq!(Ok(()), map.delete_path(path));
         assert_eq!(
             Err(DirstateMapError::PathNotFound(path.to_vec())),
             map.delete_path(path)
         );
     }
 
     #[test]
     fn test_delete_path_successful() {
         let mut map = DirsMultiset {
             inner: [("", 5), ("a", 3), ("a/b", 2), ("a/c", 1)]
                 .iter()
                 .map(|(k, v)| (k.as_bytes().to_vec(), *v))
                 .collect(),
         };
 
         assert_eq!(Ok(()), map.delete_path(b"a/b/"));
         assert_eq!(Ok(()), map.delete_path(b"a/b/"));
         assert_eq!(
             Err(DirstateMapError::PathNotFound(b"a/b/".to_vec())),
             map.delete_path(b"a/b/")
         );
 
         assert_eq!(2, *map.inner.get(&b"a".to_vec()).unwrap());
         assert_eq!(1, *map.inner.get(&b"a/c".to_vec()).unwrap());
         eprintln!("{:?}", map);
         assert_eq!(Ok(()), map.delete_path(b"a/"));
         eprintln!("{:?}", map);
 
         assert_eq!(Ok(()), map.delete_path(b"a/c/"));
         assert_eq!(
             Err(DirstateMapError::PathNotFound(b"a/c/".to_vec())),
             map.delete_path(b"a/c/")
         );
     }
 
     #[test]
     fn test_add_path_empty_path() {
         let mut map = DirsMultiset::new(DirsIterable::Manifest(&vec![]), None);
         let path = b"";
         map.add_path(path);
 
         assert_eq!(1, map.len());
     }
 
     #[test]
     fn test_add_path_successful() {
         let mut map = DirsMultiset::new(DirsIterable::Manifest(&vec![]), None);
 
         map.add_path(b"a/");
         assert_eq!(1, *map.inner.get(&b"a".to_vec()).unwrap());
         assert_eq!(1, *map.inner.get(&Vec::new()).unwrap());
         assert_eq!(2, map.len());
 
         // Non directory should be ignored
         map.add_path(b"a");
         assert_eq!(1, *map.inner.get(&b"a".to_vec()).unwrap());
         assert_eq!(2, map.len());
 
         // Non directory will still add its base
         map.add_path(b"a/b");
         assert_eq!(2, *map.inner.get(&b"a".to_vec()).unwrap());
         assert_eq!(2, map.len());
 
         // Duplicate path works
         map.add_path(b"a/");
         assert_eq!(3, *map.inner.get(&b"a".to_vec()).unwrap());
 
         // Nested dir adds to its base
         map.add_path(b"a/b/");
         assert_eq!(4, *map.inner.get(&b"a".to_vec()).unwrap());
         assert_eq!(1, *map.inner.get(&b"a/b".to_vec()).unwrap());
 
         // but not its base's base, because it already existed
         map.add_path(b"a/b/c/");
         assert_eq!(4, *map.inner.get(&b"a".to_vec()).unwrap());
         assert_eq!(2, *map.inner.get(&b"a/b".to_vec()).unwrap());
 
         map.add_path(b"a/c/");
         assert_eq!(1, *map.inner.get(&b"a/c".to_vec()).unwrap());
 
         let expected = DirsMultiset {
             inner: [("", 2), ("a", 5), ("a/b", 2), ("a/b/c", 1), ("a/c", 1)]
                 .iter()
                 .map(|(k, v)| (k.as_bytes().to_vec(), *v))
                 .collect(),
         };
         assert_eq!(map, expected);
     }
 
     #[test]
     fn test_dirsmultiset_new_empty() {
         use DirsIterable::{Dirstate, Manifest};
 
         let new = DirsMultiset::new(Manifest(&vec![]), None);
         let expected = DirsMultiset {
             inner: HashMap::new(),
         };
         assert_eq!(expected, new);
 
         let new = DirsMultiset::new(Dirstate(&HashMap::new()), None);
         let expected = DirsMultiset {
             inner: HashMap::new(),
         };
         assert_eq!(expected, new);
     }
 
     #[test]
     fn test_dirsmultiset_new_no_skip() {
         use DirsIterable::{Dirstate, Manifest};
 
         let input_vec = ["a/", "b/", "a/c", "a/d/"]
             .iter()
             .map(|e| e.as_bytes().to_vec())
             .collect();
         let expected_inner = [("", 2), ("a", 3), ("b", 1), ("a/d", 1)]
             .iter()
             .map(|(k, v)| (k.as_bytes().to_vec(), *v))
             .collect();
 
         let new = DirsMultiset::new(Manifest(&input_vec), None);
         let expected = DirsMultiset {
             inner: expected_inner,
         };
         assert_eq!(expected, new);
 
         let input_map = ["a/", "b/", "a/c", "a/d/"]
             .iter()
             .map(|f| {
                 (
                     f.as_bytes().to_vec(),
                     DirstateEntry {
                         state: EntryState::Normal,
                         mode: 0,
                         mtime: 0,
                         size: 0,
                     },
                 )
             })
             .collect();
         let expected_inner = [("", 2), ("a", 3), ("b", 1), ("a/d", 1)]
             .iter()
             .map(|(k, v)| (k.as_bytes().to_vec(), *v))
             .collect();
 
         let new = DirsMultiset::new(Dirstate(&input_map), None);
         let expected = DirsMultiset {
             inner: expected_inner,
         };
         assert_eq!(expected, new);
     }
 
     #[test]
     fn test_dirsmultiset_new_skip() {
         use DirsIterable::{Dirstate, Manifest};
 
         let input_vec = ["a/", "b/", "a/c", "a/d/"]
             .iter()
             .map(|e| e.as_bytes().to_vec())
             .collect();
         let expected_inner = [("", 2), ("a", 3), ("b", 1), ("a/d", 1)]
             .iter()
             .map(|(k, v)| (k.as_bytes().to_vec(), *v))
             .collect();
 
         let new =
             DirsMultiset::new(Manifest(&input_vec), Some(EntryState::Normal));
         let expected = DirsMultiset {
             inner: expected_inner,
         };
         // Skip does not affect a manifest
         assert_eq!(expected, new);
 
         let input_map = [
             ("a/", EntryState::Normal),
             ("a/b/", EntryState::Normal),
             ("a/c", EntryState::Removed),
             ("a/d/", EntryState::Merged),
         ]
         .iter()
         .map(|(f, state)| {
             (
                 f.as_bytes().to_vec(),
                 DirstateEntry {
                     state: *state,
                     mode: 0,
                     mtime: 0,
                     size: 0,
                 },
             )
         })
         .collect();
 
         // "a" incremented with "a/c" and "a/d/"
         let expected_inner = [("", 1), ("a", 2), ("a/d", 1)]
             .iter()
             .map(|(k, v)| (k.as_bytes().to_vec(), *v))
             .collect();
 
         let new =
             DirsMultiset::new(Dirstate(&input_map), Some(EntryState::Normal));
         let expected = DirsMultiset {
             inner: expected_inner,
         };
         assert_eq!(expected, new);
     }
 
 }

rust/hg-cpython/src/dirstate/dirs_multiset.rs

@@ -1,118 +1,113 b''
 // dirs_multiset.rs
 //
 // Copyright 2019 Raphaël Gomès <[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::dirstate::dirs_multiset` file provided by the
 //! `hg-core` package.
 
 use std::cell::RefCell;
 
 use cpython::{
     exc, ObjectProtocol, PyBytes, PyDict, PyErr, PyObject, PyResult,
-    ToPyObject,
+    PythonObject, ToPyObject,
 };
 
 use crate::dirstate::extract_dirstate;
 use hg::{
     DirsIterable, DirsMultiset, DirstateMapError, DirstateParseError,
     EntryState,
 };
 use std::convert::TryInto;
 
 py_class!(pub class Dirs |py| {
     data dirs_map: RefCell<DirsMultiset>;
 
     // `map` is either a `dict` or a flat iterator (usually a `set`, sometimes
     // a `list`)
     def __new__(
         _cls,
         map: PyObject,
         skip: Option<PyObject> = None
     ) -> PyResult<Self> {
         let mut skip_state: Option<EntryState> = None;
         if let Some(skip) = skip {
             skip_state = Some(
                 skip.extract::<PyBytes>(py)?.data(py)[0]
                     .try_into()
                     .map_err(|e: DirstateParseError| {
                         PyErr::new::<exc::ValueError, _>(py, e.to_string())
                     })?,
             );
         }
         let inner = if let Ok(map) = map.cast_as::<PyDict>(py) {
             let dirstate = extract_dirstate(py, &map)?;
             DirsMultiset::new(
                 DirsIterable::Dirstate(&dirstate),
                 skip_state,
             )
         } else {
             let map: Result<Vec<Vec<u8>>, PyErr> = map
                 .iter(py)?
                 .map(|o| Ok(o?.extract::<PyBytes>(py)?.data(py).to_owned()))
                 .collect();
             DirsMultiset::new(
                 DirsIterable::Manifest(&map?),
                 skip_state,
             )
         };
 
         Self::create_instance(py, RefCell::new(inner))
     }
 
     def addpath(&self, path: PyObject) -> PyResult<PyObject> {
         self.dirs_map(py).borrow_mut().add_path(
             path.extract::<PyBytes>(py)?.data(py),
         );
         Ok(py.None())
     }
 
     def delpath(&self, path: PyObject) -> PyResult<PyObject> {
         self.dirs_map(py).borrow_mut().delete_path(
             path.extract::<PyBytes>(py)?.data(py),
         )
             .and(Ok(py.None()))
             .or_else(|e| {
                 match e {
                     DirstateMapError::PathNotFound(_p) => {
                         Err(PyErr::new::<exc::ValueError, _>(
                             py,
                             "expected a value, found none".to_string(),
                         ))
                     }
                     DirstateMapError::EmptyPath => {
                         Ok(py.None())
                     }
                 }
             })
     }
 
     // This is really inefficient on top of being ugly, but it's an easy way
     // of having it work to continue working on the rest of the module
     // hopefully bypassing Python entirely pretty soon.
     def __iter__(&self) -> PyResult<PyObject> {
-        let dict = PyDict::new(py);
-
-        for (key, value) in self.dirs_map(py).borrow().iter() {
-            dict.set_item(
-                py,
-                PyBytes::new(py, &key[..]),
-                value.to_py_object(py),
-            )?;
-        }
-
-        let locals = PyDict::new(py);
-        locals.set_item(py, "obj", dict)?;
-
-        py.eval("iter(obj)", None, Some(&locals))
+        let dirs = self.dirs_map(py).borrow();
+        let dirs: Vec<_> = dirs
+            .iter()
+            .map(|d| PyBytes::new(py, d))
+            .collect();
+        dirs.to_py_object(py)
+            .into_object()
+            .iter(py)
+            .map(|o| o.into_object())
     }
 
     def __contains__(&self, item: PyObject) -> PyResult<bool> {
         Ok(self
             .dirs_map(py)
             .borrow()
-            .contains_key(item.extract::<PyBytes>(py)?.data(py).as_ref()))
+            .contains(item.extract::<PyBytes>(py)?.data(py).as_ref()))
     }
 });