revisionstore: add BatchedAncestorIterator class

Summary:
This moves the ancestor iteration logic into it's own class, with
support for cases where we receive bulk sets of ancestors at once. A future diff
will add similar logic for ancestor traversals where we receive one hash at a
time.

Reviewed By: quark-zju

Differential Revision: D9136985

fbshipit-source-id: 7f918476f777020b3436f5104ad3bf4b00fe9827

lib/revisionstore/src/ancestors.rs

@@ -0,0 +1,65 @@
+use std::collections::{HashMap, HashSet, VecDeque};
+use std::iter::Iterator;
+
+use error::Result;
+use historystore::{Ancestors, NodeInfo};
+use key::Key;
+
+#[derive(Debug, Fail)]
+#[fail(display = "Ancestor Iterator Error: {:?}", _0)]
+struct AncestorIteratorError(String);
+
+pub struct BatchedAncestorIterator<T: Fn(&Key, &HashSet<Key>) -> Result<Ancestors>> {
+    get_more: T,
+    seen: HashSet<Key>,
+    queue: VecDeque<Key>,
+    pending_infos: HashMap<Key, NodeInfo>,
+}
+
+impl<T: Fn(&Key, &HashSet<Key>) -> Result<Ancestors>> BatchedAncestorIterator<T> {
+    pub fn new(key: &Key, get_more: T) -> Self {
+        let mut iter = BatchedAncestorIterator {
+            get_more: get_more,
+            seen: HashSet::new(),
+            queue: VecDeque::new(),
+            pending_infos: HashMap::new(),
+        };
+        iter.queue.push_back(key.clone());
+        iter
+    }
+}
+
+impl<T: Fn(&Key, &HashSet<Key>) -> Result<Ancestors>> Iterator for BatchedAncestorIterator<T> {
+    type Item = Result<(Key, NodeInfo)>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(current) = self.queue.pop_front() {
+            if !self.pending_infos.contains_key(&current) {
+                match (self.get_more)(&current, &self.seen) {
+                    Err(e) => return Some(Err(e)),
+                    Ok(partial_ancestors) => for (key, node_info) in partial_ancestors.iter() {
+                        self.pending_infos.insert(key.clone(), node_info.clone());
+                    },
+                };
+            }
+
+            if let Some(node_info) = self.pending_infos.remove(&current) {
+                self.seen.insert(current.clone());
+                for parent in &node_info.parents {
+                    if !self.seen.contains(parent) {
+                        self.queue.push_back(parent.clone());
+                    }
+                }
+
+                Some(Ok((current.clone(), node_info.clone())))
+            } else {
+                Some(Err(AncestorIteratorError(format!(
+                    "expected {:?} ancestor",
+                    current
+                )).into()))
+            }
+        } else {
+            None
+        }
+    }
+}

lib/revisionstore/src/lib.rs

@@ -16,6 +16,7 @@ extern crate quickcheck;
 #[cfg(test)]
 extern crate rand;
 
+mod ancestors;
 mod dataindex;
 mod fanouttable;
 mod unionstore;

lib/revisionstore/src/unionhistorystore.rs

@@ -1,8 +1,8 @@
 // Copyright Facebook, Inc. 2018
 // Union history store
-use std::collections::VecDeque;
 use std::rc::Rc;
 
+use ancestors::BatchedAncestorIterator;
 use error::{KeyError, Result};
 use historystore::{Ancestors, HistoryStore, NodeInfo};
 use key::Key;
@@ -41,29 +41,7 @@ impl UnionHistoryStore {
 
 impl HistoryStore for UnionHistoryStore {
     fn get_ancestors(&self, key: &Key) -> Result<Ancestors> {
-        let mut missing_ancestors = VecDeque::new();
-        let mut ancestors = Ancestors::new();
-
-        missing_ancestors.push_back(key.clone());
-        while let Some(current) = missing_ancestors.pop_front() {
-            if ancestors.contains_key(&current) {
-                continue;
-            }
-
-            let partial_ancestors = self.get_partial_ancestors(&current)?;
-
-            for ancestor in partial_ancestors.values() {
-                for parent in &ancestor.parents {
-                    if !partial_ancestors.contains_key(parent) {
-                        missing_ancestors.push_back(parent.clone());
-                    }
-                }
-            }
-
-            ancestors.extend(partial_ancestors);
-        }
-
-        Ok(ancestors)
+        BatchedAncestorIterator::new(key, |k, _seen| self.get_partial_ancestors(k)).collect()
     }
 
     fn get_missing(&self, keys: &[Key]) -> Result<Vec<Key>> {