From 3acd091d7e482299a430620e0f84a04451ffc4ad Mon Sep 17 00:00:00 2001
From: Timo Tiuraniemi <timo.tiuraniemi@iki.fi>
Date: Tue, 18 Oct 2022 16:07:55 +0300
Subject: [PATCH] Add Iterator::next_tree() and Iterator::prev_tree()

---
 src/iterator.rs | 45 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/src/iterator.rs b/src/iterator.rs
index 9e512c6..ba25281 100644
--- a/src/iterator.rs
+++ b/src/iterator.rs
@@ -314,6 +314,51 @@ impl Iterator {
     self.offset = 2 * self.offset + 1;
     self.index
   }
+
+  /// Move to the next tree from the current position and return
+  /// new index.
+  ///
+  /// ## Examples
+  /// ```rust
+  /// assert_eq!(flat_tree::Iterator::new(0).next_tree(), 2);
+  /// assert_eq!(flat_tree::Iterator::new(1).next_tree(), 4);
+  /// assert_eq!(flat_tree::Iterator::new(3).next_tree(), 8);
+  /// assert_eq!(flat_tree::Iterator::new(7).next_tree(), 16);
+  /// ```
+  #[inline]
+  pub fn next_tree(&mut self) -> u64 {
+    self.index = self.index + self.factor / 2 + 1;
+    self.offset = self.index / 2;
+    self.factor = 2;
+    self.index
+  }
+
+  /// Move to the previous tree from the current position and return
+  /// new index.
+  ///
+  /// ## Examples
+  /// ```rust
+  /// assert_eq!(flat_tree::Iterator::new(0).prev_tree(), 0);
+  /// assert_eq!(flat_tree::Iterator::new(1).prev_tree(), 0);
+  /// assert_eq!(flat_tree::Iterator::new(2).prev_tree(), 0);
+  /// assert_eq!(flat_tree::Iterator::new(3).prev_tree(), 0);
+  /// assert_eq!(flat_tree::Iterator::new(7).prev_tree(), 0);
+  /// assert_eq!(flat_tree::Iterator::new(5).prev_tree(), 2);
+  /// assert_eq!(flat_tree::Iterator::new(9).prev_tree(), 6);
+  /// assert_eq!(flat_tree::Iterator::new(11).prev_tree(), 6);
+  /// ```
+  #[inline]
+  pub fn prev_tree(&mut self) -> u64 {
+    if self.offset == 0 {
+      self.index = 0;
+      self.factor = 2;
+    } else {
+      self.index = self.index - self.factor / 2 - 1;
+      self.offset = self.index / 2;
+      self.factor = 2;
+    }
+    self.index
+  }
 }
 
 impl iter::Iterator for Iterator {