Iterator by Node ID

src/octree/batch_iterator.rs

@@ -164,78 +164,3 @@ impl<'a> BatchIterator<'a> {
         point_stream.callback()
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::color::Color;
-    use crate::generation::build_octree;
-    use cgmath::Point3;
-    use tempdir::TempDir;
-
-    fn build_test_octree(batch_size: usize) -> Box<octree::Octree> {
-        let default_point = Point {
-            position: Vector3::new(-2_699_182.0, -4_294_938.0, 3_853_373.0), //ECEF parking lot porter dr
-            color: Color {
-                red: 255,
-                green: 0,
-                blue: 0,
-                alpha: 255,
-            },
-            intensity: None,
-        };
-
-        let mut points = vec![default_point; 4 * batch_size + 1];
-        points[3 * batch_size].position = Vector3::new(-2_702_846.0, -4_291_151.0, 3_855_012.0); // ECEF STANFORD
-
-        let p = Point3::new(6_400_000.0, 6_400_000.0, 6_400_000.0);
-        let bounding_box = Aabb3::new(-1.0 * p, p);
-
-        let pool = scoped_pool::Pool::new(10);
-        let tmp_dir = TempDir::new("octree").unwrap();
-
-        build_octree(&pool, &tmp_dir, 1.0, bounding_box, points.into_iter());
-        crate::octree::on_disk::octree_from_directory(tmp_dir.into_path()).unwrap()
-    }
-
-    #[test]
-    //#[ignore]
-    fn test_batch_iterator() {
-        let batch_size = NUM_POINTS_PER_BATCH / 10;
-        // define function
-        let mut point_count: usize = 0;
-        let mut print_count: usize = 1;
-        let num_points = 25 * batch_size / 10;
-        println!("batch_size= {} ,  num_points= {}", batch_size, num_points);
-        let callback_func = |point_data: PointData| -> Result<()> {
-            point_count += point_data.position.len();
-            if point_count >= print_count * 2 * batch_size {
-                print_count += 1;
-                println!("Streamed {} points", point_count);
-            }
-            if point_count >= num_points {
-                return Err(std::io::Error::new(
-                    std::io::ErrorKind::Interrupted,
-                    format!("Maximum number of {} points reached.", num_points),
-                )
-                .into());
-            }
-            Ok(())
-        };
-
-        // octree and iterator
-        let octree = build_test_octree(batch_size);
-        let location = PointLocation {
-            culling: PointCulling::Any(),
-            global_from_local: None,
-        };
-        let mut batch_iterator = BatchIterator::new(&octree, &location, batch_size);
-
-        let _err_stop = batch_iterator
-            .try_for_each_batch(callback_func)
-            .expect_err("Test error");
-
-        assert_eq!(3 * batch_size, point_count);
-        assert_eq!(2, print_count);
-    }
-}

src/octree/mod.rs

@@ -43,12 +43,15 @@ pub use self::factory::OctreeFactory;
 
 mod octree_iterator;
 pub use self::octree_iterator::{
-    contains, intersecting_node_ids, AllPointsIterator, FilteredPointsIterator,
+    contains, intersecting_node_ids, AllPointsIterator, FilteredPointsIterator, NodeIdsIterator,
 };
 
 mod batch_iterator;
 pub use self::batch_iterator::{BatchIterator, PointCulling, PointLocation, NUM_POINTS_PER_BATCH};
 
+#[cfg(test)]
+mod octree_test;
+
 // Version 9 -> 10: Change in NodeId proto from level (u8) and index (u64) to high (u64) and low
 // (u64). We are able to convert the proto on read, so the tools can still read version 9.
 // Version 10 -> 11: Change in AxisAlignedCuboid proto from Vector3f min/max to Vector3d min/max.
@@ -338,11 +341,11 @@ impl Octree {
         AllPointsIterator::new(&self)
     }
 
+    /// return the bounding box saved in meta
     pub fn bounding_box(&self) -> &Aabb3<f64> {
         &self.meta.bounding_box
     }
 }
-
 struct OpenNode {
     node: Node,
     relation: Relation,

src/octree/octree_iterator.rs

@@ -130,3 +130,47 @@ pub fn intersecting_node_ids(
     }
     node_ids
 }
+
+pub struct NodeIdsIterator<'a, F>
+where
+    F: Fn(&NodeId, &Octree) -> bool + 'a,
+{
+    octree: &'a Octree,
+    filter_func: F,
+    node_ids: VecDeque<NodeId>,
+}
+
+impl<'a, F> NodeIdsIterator<'a, F>
+where
+    F: Fn(&NodeId, &Octree) -> bool + 'a,
+{
+    pub fn new(octree: &'a Octree, filter_func: F) -> Self {
+        NodeIdsIterator {
+            octree,
+            node_ids: vec![NodeId::from_level_index(0, 0)].into(),
+            filter_func,
+        }
+    }
+}
+
+impl<'a, F> Iterator for NodeIdsIterator<'a, F>
+where
+    F: Fn(&NodeId, &Octree) -> bool + 'a,
+{
+    type Item = NodeId;
+
+    fn next(&mut self) -> Option<NodeId> {
+        while let Some(current) = self.node_ids.pop_front() {
+            if (self.filter_func)(&current, &self.octree) {
+                for child_index in 0..8 {
+                    let child_id = current.get_child_id(ChildIndex::from_u8(child_index));
+                    if self.octree.nodes.contains_key(&child_id) {
+                        self.node_ids.push_back(child_id);
+                    }
+                }
+                return Some(current);
+            }
+        }
+        None
+    }
+}

