From 8eaf7d3904dff4545ea84e0f9011d9914e9b1a30 Mon Sep 17 00:00:00 2001
From: Jakob Stoklund Olesen <jolesen@mozilla.com>
Date: Mon, 20 Nov 2017 13:24:19 -0800
Subject: [PATCH] Add iterators for bforest::{Set,Map}.

The iter() methods return an iterator that traverses all set elements /
map key-value pairs. The iterator doesn't require a mutable container
and forest reference, unlike the cursor types.
---
 lib/cretonne/src/bforest/map.rs | 54 ++++++++++++++++++++++++++++++++-
 lib/cretonne/src/bforest/mod.rs |  4 +--
 lib/cretonne/src/bforest/set.rs | 47 ++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 3 deletions(-)

diff --git a/lib/cretonne/src/bforest/map.rs b/lib/cretonne/src/bforest/map.rs
index 0cc1364c47..23cac5600d 100644
--- a/lib/cretonne/src/bforest/map.rs
+++ b/lib/cretonne/src/bforest/map.rs
@@ -134,6 +134,15 @@ where
     ) -> MapCursor<'a, K, V, C> {
         MapCursor::new(self, forest, comp)
     }
+
+    /// Create an iterator traversing this map. The iterator type is `(K, V)`.
+    pub fn iter<'a>(&'a self, forest: &'a MapForest<K, V, C>) -> MapIter<'a, K, V, C> {
+        MapIter {
+            root: self.root,
+            pool: &forest.nodes,
+            path: Path::default(),
+        }
+    }
 }
 
 #[cfg(test)]
@@ -296,6 +305,37 @@ where
     }
 }
 
+/// An iterator visiting the key-value pairs of a `Map`.
+pub struct MapIter<'a, K, V, C>
+where
+    K: 'a + Copy,
+    V: 'a + Copy,
+    C: 'a + Comparator<K>,
+{
+    root: PackedOption<Node>,
+    pool: &'a NodePool<MapTypes<K, V, C>>,
+    path: Path<MapTypes<K, V, C>>,
+}
+
+impl<'a, K, V, C> Iterator for MapIter<'a, K, V, C>
+where
+    K: 'a + Copy,
+    V: 'a + Copy,
+    C: 'a + Comparator<K>,
+{
+    type Item = (K, V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // We use `self.root` to indicate if we need to go to the first element. Reset to `None`
+        // once we've returned the first element. This also works for an empty tree since the
+        // `path.next()` call returns `None` when the path is empty. This also fuses the iterator.
+        match self.root.take() {
+            Some(root) => Some(self.path.first(root, self.pool)),
+            None => self.path.next(self.pool),
+        }
+    }
+}
+
 #[cfg(test)]
 impl<'a, K, V, C> MapCursor<'a, K, V, C>
 where
@@ -337,6 +377,7 @@ mod test {
         m.clear(&mut f);
 
         assert_eq!(m.get(7, &f, &()), None);
+        assert_eq!(m.iter(&f).next(), None);
 
         let mut c = m.cursor(&mut f, &());
         assert!(c.is_empty());
@@ -367,7 +408,18 @@ mod test {
 
         m.verify(f, &());
 
-        // [ 20 40 50 60 80 90 200 ]
+        assert_eq!(
+            m.iter(f).collect::<Vec<_>>(),
+            [
+                (20, 2.0),
+                (40, 4.0),
+                (50, 5.5),
+                (60, 6.0),
+                (80, 8.0),
+                (90, 9.0),
+                (200, 20.0),
+            ]
+        );
 
         assert_eq!(m.get(0, f, &()), None);
         assert_eq!(m.get(20, f, &()), Some(2.0));
diff --git a/lib/cretonne/src/bforest/mod.rs b/lib/cretonne/src/bforest/mod.rs
index 8234551b23..b1d4aa3c28 100644
--- a/lib/cretonne/src/bforest/mod.rs
+++ b/lib/cretonne/src/bforest/mod.rs
@@ -22,8 +22,8 @@ mod path;
 mod pool;
 mod set;
 
-pub use self::map::{MapForest, Map, MapCursor};
-pub use self::set::{SetForest, Set, SetCursor};
+pub use self::map::{MapForest, Map, MapCursor, MapIter};
+pub use self::set::{SetForest, Set, SetCursor, SetIter};
 
 use self::node::NodeData;
 use self::path::Path;
diff --git a/lib/cretonne/src/bforest/set.rs b/lib/cretonne/src/bforest/set.rs
index ce809ac648..4956a26147 100644
--- a/lib/cretonne/src/bforest/set.rs
+++ b/lib/cretonne/src/bforest/set.rs
@@ -131,6 +131,15 @@ where
     ) -> SetCursor<'a, K, C> {
         SetCursor::new(self, forest, comp)
     }
+
+    /// Create an iterator traversing this set. The iterator type is `K`.
+    pub fn iter<'a>(&'a self, forest: &'a SetForest<K, C>) -> SetIter<'a, K, C> {
+        SetIter {
+            root: self.root,
+            pool: &forest.nodes,
+            path: Path::default(),
+        }
+    }
 }
 
 /// A position in a `Set` used to navigate and modify the ordered set.
@@ -275,6 +284,35 @@ where
     }
 }
 
+/// An iterator visiting the elements of a `Set`.
+pub struct SetIter<'a, K, C>
+where
+    K: 'a + Copy,
+    C: 'a + Comparator<K>,
+{
+    root: PackedOption<Node>,
+    pool: &'a NodePool<SetTypes<K, C>>,
+    path: Path<SetTypes<K, C>>,
+}
+
+impl<'a, K, C> Iterator for SetIter<'a, K, C>
+where
+    K: 'a + Copy,
+    C: 'a + Comparator<K>,
+{
+    type Item = K;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // We use `self.root` to indicate if we need to go to the first element. Reset to `None`
+        // once we've returned the first element. This also works for an empty tree since the
+        // `path.next()` call returns `None` when the path is empty. This also fuses the iterator.
+        match self.root.take() {
+            Some(root) => Some(self.path.first(root, self.pool).0),
+            None => self.path.next(self.pool).map(|(k, _)| k),
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::mem;
@@ -298,6 +336,9 @@ mod test {
         s.clear(&mut f);
         assert!(!s.contains(7, &f, &()));
 
+        // Iterator for an empty set.
+        assert_eq!(s.iter(&f).next(), None);
+
         let mut c = SetCursor::new(&mut s, &mut f, &());
         c.verify();
         assert_eq!(c.elem(), None);
@@ -465,6 +506,12 @@ mod test {
     fn four_level() {
         let mut f = SetForest::<i32, ()>::new();
         let mut s = dense4l(&mut f);
+
+        assert_eq!(
+            s.iter(&f).collect::<Vec<_>>()[0..10],
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+        );
+
         let mut c = s.cursor(&mut f, &());
 
         c.verify();