From c18a1d04f1a8d7068bc67a12eeb32aa28444d787 Mon Sep 17 00:00:00 2001
From: Durham Goode <durham@fb.com>
Date: Thu, 26 Jul 2018 12:14:20 -0700
Subject: [PATCH] revisionstore: implement iter for DataStore trait

Summary:
We need the ability to iterate over a datastore so we can implement
repack and cleanup. In a later diff we'll use this trait to implement repack
functionality in a way that it can apply to any store that implements
IterableStore.

Reviewed By: quark-zju

Differential Revision: D8885094

fbshipit-source-id: 0a2b1ab8cf524392d890302c33e386f1cd218d24
---
 lib/revisionstore/src/datapack.rs | 99 +++++++++++++++++++++++++++++++
 lib/revisionstore/src/lib.rs      |  1 +
 lib/revisionstore/src/repack.rs   |  6 ++
 3 files changed, 106 insertions(+)
 create mode 100644 lib/revisionstore/src/repack.rs

diff --git a/lib/revisionstore/src/datapack.rs b/lib/revisionstore/src/datapack.rs
index 7607b88034..ee619e66f3 100644
--- a/lib/revisionstore/src/datapack.rs
+++ b/lib/revisionstore/src/datapack.rs
@@ -87,6 +87,7 @@ use datastore::{DataStore, Delta, Metadata};
 use error::Result;
 use key::Key;
 use node::Node;
+use repack::IterableStore;
 
 #[derive(Debug, Fail)]
 #[fail(display = "Datapack Error: {:?}", _0)]
@@ -357,6 +358,44 @@ impl DataStore for DataPack {
     }
 }
 
+impl IterableStore for DataPack {
+    fn iter<'a>(&'a self) -> Box<Iterator<Item = Result<Key>> + 'a> {
+        Box::new(DataPackIterator::new(self))
+    }
+}
+
+struct DataPackIterator<'a> {
+    pack: &'a DataPack,
+    offset: u64,
+}
+
+impl<'a> DataPackIterator<'a> {
+    pub fn new(pack: &'a DataPack) -> Self {
+        DataPackIterator {
+            pack: pack,
+            offset: 1, // Start after the header byte
+        }
+    }
+}
+
+impl<'a> Iterator for DataPackIterator<'a> {
+    type Item = Result<Key>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.offset as usize >= self.pack.len() {
+            return None;
+        }
+        let entry = self.pack.read_entry(self.offset);
+        Some(match entry {
+            Ok(ref e) => {
+                self.offset = e.next_offset;
+                Ok(Key::new(e.filename.to_vec().into_boxed_slice(), e.node))
+            }
+            Err(e) => Err(e),
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -562,4 +601,64 @@ mod tests {
             assert_eq!(&chains[i], &chain);
         }
     }
+
+    #[test]
+    fn test_iter() {
+        let mut rng = ChaChaRng::from_seed([0u8; 32]);
+        let tempdir = TempDir::new().unwrap();
+
+        let revisions = vec![
+            (
+                Delta {
+                    data: Rc::new([1, 2, 3, 4]),
+                    base: Some(Key::new(Box::new([0]), Node::random(&mut rng))),
+                    key: Key::new(Box::new([0]), Node::random(&mut rng)),
+                },
+                None,
+            ),
+            (
+                Delta {
+                    data: Rc::new([1, 2, 3, 4]),
+                    base: Some(Key::new(Box::new([0]), Node::random(&mut rng))),
+                    key: Key::new(Box::new([0]), Node::random(&mut rng)),
+                },
+                None,
+            ),
+        ];
+
+        let pack = make_pack(&tempdir, &revisions);
+        assert_eq!(
+            pack.iter().collect::<Result<Vec<Key>>>().unwrap(),
+            revisions
+                .iter()
+                .map(|d| d.0.key.clone())
+                .collect::<Vec<Key>>()
+        );
+    }
+
+    quickcheck! {
+        fn test_iter_quickcheck(keys: Vec<(Vec<u8>, Key)>) -> bool {
+            let tempdir = TempDir::new().unwrap();
+
+            let mut revisions: Vec<(Delta, Option<Metadata>)> = vec![];
+            for (data, key) in keys {
+                revisions.push((
+                    Delta {
+                        data: Rc::from(data.clone()),
+                        base: None,
+                        key: key.clone(),
+                    },
+                    None,
+                ));
+            }
+
+            let pack = make_pack(&tempdir, &revisions);
+            let same = pack.iter().collect::<Result<Vec<Key>>>().unwrap()
+                == revisions
+                    .iter()
+                    .map(|d| d.0.key.clone())
+                    .collect::<Vec<Key>>();
+            same
+        }
+    }
 }
diff --git a/lib/revisionstore/src/lib.rs b/lib/revisionstore/src/lib.rs
index 413f65b943..249650a06d 100644
--- a/lib/revisionstore/src/lib.rs
+++ b/lib/revisionstore/src/lib.rs
@@ -28,5 +28,6 @@ pub mod key;
 pub mod loosefile;
 pub mod mutabledatapack;
 pub mod node;
+pub mod repack;
 pub mod uniondatastore;
 pub mod unionhistorystore;
diff --git a/lib/revisionstore/src/repack.rs b/lib/revisionstore/src/repack.rs
new file mode 100644
index 0000000000..bff7278c59
--- /dev/null
+++ b/lib/revisionstore/src/repack.rs
@@ -0,0 +1,6 @@
+use error::Result;
+use key::Key;
+
+pub trait IterableStore {
+    fn iter<'a>(&'a self) -> Box<Iterator<Item = Result<Key>> + 'a>;
+}