From 3b5966895f75f23b7e26153078be30cb77ca9fbe Mon Sep 17 00:00:00 2001
From: Durham Goode <durham@fb.com>
Date: Wed, 15 Aug 2018 15:05:35 -0700
Subject: [PATCH] mutablehistorypack: implement get_ancestors

Summary: Implements the HistoryStore get_ancestors api.

Reviewed By: quark-zju

Differential Revision: D9136980

fbshipit-source-id: 59b7a1d51c4bf95edec452fcb912fb7647151d24
---
 lib/revisionstore/src/ancestors.rs          |  6 ++
 lib/revisionstore/src/key.rs                |  2 +-
 lib/revisionstore/src/mutablehistorypack.rs | 62 ++++++++++++++++++++-
 3 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/lib/revisionstore/src/ancestors.rs b/lib/revisionstore/src/ancestors.rs
index 68319d8a04..e5fdd5a1d9 100644
--- a/lib/revisionstore/src/ancestors.rs
+++ b/lib/revisionstore/src/ancestors.rs
@@ -30,6 +30,9 @@ impl<T: Fn(&Key, &HashSet<Key>) -> Result<NodeInfo>> AncestorIterator<T> {
             queue: VecDeque::new(),
         };
         iter.queue.push_back(key.clone());
+
+        // Insert the null id so we stop iterating there
+        iter.seen.insert(Key::default());
         iter
     }
 }
@@ -43,6 +46,9 @@ impl<T: Fn(&Key, &HashSet<Key>) -> Result<Ancestors>> BatchedAncestorIterator<T>
             pending_infos: HashMap::new(),
         };
         iter.queue.push_back(key.clone());
+
+        // Insert the null id so we stop iterating there
+        iter.seen.insert(Key::default());
         iter
     }
 }
diff --git a/lib/revisionstore/src/key.rs b/lib/revisionstore/src/key.rs
index a6a89412ca..9a028b6ab1 100644
--- a/lib/revisionstore/src/key.rs
+++ b/lib/revisionstore/src/key.rs
@@ -1,6 +1,6 @@
 use node::Node;
 
-#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
+#[derive(Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
 pub struct Key {
     // Name is usually a file or directory path
     name: Box<[u8]>,
diff --git a/lib/revisionstore/src/mutablehistorypack.rs b/lib/revisionstore/src/mutablehistorypack.rs
index 43e452be1d..1a4445c89a 100644
--- a/lib/revisionstore/src/mutablehistorypack.rs
+++ b/lib/revisionstore/src/mutablehistorypack.rs
@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 
+use ancestors::AncestorIterator;
 use error::Result;
 use historypack::HistoryPackVersion;
 use historystore::{Ancestors, HistoryStore, NodeInfo};
@@ -47,7 +48,7 @@ impl MutableHistoryPack {
 
 impl HistoryStore for MutableHistoryPack {
     fn get_ancestors(&self, key: &Key) -> Result<Ancestors> {
-        unimplemented!();
+        AncestorIterator::new(key, |k, _seen| self.get_node_info(k)).collect()
     }
 
     fn get_node_info(&self, key: &Key) -> Result<NodeInfo> {
@@ -79,10 +80,67 @@ impl HistoryStore for MutableHistoryPack {
 #[cfg(test)]
 mod tests {
     use super::*;
-
+    use rand::Rng;
+    use rand::SeedableRng;
+    use rand::chacha::ChaChaRng;
     use tempfile::tempdir;
 
+    use node::Node;
+
     quickcheck! {
+        fn test_get_ancestors(keys: Vec<(Key, bool)>) -> bool {
+            let mut rng = ChaChaRng::from_seed([0u8; 32]);
+            let tempdir = tempdir().unwrap();
+            let mut muthistorypack =
+                MutableHistoryPack::new(tempdir.path(), HistoryPackVersion::One).unwrap();
+
+            // Insert all the keys, randomly choosing nodes from the already inserted keys
+            let mut chains = HashMap::<Key, Ancestors>::new();
+            chains.insert(Key::default(), Ancestors::new());
+            for &(ref key, ref has_p2) in keys.iter() {
+                let mut p1 = Key::default();
+                let mut p2 = Key::default();
+                let available_parents = chains.keys().map(|k| k.clone()).collect::<Vec<Key>>();
+
+                if !chains.is_empty() {
+                    p1 = rng.choose(&available_parents[..])
+                        .expect("choose p1")
+                        .clone();
+
+                    if *has_p2 {
+                        p2 = rng.choose(&available_parents[..])
+                            .expect("choose p2")
+                            .clone();
+                    }
+                }
+
+                // Insert into the history pack
+                let info = NodeInfo {
+                    parents: [p1.clone(), p2.clone()],
+                    linknode: Node::random(&mut rng),
+                };
+                muthistorypack.add(&key, &info).unwrap();
+
+                // Compute the ancestors for the inserted key
+                let p1_ancestors = chains.get(&p1).expect("get p1 ancestors").clone();
+                let p2_ancestors = chains.get(&p2).expect("get p2 ancestors").clone();
+                let mut ancestors = Ancestors::new();
+                ancestors.extend(p1_ancestors);
+                ancestors.extend(p2_ancestors);
+                ancestors.insert(key.clone(), info.clone());
+                chains.insert(key.clone(), ancestors);
+            }
+
+            for &(ref key, _) in keys.iter() {
+                let in_pack = muthistorypack.get_ancestors(&key).expect("get ancestors");
+                if in_pack != chains[&key] {
+                    return false;
+                }
+            }
+
+            true
+        }
+
         fn test_get_node_info(insert: HashMap<Key, NodeInfo>, notinsert: Vec<Key>) -> bool {
             let tempdir = tempdir().unwrap();
             let mut muthistorypack =