view rust/hg-core/src/dirstate/dirs_multiset.rs @ 47093:787ff5d21bcd

dirstate-tree: Make Rust DirstateMap bindings go through a trait object

This changeset starts a series that adds an experiment to make status faster
by changing the dirstate (first only in memory and later also on disk) to
be shaped as a tree matching the directory structure, instead of the current
flat collection of entries. The status algorithm can then traverse this tree
dirstate at the same time as it traverses the filesystem.

We (Octobus) have made prototypes that show promising results but are prone
to bitrot. We would like to start upstreaming some experimental Rust code that
goes in this direction, but to avoid disrupting users it should only be
enabled by some run-time opt-in while keeping the existing dirstate structure
and status algorithm as-is.

The `DirstateMap` type and `status` function look like the appropriate
boundary. This adds a new trait that abstracts everything Python bindings need
and makes those bindings go through a `dyn` trait object. Later we’ll have two
implementations of this trait, and the same bindings can use either.

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

// 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,
        hg_path::{HgPath, HgPathBuf, HgPathError},
    },
    DirstateEntry, DirstateMapError, FastHashMap,
};
use std::collections::{hash_map, hash_map::Entry, HashMap, HashSet};

// could be encapsulated if we care API stability more seriously
pub type DirsMultisetIter<'a> = hash_map::Keys<'a, HgPathBuf, u32>;

#[derive(PartialEq, Debug)]
pub struct DirsMultiset {
    inner: FastHashMap<HgPathBuf, u32>,
}

impl DirsMultiset {
    /// Initializes the multiset from a dirstate.
    ///
    /// If `skip_state` is provided, skips dirstate entries with equal state.
    pub fn from_dirstate<'a>(
        dirstate: impl IntoIterator<Item = (&'a HgPathBuf, &'a DirstateEntry)>,
        skip_state: Option<EntryState>,
    ) -> Result<Self, DirstateMapError> {
        let mut multiset = DirsMultiset {
            inner: FastHashMap::default(),
        };
        for (filename, DirstateEntry { state, .. }) in dirstate {
            // 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)?;
            }
        }

        Ok(multiset)
    }

    /// Initializes the multiset from a manifest.
    pub fn from_manifest(
        manifest: &[impl AsRef<HgPath>],
    ) -> Result<Self, DirstateMapError> {
        let mut multiset = DirsMultiset {
            inner: FastHashMap::default(),
        };

        for filename in manifest {
            multiset.add_path(filename.as_ref())?;
        }

        Ok(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: impl AsRef<HgPath>,
    ) -> Result<(), DirstateMapError> {
        for subpath in files::find_dirs(path.as_ref()) {
            if subpath.as_bytes().last() == Some(&b'/') {
                // TODO Remove this once PathAuditor is certified
                // as the only entrypoint for path data
                let second_slash_index = subpath.len() - 1;

                return Err(DirstateMapError::InvalidPath(
                    HgPathError::ConsecutiveSlashes {
                        bytes: path.as_ref().as_bytes().to_owned(),
                        second_slash_index,
                    },
                ));
            }
            if let Some(val) = self.inner.get_mut(subpath) {
                *val += 1;
                break;
            }
            self.inner.insert(subpath.to_owned(), 1);
        }
        Ok(())
    }

    /// 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: impl AsRef<HgPath>,
    ) -> Result<(), DirstateMapError> {
        for subpath in files::find_dirs(path.as_ref()) {
            match self.inner.entry(subpath.to_owned()) {
                Entry::Occupied(mut entry) => {
                    let val = *entry.get();
                    if val > 1 {
                        entry.insert(val - 1);
                        break;
                    }
                    entry.remove();
                }
                Entry::Vacant(_) => {
                    return Err(DirstateMapError::PathNotFound(
                        path.as_ref().to_owned(),
                    ))
                }
            };
        }

        Ok(())
    }

    pub fn contains(&self, key: impl AsRef<HgPath>) -> bool {
        self.inner.contains_key(key.as_ref())
    }

    pub fn iter(&self) -> DirsMultisetIter {
        self.inner.keys()
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// This is basically a reimplementation of `DirsMultiset` that stores the
/// children instead of just a count of them, plus a small optional
/// optimization to avoid some directories we don't need.
#[derive(PartialEq, Debug)]
pub struct DirsChildrenMultiset<'a> {
    inner: FastHashMap<&'a HgPath, HashSet<&'a HgPath>>,
    only_include: Option<HashSet<&'a HgPath>>,
}