src/octree/octree_test.rs

@@ -0,0 +1,146 @@
+#[cfg(test)]
+mod tests {
+    use crate::color::Color;
+    use crate::errors::Result;
+    use crate::generation::build_octree;
+    use crate::octree::{self, BatchIterator, PointCulling, PointLocation};
+    use crate::{Point, PointData};
+    use cgmath::{EuclideanSpace, Point3, Vector3};
+    use collision::{Aabb, Aabb3};
+    use tempdir::TempDir;
+
+    fn build_big_test_octree() -> Box<octree::Octree> {
+        let default_point = Point {
+            position: Vector3::new(-2_699_182.0, -4_294_938.0, 3_853_373.0), //ECEF parking lot porter dr
+            color: Color {
+                red: 255,
+                green: 0,
+                blue: 0,
+                alpha: 255,
+            },
+            intensity: None,
+        };
+
+        let mut points = vec![default_point; 200_000];
+        points[3].position = Vector3::new(-2_702_846.0, -4_291_151.0, 3_855_012.0); // ECEF STANFORD
+
+        let p = Point3::new(6_400_000.0, 6_400_000.0, 6_400_000.0);
+        let bounding_box = Aabb3::new(-1.0 * p, p);
+
+        let pool = scoped_pool::Pool::new(10);
+        let tmp_dir = TempDir::new("octree").unwrap();
+
+        build_octree(&pool, &tmp_dir, 1.0, bounding_box, points.into_iter());
+        crate::octree::on_disk::octree_from_directory(tmp_dir.into_path()).unwrap()
+    }
+
+    fn build_test_octree() -> Box<octree::Octree> {
+        let default_point = Point {
+            position: Vector3::new(0.0, 0.0, 0.0),
+            color: Color {
+                red: 255,
+                green: 0,
+                blue: 0,
+                alpha: 255,
+            },
+            intensity: None,
+        };
+
+        let mut points = vec![default_point; 100_001];
+        points[100_000].position = Vector3::new(-200., -40., 30.);
+
+        let bounding_box = Aabb3::zero().grow(Point3::from_vec(points[100_000].position));
+
+        let pool = scoped_pool::Pool::new(10);
+        let tmp_dir = TempDir::new("octree").unwrap();
+
+        build_octree(&pool, &tmp_dir, 1.0, bounding_box, points.into_iter());
+        crate::octree::on_disk::octree_from_directory(tmp_dir.into_path()).unwrap()
+    }
+
+    #[test]
+    fn test_batch_iterator() {
+        let batch_size = 5000;
+        // define function
+        let mut callback_count: usize = 0;
+        let max_num_points = 13_000; // 2*batch size + 3000
+        let mut delivered_points = 0;
+        println!(
+            "batch_size= {} ,  num_points= {}",
+            batch_size, max_num_points
+        );
+        let callback_func = |point_data: PointData| -> Result<()> {
+            callback_count += 1;
+            delivered_points += point_data.position.len();
+
+            if delivered_points >= max_num_points {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::Interrupted,
+                    format!("Maximum number of {} points reached.", max_num_points),
+                )
+                .into());
+            }
+            Ok(())
+        };
+
+        // octree and iterator
+        let octree = build_test_octree();
+        let location = PointLocation {
+            culling: PointCulling::Any(),
+            global_from_local: None,
+        };
+        let mut batch_iterator = BatchIterator::new(&octree, &location, batch_size);
+
+        let _err_stop = batch_iterator
+            .try_for_each_batch(callback_func)
+            .expect_err("Test error");
+
+        assert_eq!(3, callback_count);
+        assert_eq!(3 * batch_size, delivered_points);
+        assert!(delivered_points as i32 - max_num_points as i32 >= 0);
+    }
+
+    #[test]
+    #[ignore]
+    fn test_batch_iterator_big_octree() {
+        let batch_size = 5000;
+        // define function
+        let mut callback_count: usize = 0;
+        let max_num_points = 13_000; // 2*batch size + 3000
+        let mut delivered_points = 0;
+        println!(
+            "batch_size= {} ,  num_points= {}",
+            batch_size, max_num_points
+        );
+        let callback_func = |point_data: PointData| -> Result<()> {
+            callback_count += 1;
+            delivered_points += point_data.position.len();
+
+            if delivered_points >= max_num_points {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::Interrupted,
+                    format!("Maximum number of {} points reached.", max_num_points),
+                )
+                .into());
+            }
+            Ok(())
+        };
+
+        // octree and iterator
+        let octree = build_big_test_octree();
+        let location = PointLocation {
+            culling: PointCulling::Any(),
+            global_from_local: None,
+        };
+        let mut batch_iterator = BatchIterator::new(&octree, &location, batch_size);
+
+        let _err_stop = batch_iterator
+            .try_for_each_batch(callback_func)
+            .expect_err("Test error");
+
+        assert_eq!(3, callback_count);
+        assert_eq!(3 * batch_size, delivered_points);
+        assert!(delivered_points as i32 - max_num_points as i32 >= 0);
+    }
+
+}