extend find_entries to work with prefixes

Summary:
Find entries now support prefixes as well as paths:
- selector can be anything that is convertable to `PathOrPrefix`
- returns `Option<MPath>` instead of `MPath`

Reviewed By: StanislavGlebik, farnz

Differential Revision: D16937401

fbshipit-source-id: 24cce9719d52ce3cbc493575c2a6db5ebf121315

manifest/src/lib.rs

@@ -7,7 +7,7 @@
 #![deny(warnings)]
 
 pub use crate::derive::{derive_manifest, LeafInfo, TreeInfo};
-pub use crate::ops::{Diff, ManifestOps};
+pub use crate::ops::{Diff, ManifestOps, PathOrPrefix};
 pub use crate::types::{Entry, Manifest, PathTree};
 
 mod derive;

manifest/src/ops.rs

@@ -11,7 +11,6 @@ use failure::Error;
 use futures::{stream, Future, Stream};
 use futures_ext::{bounded_traversal::bounded_traversal_stream, BoxStream, FutureExt, StreamExt};
 use mononoke_types::MPath;
-use std::iter::FromIterator;
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum Diff<Entry> {
@@ -20,45 +19,131 @@ pub enum Diff<Entry> {
     Changed(Option<MPath>, Entry, Entry),
 }
 
+#[derive(Debug, Clone)]
+pub enum PathOrPrefix {
+    Path(Option<MPath>),
+    Prefix(Option<MPath>),
+}
+
+impl From<MPath> for PathOrPrefix {
+    fn from(path: MPath) -> Self {
+        PathOrPrefix::Path(Some(path))
+    }
+}
+
+impl From<Option<MPath>> for PathOrPrefix {
+    fn from(path: Option<MPath>) -> Self {
+        PathOrPrefix::Path(path)
+    }
+}
+
 pub trait ManifestOps
 where
     Self: Loadable + Copy + Send + Eq,
     <Self as Loadable>::Value: Manifest<TreeId = Self> + Send,
     <<Self as Loadable>::Value as Manifest>::LeafId: Copy + Send + Eq,
 {
-    fn find_entries(
+    fn find_entries<I, P>(
         &self,
         ctx: CoreContext,
         blobstore: impl Blobstore + Clone,
-        paths: impl IntoIterator<Item = MPath>,
+        paths_or_prefixes: I,
     ) -> BoxStream<
         (
-            MPath,
+            Option<MPath>,
             Entry<Self, <<Self as Loadable>::Value as Manifest>::LeafId>,
         ),
         Error,
-    > {
-        let selector = PathTree::from_iter(paths.into_iter().map(|path| (path, true)));
+    >
+    where
+        I: IntoIterator<Item = P>,
+        PathOrPrefix: From<P>,
+    {
+        enum Select {
+            Single,    // single entry selected
+            Recursive, // whole subtree selected
+            Skip,      // not selected
+        }
+
+        impl Select {
+            fn is_selected(&self) -> bool {
+                match self {
+                    Select::Single | Select::Recursive => true,
+                    Select::Skip => false,
+                }
+            }
+
+            fn is_recursive(&self) -> bool {
+                match self {
+                    Select::Recursive => true,
+                    _ => false,
+                }
+            }
+        }
+
+        impl Default for Select {
+            fn default() -> Select {
+                Select::Skip
+            }
+        }
+
+        let selector: PathTree<Select> = paths_or_prefixes
+            .into_iter()
+            .map(|path_or_prefix| match PathOrPrefix::from(path_or_prefix) {
+                PathOrPrefix::Path(path) => (path, Select::Single),
+                PathOrPrefix::Prefix(path) => (path, Select::Recursive),
+            })
+            .collect();
+
         bounded_traversal_stream(
             256,
-            (selector, None, self.clone()),
-            move |(PathTree { subentries, .. }, path, manifest_id)| {
+            (self.clone(), selector, None, false),
+            move |(manifest_id, selector, path, recursive)| {
+                let PathTree {
+                    subentries,
+                    value: select,
+                } = selector;
+
                 manifest_id
                     .load(ctx.clone(), &blobstore)
                     .map(move |manifest| {
                         let mut output = Vec::new();
                         let mut recurse = Vec::new();
-                        for (name, subentry) in subentries {
-                            if let Some(entry) = manifest.lookup(&name) {
-                                let path = MPath::join_opt_element(path.as_ref(), &name);
-                                if subentry.value {
-                                    output.push((path.clone(), entry.clone()));
+
+                        if recursive || select.is_recursive() {
+                            output.push((path.clone(), Entry::Tree(manifest_id)));
+                            for (name, entry) in manifest.list() {
+                                let path = Some(MPath::join_opt_element(path.as_ref(), &name));
+                                match entry {
+                                    Entry::Leaf(_) => {
+                                        output.push((path.clone(), entry));
+                                    }
+                                    Entry::Tree(manifest_id) => {
+                                        recurse.push((manifest_id, Default::default(), path, true));
+                                    }
                                 }
-                                if let Entry::Tree(manifest_id) = entry {
-                                    recurse.push((subentry, Some(path), manifest_id));
+                            }
+                        } else {
+                            if select.is_selected() {
+                                output.push((path.clone(), Entry::Tree(manifest_id)));
+                            }
+                            for (name, selector) in subentries {
+                                if let Some(entry) = manifest.lookup(&name) {
+                                    let path = Some(MPath::join_opt_element(path.as_ref(), &name));
+                                    match entry {
+                                        Entry::Leaf(_) => {
+                                            if selector.value.is_selected() {
+                                                output.push((path.clone(), entry));
+                                            }
+                                        }
+                                        Entry::Tree(manifest_id) => {
+                                            recurse.push((manifest_id, selector, path, false));
+                                        }
+                                    }
                                 }
                             }
                         }
+
                         (output, recurse)
                     })
             },

manifest/src/tests.rs

@@ -4,7 +4,9 @@
 // This software may be used and distributed according to the terms of the
 // GNU General Public License version 2 or any later version.
 
-use crate::{derive_manifest, Diff, Entry, Manifest, ManifestOps, PathTree, TreeInfo};
+use crate::{
+    derive_manifest, Diff, Entry, Manifest, ManifestOps, PathOrPrefix, PathTree, TreeInfo,
+};
 use blobstore::{Blobstore, Loadable, LoadableError, Storable};
 use context::CoreContext;
 use failure::{err_msg, Error};
@@ -17,7 +19,7 @@ use mononoke_types::{BlobstoreBytes, MPath, MPathElement};
 use pretty_assertions::assert_eq;
 use serde_derive::{Deserialize, Serialize};
 use std::{
-    collections::{hash_map::DefaultHasher, BTreeMap, HashSet},
+    collections::{hash_map::DefaultHasher, BTreeMap, BTreeSet},
     hash::{Hash, Hasher},
     iter::FromIterator,
     sync::{Arc, Mutex},
@@ -444,21 +446,27 @@ fn test_derive_manifest() -> Result<(), Error> {
     Ok(())
 }
 
-fn make_paths(paths_str: &[&str]) -> Result<HashSet<MPath>, Error> {
-    paths_str.into_iter().map(MPath::new).collect()
+fn make_paths(paths_str: &[&str]) -> Result<BTreeSet<Option<MPath>>, Error> {
+    paths_str
+        .into_iter()
+        .map(|path_str| match path_str {
+            &"/" => Ok(None),
+            _ => MPath::new(path_str).map(Some),
+        })
+        .collect()
 }
 
 #[test]
 fn test_find_entries() -> Result<(), Error> {
-    let runtime = Arc::new(Mutex::new(Runtime::new()?));
+    let rt = Arc::new(Mutex::new(Runtime::new()?));
     let blobstore: Arc<dyn Blobstore> = Arc::new(LazyMemblob::new());
     let ctx = CoreContext::test_mock();
 
     // derive manifest
     let derive = |parents, changes| -> Result<TestManifestId, Error> {
-        let manifest_id = runtime
-            .with(|runtime| {
-                runtime.block_on(derive_test_manifest(
+        let manifest_id = rt
+            .with(|rt| {
+                rt.block_on(derive_test_manifest(
                     ctx.clone(),
                     blobstore.clone(),
                     parents,
@@ -479,31 +487,81 @@ fn test_find_entries() -> Result<(), Error> {
             "two/three/5" => Some("5"),
             "two/three/four/6" => Some("6"),
             "two/three/four/7" => Some("7"),
+            "five/six/8" => Some("8"),
+            "five/seven/eight/9" => Some("9"),
+            "five/seven/eight/nine/10" => Some("10"),
+            "five/seven/11" => Some("11"),
         },
     )?;
 
-    let paths = make_paths(&[
-        "one/1",
-        "two/three",
-        "none",
-        "two/three/four/7",
-        "two/three/6",
-    ])?;
+    // use single select
+    {
+        let paths = make_paths(&[
+            "one/1",
+            "two/three",
+            "none",
+            "two/three/four/7",
+            "two/three/6",
+        ])?;
 
-    let results =
-        runtime.with(|rt| rt.block_on(mf0.find_entries(ctx, blobstore, paths).collect()))?;
+        let results = rt.with(|rt| {
+            rt.block_on(
+                mf0.find_entries(ctx.clone(), blobstore.clone(), paths)
+                    .collect(),
+            )
+        })?;
 
-    let mut leafs = HashSet::new();
-    let mut trees = HashSet::new();
-    for (path, entry) in results {
-        match entry {
-            Entry::Tree(_) => trees.insert(path),
-            Entry::Leaf(_) => leafs.insert(path),
-        };
+        let mut leafs = BTreeSet::new();
+        let mut trees = BTreeSet::new();
+        for (path, entry) in results {
+            match entry {
+                Entry::Tree(_) => trees.insert(path),
+                Entry::Leaf(_) => leafs.insert(path),
+            };
+        }
+
+        assert_eq!(leafs, make_paths(&["one/1", "two/three/four/7",])?);
+        assert_eq!(trees, make_paths(&["two/three"])?);
     }
 
-    assert_eq!(leafs, make_paths(&["one/1", "two/three/four/7",])?);
-    assert_eq!(trees, make_paths(&["two/three"])?);
+    // use prefix + single select
+    {
+        let paths = vec![
+            PathOrPrefix::Path(Some(MPath::new("two/three/5")?)),
+            PathOrPrefix::Path(Some(MPath::new("five/seven/11")?)),
+            PathOrPrefix::Prefix(Some(MPath::new("five/seven/eight")?)),
+        ];
+
+        let results = rt.with(|rt| {
+            rt.block_on(
+                mf0.find_entries(ctx.clone(), blobstore.clone(), paths)
+                    .collect(),
+            )
+        })?;
+
+        let mut leafs = BTreeSet::new();
+        let mut trees = BTreeSet::new();
+        for (path, entry) in results {
+            match entry {
+                Entry::Tree(_) => trees.insert(path),
+                Entry::Leaf(_) => leafs.insert(path),
+            };
+        }
+
+        assert_eq!(
+            leafs,
+            make_paths(&[
+                "two/three/5",
+                "five/seven/11",
+                "five/seven/eight/9",
+                "five/seven/eight/nine/10"
+            ])?
+        );
+        assert_eq!(
+            trees,
+            make_paths(&["five/seven/eight", "five/seven/eight/nine"])?
+        );
+    }
 
     Ok(())
 }
@@ -557,21 +615,20 @@ fn test_diff() -> Result<(), Error> {
 
     let diffs = runtime.with(|rt| rt.block_on(mf0.diff(ctx, blobstore, mf1).collect()))?;
 
-    let mut added = HashSet::new();
-    let mut removed = HashSet::new();
-    let mut changed = HashSet::new();
+    let mut added = BTreeSet::new();
+    let mut removed = BTreeSet::new();
+    let mut changed = BTreeSet::new();
     for diff in diffs {
         match diff {
-            Diff::Added(Some(path), _) => {
+            Diff::Added(path, _) => {
                 added.insert(path);
             }
-            Diff::Removed(Some(path), _) => {
+            Diff::Removed(path, _) => {
                 removed.insert(path);
             }
-            Diff::Changed(Some(path), _, _) => {
+            Diff::Changed(path, _, _) => {
                 changed.insert(path);
             }
-            _ => {}
         };
     }
 
@@ -594,7 +651,7 @@ fn test_diff() -> Result<(), Error> {
             "dir_file_conflict/5"
         ])?
     );
-    assert_eq!(changed, make_paths(&["dir", "dir/changed_file"])?);
+    assert_eq!(changed, make_paths(&["/", "dir", "dir/changed_file"])?);
 
     Ok(())
 }

manifest/src/types.rs

@@ -117,8 +117,8 @@ impl<V> PathTree<V>
 where
     V: Default,
 {
-    pub fn insert(&mut self, path: MPath, value: V) {
-        let mut node = path.into_iter().fold(self, |node, element| {
+    pub fn insert(&mut self, path: Option<MPath>, value: V) {
+        let mut node = path.into_iter().flatten().fold(self, |node, element| {
             node.subentries
                 .entry(element)
                 .or_insert_with(Default::default)
@@ -146,6 +146,22 @@ where
     fn from_iter<I>(iter: I) -> Self
     where
         I: IntoIterator<Item = (MPath, V)>,
+    {
+        let mut tree: Self = Default::default();
+        for (path, value) in iter {
+            tree.insert(Some(path), value);
+        }
+        tree
+    }
+}
+
+impl<V> FromIterator<(Option<MPath>, V)> for PathTree<V>
+where
+    V: Default,
+{
+    fn from_iter<I>(iter: I) -> Self
+    where
+        I: IntoIterator<Item = (Option<MPath>, V)>,
     {
         let mut tree: Self = Default::default();
         for (path, value) in iter {