impl<'a> DirsChildrenMultiset<'a> {
    pub fn new(
        paths: impl Iterator<Item = &'a HgPathBuf>,
        only_include: Option<&'a HashSet<impl AsRef<HgPath> + 'a>>,
    ) -> Self {
        let mut new = Self {
            inner: HashMap::default(),
            only_include: only_include
                .map(|s| s.iter().map(AsRef::as_ref).collect()),
        };

        for path in paths {
            new.add_path(path)
        }

        new
    }
    fn add_path(&mut self, path: &'a (impl AsRef<HgPath> + 'a)) {
        if path.as_ref().is_empty() {
            return;
        }
        for (directory, basename) in files::find_dirs_with_base(path.as_ref())
        {
            if !self.is_dir_included(directory) {
                continue;
            }
            self.inner
                .entry(directory)
                .and_modify(|e| {
                    e.insert(basename);
                })
                .or_insert_with(|| {
                    let mut set = HashSet::new();
                    set.insert(basename);
                    set
                });
        }
    }
    fn is_dir_included(&self, dir: impl AsRef<HgPath>) -> bool {
        match &self.only_include {
            None => false,
            Some(i) => i.contains(dir.as_ref()),
        }
    }

    pub fn get(
        &self,
        path: impl AsRef<HgPath>,
    ) -> Option<&HashSet<&'a HgPath>> {
        self.inner.get(path.as_ref())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::StateMap;

    #[test]
    fn test_delete_path_path_not_found() {
        let manifest: Vec<HgPathBuf> = vec![];
        let mut map = DirsMultiset::from_manifest(&manifest).unwrap();
        let path = HgPathBuf::from_bytes(b"doesnotexist/");
        assert_eq!(
            Err(DirstateMapError::PathNotFound(path.to_owned())),
            map.delete_path(&path)
        );
    }

    #[test]
    fn test_delete_path_empty_path() {
        let mut map =
            DirsMultiset::from_manifest(&vec![HgPathBuf::new()]).unwrap();
        let path = HgPath::new(b"");
        assert_eq!(Ok(()), map.delete_path(path));
        assert_eq!(
            Err(DirstateMapError::PathNotFound(path.to_owned())),
            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)| (HgPathBuf::from_bytes(k.as_bytes()), *v))
                .collect(),
        };

        assert_eq!(Ok(()), map.delete_path(HgPath::new(b"a/b/")));
        eprintln!("{:?}", map);
        assert_eq!(Ok(()), map.delete_path(HgPath::new(b"a/b/")));
        eprintln!("{:?}", map);
        assert_eq!(
            Err(DirstateMapError::PathNotFound(HgPathBuf::from_bytes(
                b"a/b/"
            ))),
            map.delete_path(HgPath::new(b"a/b/"))
        );

        assert_eq!(2, *map.inner.get(HgPath::new(b"a")).unwrap());
        assert_eq!(1, *map.inner.get(HgPath::new(b"a/c")).unwrap());
        eprintln!("{:?}", map);
        assert_eq!(Ok(()), map.delete_path(HgPath::new(b"a/")));
        eprintln!("{:?}", map);

        assert_eq!(Ok(()), map.delete_path(HgPath::new(b"a/c/")));
        assert_eq!(
            Err(DirstateMapError::PathNotFound(HgPathBuf::from_bytes(
                b"a/c/"
            ))),
            map.delete_path(HgPath::new(b"a/c/"))
        );
    }

    #[test]
    fn test_add_path_empty_path() {
        let manifest: Vec<HgPathBuf> = vec![];
        let mut map = DirsMultiset::from_manifest(&manifest).unwrap();
        let path = HgPath::new(b"");
        map.add_path(path).unwrap();

        assert_eq!(1, map.len());
    }

    #[test]
    fn test_add_path_successful() {
        let manifest: Vec<HgPathBuf> = vec![];
        let mut map = DirsMultiset::from_manifest(&manifest).unwrap();

        map.add_path(HgPath::new(b"a/")).unwrap();
        assert_eq!(1, *map.inner.get(HgPath::new(b"a")).unwrap());
        assert_eq!(1, *map.inner.get(HgPath::new(b"")).unwrap());
        assert_eq!(2, map.len());

        // Non directory should be ignored
        map.add_path(HgPath::new(b"a")).unwrap();
        assert_eq!(1, *map.inner.get(HgPath::new(b"a")).unwrap());
        assert_eq!(2, map.len());

        // Non directory will still add its base
        map.add_path(HgPath::new(b"a/b")).unwrap();
        assert_eq!(2, *map.inner.get(HgPath::new(b"a")).unwrap());
        assert_eq!(2, map.len());

        // Duplicate path works
        map.add_path(HgPath::new(b"a/")).unwrap();
        assert_eq!(3, *map.inner.get(HgPath::new(b"a")).unwrap());

        // Nested dir adds to its base
        map.add_path(HgPath::new(b"a/b/")).unwrap();
        assert_eq!(4, *map.inner.get(HgPath::new(b"a")).unwrap());
        assert_eq!(1, *map.inner.get(HgPath::new(b"a/b")).unwrap());

        // but not its base's base, because it already existed
        map.add_path(HgPath::new(b"a/b/c/")).unwrap();
        assert_eq!(4, *map.inner.get(HgPath::new(b"a")).unwrap());
        assert_eq!(2, *map.inner.get(HgPath::new(b"a/b")).unwrap());

        map.add_path(HgPath::new(b"a/c/")).unwrap();
        assert_eq!(1, *map.inner.get(HgPath::new(b"a/c")).unwrap());

        let expected = DirsMultiset {
            inner: [("", 2), ("a", 5), ("a/b", 2), ("a/b/c", 1), ("a/c", 1)]
                .iter()
                .map(|(k, v)| (HgPathBuf::from_bytes(k.as_bytes()), *v))
                .collect(),
        };
        assert_eq!(map, expected);
    }

    #[test]
    fn test_dirsmultiset_new_empty() {
        let manifest: Vec<HgPathBuf> = vec![];
        let new = DirsMultiset::from_manifest(&manifest).unwrap();
        let expected = DirsMultiset {
            inner: FastHashMap::default(),
        };
        assert_eq!(expected, new);

        let new =
            DirsMultiset::from_dirstate(&StateMap::default(), None).unwrap();
        let expected = DirsMultiset {
            inner: FastHashMap::default(),
        };
        assert_eq!(expected, new);
    }

    #[test]
    fn test_dirsmultiset_new_no_skip() {
        let input_vec: Vec<HgPathBuf> = ["a/", "b/", "a/c", "a/d/"]
            .iter()
            .map(|e| HgPathBuf::from_bytes(e.as_bytes()))
            .collect();
        let expected_inner = [("", 2), ("a", 3), ("b", 1), ("a/d", 1)]
            .iter()
            .map(|(k, v)| (HgPathBuf::from_bytes(k.as_bytes()), *v))
            .collect();

        let new = DirsMultiset::from_manifest(&input_vec).unwrap();
        let expected = DirsMultiset {
            inner: expected_inner,
        };
        assert_eq!(expected, new);

        let input_map: HashMap<_, _> = ["b/x", "a/c", "a/d/x"]
            .iter()
            .map(|f| {
                (
                    HgPathBuf::from_bytes(f.as_bytes()),
                    DirstateEntry {
                        state: EntryState::Normal,
                        mode: 0,
                        mtime: 0,
                        size: 0,
                    },
                )
            })
            .collect();
        let expected_inner = [("", 2), ("a", 2), ("b", 1), ("a/d", 1)]
            .iter()
            .map(|(k, v)| (HgPathBuf::from_bytes(k.as_bytes()), *v))
            .collect();

        let new = DirsMultiset::from_dirstate(&input_map, None).unwrap();
        let expected = DirsMultiset {
            inner: expected_inner,
        };
        assert_eq!(expected, new);
    }

    #[test]
    fn test_dirsmultiset_new_skip() {
        let input_map: HashMap<_, _> = [
            ("a/", EntryState::Normal),
            ("a/b", EntryState::Normal),
            ("a/c", EntryState::Removed),
            ("a/d", EntryState::Merged),
        ]
        .iter()
        .map(|(f, state)| {
            (
                HgPathBuf::from_bytes(f.as_bytes()),
                DirstateEntry {
                    state: *state,
                    mode: 0,
                    mtime: 0,
                    size: 0,
                },
            )
        })
        .collect();

        // "a" incremented with "a/c" and "a/d/"
        let expected_inner = [("", 1), ("a", 2)]
            .iter()
            .map(|(k, v)| (HgPathBuf::from_bytes(k.as_bytes()), *v))
            .collect();

        let new =
            DirsMultiset::from_dirstate(&input_map, Some(EntryState::Normal))
                .unwrap();
        let expected = DirsMultiset {
            inner: expected_inner,
        };
        assert_eq!(expected, new);
    }
}