mononoke_api: Add trees_under_path method to HgRepoContext

Summary: Add a `trees_under_path` method to `HgRepoContext` which is basically just a way to call Mononoke's implementation of `gettreepack`. This is required for EdenAPI's `/trees/prefetch` endpoint, which basically returns the results of a `gettreepack` request via HTTP.

Reviewed By: StanislavGlebik

Differential Revision: D20414860

fbshipit-source-id: 840bd5d19a10625b78758949b290f1947bdb94be

eden/mononoke/mononoke_api/src/hg/repo.rs

@@ -7,9 +7,14 @@
 
 use blobrepo::BlobRepo;
 use context::CoreContext;
+use futures::{compat::Stream01CompatExt, TryStream, TryStreamExt};
+use hgproto::GettreepackArgs;
 use mercurial_types::{HgFileNodeId, HgManifestId};
+use mononoke_types::MPath;
+use repo_client::gettreepack_entries;
 
 use crate::errors::MononokeError;
+use crate::path::MononokePath;
 use crate::repo::RepoContext;
 
 use super::{HgFileContext, HgTreeContext};
@@ -47,12 +52,66 @@ impl HgRepoContext {
         HgFileContext::new_check_exists(self.clone(), filenode_id).await
     }
 
+    /// Look up a tree in the repo by `HgManifestId`.
     pub async fn tree(
         &self,
         manifest_id: HgManifestId,
     ) -> Result<Option<HgTreeContext>, MononokeError> {
         HgTreeContext::new_check_exists(self.clone(), manifest_id).await
     }
+
+    /// Request all of the tree nodes in the repo under a given path.
+    ///
+    /// The caller must specify a list of desired versions of the subtree for
+    /// this path, specified as a list of manifest IDs of tree nodes
+    /// corresponding to different versions of the root node of the subtree.
+    ///
+    /// The caller may also specify a list of versions of the subtree to
+    /// delta against. The server will only return tree nodes that are in
+    /// the requested subtrees that are not in the base subtrees.
+    ///
+    /// Returns a stream of `HgTreeContext`s, each corresponding to a node in
+    /// the requested versions of the subtree, along with its associated path.
+    ///
+    /// This method is equivalent to Mercurial's `gettreepack` wire protocol
+    /// command.
+    pub fn trees_under_path(
+        &self,
+        path: MononokePath,
+        root_versions: impl IntoIterator<Item = HgManifestId>,
+        base_versions: impl IntoIterator<Item = HgManifestId>,
+        depth: Option<usize>,
+    ) -> impl TryStream<Ok = (HgTreeContext, MononokePath), Error = MononokeError> {
+        let ctx = self.ctx().clone();
+        let blob_repo = self.blob_repo();
+        let path_bytes = path
+            .into_mpath()
+            .map(|mpath| mpath.to_vec().into())
+            .unwrap_or_default();
+
+        let args = GettreepackArgs {
+            rootdir: path_bytes,
+            mfnodes: root_versions.into_iter().collect(),
+            basemfnodes: base_versions.into_iter().collect(),
+            directories: vec![], // Not supported.
+            depth,
+        };
+
+        gettreepack_entries(ctx, blob_repo, args)
+            .compat()
+            .map_err(MononokeError::from)
+            .and_then({
+                let repo = self.clone();
+                move |(mfid, path): (HgManifestId, Option<MPath>)| {
+                    let repo = repo.clone();
+                    async move {
+                        let tree = HgTreeContext::new(repo, mfid).await?;
+                        let path = MononokePath::new(path);
+                        Ok((tree, path))
+                    }
+                }
+            })
+    }
 }
 
 #[cfg(test)]