From 5f78e5fc5f335d32a5fe94c6a75799ea0bad9b04 Mon Sep 17 00:00:00 2001
From: Thomas Forgione <thomas@tforgione.fr>
Date: Thu, 8 Mar 2018 16:33:03 +0100
Subject: [PATCH] Give info about the tree

---
 src/lib.rs | 39 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 395fdb1..1f0787c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,11 @@
 use std::fmt;
 use std::cmp::Ordering;
 
+pub enum Side {
+    Left,
+    Right,
+}
+
 /// A k-dimension kd tree
 pub enum KdTree<T> {
     Node {
@@ -18,6 +23,7 @@ impl<T: Clone> KdTree<T> where {
         comparator: F,
         leaf_size: usize,
     ) -> KdTree<T> {
+        // Compute bounding box of points
         KdTree::new_with_depth(points, dimension, &comparator, leaf_size, 0)
     }
 
@@ -77,16 +83,39 @@ impl<T: Clone> KdTree<T> where {
             )),
         }
     }
+
+    // Fn(location, elements)
+    pub fn traverse_leaves<F: FnMut(&Vec<(T, Side)>, &Vec<T>)>(&self, callback: &mut F) {
+        self.traverse_leaves_aux(callback, &mut vec![])
+    }
+
+    pub fn traverse_leaves_aux<F:FnMut(&Vec<(T, Side)>, &Vec<T>)>(&self, callback: &mut F, locations: &mut Vec<(T, Side)>) {
+
+        // Compute the full bounding box of the points
+        match *self {
+            KdTree::Node { ref location , ref left, ref right } => {
+
+                locations.push((location.clone(), Side::Left));
+                left.traverse_leaves_aux(callback, locations);
+                locations.pop();
+
+                locations.push((location.clone(), Side::Right));
+                right.traverse_leaves_aux(callback, locations);
+                locations.pop();
+
+            },
+            KdTree::Leaf(ref elements) => {
+                callback(locations, elements);
+            }
+
+        }
+    }
 }
 
 impl<T: fmt::Display> KdTree<T> {
     fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
         match self {
-            &KdTree::Node {
-                ref location,
-                ref left,
-                ref right,
-            } => {
+            &KdTree::Node { ref location, ref left, ref right, } => {
                 writeln!(formatter, "{}{}", indent, location)?;
                 left.print(formatter, &format!("{}  ", indent))?;
                 right.print(formatter, &format!("{}  ", indent))?;