Skip to main content

zenoh_keyexpr/keyexpr_tree/iters/
tree_iter.rs

1//
2// Copyright (c) 2023 ZettaScale Technology
3//
4// This program and the accompanying materials are made available under the
5// terms of the Eclipse Public License 2.0 which is available at
6// http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
7// which is available at https://www.apache.org/licenses/LICENSE-2.0.
8//
9// SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
10//
11// Contributors:
12//   ZettaScale Zenoh Team, <zenoh@zettascale.tech>
13//
14
15use alloc::vec::Vec;
16use core::num::NonZeroUsize;
17
18use crate::keyexpr_tree::*;
19pub struct TreeIter<'a, Children: IChildrenProvider<Node>, Node: UIKeyExprTreeNode<Weight>, Weight>
20where
21    Children::Assoc: IChildren<Node> + 'a,
22    <Children::Assoc as IChildren<Node>>::Node: 'a,
23{
24    iterators: Vec<<Children::Assoc as IChildren<Node>>::Iter<'a>>,
25    _marker: core::marker::PhantomData<Weight>,
26}
27
28impl<'a, Children: IChildrenProvider<Node>, Node: UIKeyExprTreeNode<Weight>, Weight>
29    TreeIter<'a, Children, Node, Weight>
30where
31    Children::Assoc: IChildren<Node> + 'a,
32{
33    pub(crate) fn new(children: &'a Children::Assoc) -> Self {
34        let mut iterators = Vec::with_capacity(16);
35        iterators.push(children.children());
36        Self {
37            iterators,
38            _marker: Default::default(),
39        }
40    }
41    pub fn with_depth(self) -> DepthInstrumented<Self> {
42        DepthInstrumented(self)
43    }
44}
45
46impl<
47        'a,
48        Children: IChildrenProvider<Node>,
49        Node: UIKeyExprTreeNode<Weight, Children = Children::Assoc> + 'a,
50        Weight,
51    > Iterator for TreeIter<'a, Children, Node, Weight>
52where
53    Children::Assoc: IChildren<Node> + 'a,
54{
55    type Item = &'a Node;
56    fn next(&mut self) -> Option<Self::Item> {
57        loop {
58            match self.iterators.last_mut()?.next() {
59                Some(node) => {
60                    // SAFETY: upheld by the surrounding invariants and prior validation.
61                    let iterator = unsafe { node.as_node().__children() }.children();
62                    self.iterators.push(iterator);
63                    return Some(node.as_node());
64                }
65                None => {
66                    self.iterators.pop();
67                }
68            }
69        }
70    }
71}
72pub struct TreeIterMut<
73    'a,
74    Children: IChildrenProvider<Node>,
75    Node: IKeyExprTreeNode<Weight>,
76    Weight,
77> where
78    Children::Assoc: IChildren<Node> + 'a,
79    <Children::Assoc as IChildren<Node>>::Node: 'a,
80{
81    iterators: Vec<<Children::Assoc as IChildren<Node>>::IterMut<'a>>,
82    _marker: core::marker::PhantomData<Weight>,
83}
84
85impl<'a, Children: IChildrenProvider<Node>, Node: IKeyExprTreeNode<Weight>, Weight>
86    TreeIterMut<'a, Children, Node, Weight>
87where
88    Children::Assoc: IChildren<Node> + 'a,
89{
90    pub(crate) fn new(children: &'a mut Children::Assoc) -> Self {
91        let mut iterators = Vec::with_capacity(16);
92        iterators.push(children.children_mut());
93        Self {
94            iterators,
95            _marker: Default::default(),
96        }
97    }
98}
99
100impl<
101        'a,
102        Children: IChildrenProvider<Node>,
103        Node: IKeyExprTreeNodeMut<Weight, Children = Children::Assoc> + 'a,
104        Weight,
105    > Iterator for TreeIterMut<'a, Children, Node, Weight>
106where
107    Children::Assoc: IChildren<Node> + 'a,
108{
109    type Item = &'a mut <Children::Assoc as IChildren<Node>>::Node;
110    fn next(&mut self) -> Option<Self::Item> {
111        loop {
112            match self.iterators.last_mut()?.next() {
113                Some(node) => {
114                    // SAFETY: upheld by the surrounding invariants and prior validation.
115                    let iterator = unsafe { &mut *(node.as_node_mut() as *mut Node) }
116                        .children_mut()
117                        .children_mut();
118                    self.iterators.push(iterator);
119                    return Some(node);
120                }
121                None => {
122                    self.iterators.pop();
123                }
124            }
125        }
126    }
127}
128
129pub struct DepthInstrumented<T>(T);
130impl<
131        'a,
132        Children: IChildrenProvider<Node>,
133        Node: IKeyExprTreeNode<Weight, Children = Children::Assoc> + 'a,
134        Weight,
135    > Iterator for DepthInstrumented<TreeIter<'a, Children, Node, Weight>>
136where
137    Children::Assoc: IChildren<Node> + 'a,
138{
139    type Item = (NonZeroUsize, &'a <Children::Assoc as IChildren<Node>>::Node);
140    fn next(&mut self) -> Option<Self::Item> {
141        loop {
142            let depth = self.0.iterators.len();
143            match self.0.iterators.last_mut()?.next() {
144                Some(node) => {
145                    // SAFETY: upheld by the surrounding invariants and prior validation.
146                    let iterator = unsafe { node.as_node().__children() }.children();
147                    self.0.iterators.push(iterator);
148                    // SAFETY: upheld by the surrounding invariants and prior validation.
149                    return Some((unsafe { NonZeroUsize::new_unchecked(depth) }, node));
150                }
151                None => {
152                    self.0.iterators.pop();
153                }
154            }
155        }
156    }
157}