diff --git a/rust/crates/tidy-tree/src/iter.rs b/rust/crates/tidy-tree/src/iter.rs new file mode 100644 index 0000000..e0363e3 --- /dev/null +++ b/rust/crates/tidy-tree/src/iter.rs @@ -0,0 +1,72 @@ +use crate::Node; + +struct Iter<'a> { + node: &'a Node, + slot_stack: Vec, + finished: bool, +} + +impl<'a> Iterator for Iter<'a> { + type Item = &'a Node; + + fn next(&mut self) -> Option { + if self.finished { + return None; + } + + let ans = self.node; + if self.node.children.len() > 0 { + self.node = &self.node.children[0]; + self.slot_stack.push(0); + return Some(ans); + } else if self.node.parent.is_some() { + let mut parent = unsafe { self.node.parent.unwrap().as_ref() }; + let mut index_slot = self.slot_stack.pop().unwrap(); + while parent.children.len() <= index_slot + 1 + && self.slot_stack.len() > 0 + && parent.parent.is_some() + { + parent = unsafe { parent.parent.unwrap().as_ref() }; + index_slot = self.slot_stack.pop().unwrap(); + } + + if parent.children.len() > index_slot + 1 { + self.node = &parent.children[index_slot + 1]; + self.slot_stack.push(index_slot + 1); + return Some(ans); + } else { + self.finished = true; + return Some(ans); + } + } + + return Some(ans); + } +} + +impl Node { + #[inline] + fn iter(&self) -> Iter { + Iter { + node: self, + slot_stack: vec![], + finished: false, + } + } +} + +#[cfg(test)] +mod iter_test { + use super::*; + + #[test] + fn test_node_iter() { + let mut root = Node::new_with_child(0, 1., 1., Node::new(1, 2., 2.)); + assert_eq!(root.iter().count(), 2); + root.append_child(Node::new(2, 3., 3.)); + assert_eq!(root.iter().count(), 3); + for (i, node) in root.iter().enumerate() { + assert_eq!(i, node.id); + } + } +} diff --git a/rust/crates/tidy-tree/src/lib.rs b/rust/crates/tidy-tree/src/lib.rs index 065f083..1e08ebe 100644 --- a/rust/crates/tidy-tree/src/lib.rs +++ b/rust/crates/tidy-tree/src/lib.rs @@ -1,5 +1,6 @@ #![allow(dead_code, unused_imports, unused_variables)] pub mod geometry; +mod iter; mod layout; mod node; mod utils; diff --git a/rust/crates/tidy-tree/tests/layout_test.rs b/rust/crates/tidy-tree/tests/layout_test.rs index 4ad8a03..bb34e0e 100644 --- a/rust/crates/tidy-tree/tests/layout_test.rs +++ b/rust/crates/tidy-tree/tests/layout_test.rs @@ -1,4 +1,4 @@ -use std::{collections::HashSet, panic::catch_unwind, ptr::NonNull}; +use std::{panic::catch_unwind, ptr::NonNull}; mod aesthetic_rules; use rand::prelude::*;