scene_graph/
iter_mut.rs

1use thunderdome::Index;
2
3use crate::{Node, NodeIndex, SceneGraph};
4
5/// A mutable iterator over the children of a node in a [SceneGraph].
6/// See [SceneGraph::iter_mut] for more information.
7pub struct SceneGraphIterMut<'a, T> {
8    sg: &'a mut SceneGraph<T>,
9    stacks: Vec<StackState>,
10}
11
12impl<'a, T> SceneGraphIterMut<'a, T> {
13    pub(crate) fn new(sg: &'a mut SceneGraph<T>, root_node_idx: NodeIndex) -> Self {
14        let mut stacks = Vec::new();
15
16        let first_child = match root_node_idx {
17            NodeIndex::Root => sg.root_children.map(|v| v.first),
18            NodeIndex::Branch(idx) => sg.arena.get(idx).and_then(|v| v.children.map(|v| v.first)),
19        };
20
21        if let Some(first_child) = first_child {
22            stacks.push(StackState::new(root_node_idx, first_child));
23        };
24        SceneGraphIterMut { sg, stacks }
25    }
26}
27
28impl<'a, T> Iterator for SceneGraphIterMut<'a, T> {
29    type Item = (&'a mut T, &'a mut T);
30
31    fn next(&mut self) -> Option<Self::Item> {
32        // if we're out of stack frames, we die here
33        let stack_frame = self.stacks.pop()?;
34
35        let (parent, current_child) = match stack_frame.parent {
36            NodeIndex::Root => {
37                let parent = &mut self.sg.root;
38
39                let child = self.sg.arena.get_mut(stack_frame.current_child).unwrap();
40
41                (parent, child)
42            }
43            NodeIndex::Branch(idx) => {
44                let (parent, current_child) = self.sg.arena.get2_mut(idx, stack_frame.current_child);
45
46                (&mut parent.unwrap().value, current_child.unwrap())
47            }
48        };
49
50        // safety:  this is a lifetime extension, which i know is valid because get2_mut
51        // panics when we pass in two of the same things, and this iterator requires `&mut SG`
52        // to call `next`.
53        let (parent, current_child): (&mut T, &mut Node<T>) =
54            unsafe { (&mut *(parent as *mut _), &mut *(current_child as *mut _)) };
55
56        // if there's a sibling, push it onto the to do list!
57        if let Some(next_sibling) = current_child.next_sibling {
58            self.stacks.push(StackState::new(stack_frame.parent, next_sibling));
59        }
60
61        if let Some(first_child) = current_child.children.map(|v| v.first) {
62            self.stacks.push(StackState::new(
63                NodeIndex::Branch(stack_frame.current_child),
64                first_child,
65            ));
66        }
67
68        Some((parent, &mut current_child.value))
69    }
70}
71
72#[derive(Debug)]
73struct StackState {
74    parent: NodeIndex,
75    current_child: Index,
76}
77
78impl StackState {
79    fn new(parent: NodeIndex, first_child: Index) -> Self {
80        Self {
81            parent,
82            current_child: first_child,
83        }
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn scene_graph_returns_nothing_on_empty_iteration() {
93        let mut scene_graph = SceneGraph::new("Root");
94
95        assert!(scene_graph.iter_mut().next().is_none());
96    }
97
98    #[test]
99    fn normal_iteration() {
100        let mut sg = SceneGraph::new("Root");
101        let root_idx = NodeIndex::Root;
102        sg.attach(root_idx, "First Child").unwrap();
103
104        let second_child = sg.attach(root_idx, "Second Child").unwrap();
105        sg.attach(second_child, "First Grandchild").unwrap();
106
107        assert_eq!(
108            Vec::from_iter(sg.iter_mut().map(|(_parent, value)| &*value).copied()),
109            vec!["First Child", "Second Child", "First Grandchild"]
110        );
111    }
112
113    #[test]
114    fn stagger_iteration() {
115        let mut sg = SceneGraph::new("Root");
116        let root_idx = NodeIndex::Root;
117        let child = sg.attach(root_idx, "First Child").unwrap();
118        sg.attach(child, "Second Child").unwrap();
119
120        assert_eq!(
121            Vec::from_iter(sg.iter_mut().map(|(_parent, value)| &*value).copied()),
122            vec!["First Child", "Second Child"]
123        );
124    }
125
126    #[test]
127    fn single_iteration() {
128        let mut sg = SceneGraph::new("Root");
129        let root_idx = NodeIndex::Root;
130        sg.attach(root_idx, "First Child").unwrap();
131
132        assert_eq!(
133            Vec::from_iter(sg.iter_mut().map(|(_parent, value)| &*value).copied()),
134            vec!["First Child",]
135        );
136    }
137}
scene_graph/iter_mut.rs

scene_graph/
iter_mut.rs
