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                    let iterator = unsafe { node.as_node().__children() }.children();
61                    self.iterators.push(iterator);
62                    return Some(node.as_node());
63                }
64                None => {
65                    self.iterators.pop();
66                }
67            }
68        }
69    }
70}
71pub struct TreeIterMut<
72    'a,
73    Children: IChildrenProvider<Node>,
74    Node: IKeyExprTreeNode<Weight>,
75    Weight,
76> where
77    Children::Assoc: IChildren<Node> + 'a,
78    <Children::Assoc as IChildren<Node>>::Node: 'a,
79{
80    iterators: Vec<<Children::Assoc as IChildren<Node>>::IterMut<'a>>,
81    _marker: core::marker::PhantomData<Weight>,
82}
83
84impl<'a, Children: IChildrenProvider<Node>, Node: IKeyExprTreeNode<Weight>, Weight>
85    TreeIterMut<'a, Children, Node, Weight>
86where
87    Children::Assoc: IChildren<Node> + 'a,
88{
89    pub(crate) fn new(children: &'a mut Children::Assoc) -> Self {
90        let mut iterators = Vec::with_capacity(16);
91        iterators.push(children.children_mut());
92        Self {
93            iterators,
94            _marker: Default::default(),
95        }
96    }
97}
98
99impl<
100        'a,
101        Children: IChildrenProvider<Node>,
102        Node: IKeyExprTreeNodeMut<Weight, Children = Children::Assoc> + 'a,
103        Weight,
104    > Iterator for TreeIterMut<'a, Children, Node, Weight>
105where
106    Children::Assoc: IChildren<Node> + 'a,
107{
108    type Item = &'a mut <Children::Assoc as IChildren<Node>>::Node;
109    fn next(&mut self) -> Option<Self::Item> {
110        loop {
111            match self.iterators.last_mut()?.next() {
112                Some(node) => {
113                    let iterator = unsafe { &mut *(node.as_node_mut() as *mut Node) }
114                        .children_mut()
115                        .children_mut();
116                    self.iterators.push(iterator);
117                    return Some(node);
118                }
119                None => {
120                    self.iterators.pop();
121                }
122            }
123        }
124    }
125}
126
127pub struct DepthInstrumented<T>(T);
128impl<
129        'a,
130        Children: IChildrenProvider<Node>,
131        Node: IKeyExprTreeNode<Weight, Children = Children::Assoc> + 'a,
132        Weight,
133    > Iterator for DepthInstrumented<TreeIter<'a, Children, Node, Weight>>
134where
135    Children::Assoc: IChildren<Node> + 'a,
136{
137    type Item = (NonZeroUsize, &'a <Children::Assoc as IChildren<Node>>::Node);
138    fn next(&mut self) -> Option<Self::Item> {
139        loop {
140            let depth = self.0.iterators.len();
141            match self.0.iterators.last_mut()?.next() {
142                Some(node) => {
143                    let iterator = unsafe { node.as_node().__children() }.children();
144                    self.0.iterators.push(iterator);
145                    return Some((unsafe { NonZeroUsize::new_unchecked(depth) }, node));
146                }
147                None => {
148                    self.0.iterators.pop();
149                }
150            }
151        }
152    }
153}