tracing_capture/
iter.rs

Help
1use id_arena::{DefaultArenaBehavior, Id};
2
3use std::{iter::FlatMap, slice};
4
5use crate::{CapturedEvent, CapturedEventInner, CapturedSpan, CapturedSpanInner, Storage};
6
7#[derive(Debug)]
8enum IdsIter<'a, T> {
9    Arena(id_arena::Iter<'a, T, DefaultArenaBehavior<T>>),
10    Slice(slice::Iter<'a, Id<T>>),
11}
12
13/// Iterator over [`CapturedSpan`]s returned from [`Storage::all_spans()`] etc.
14#[derive(Debug)]
15pub struct CapturedSpans<'a> {
16    storage: &'a Storage,
17    ids_iter: IdsIter<'a, CapturedSpanInner>,
18}
19
20impl<'a> CapturedSpans<'a> {
21    pub(crate) fn from_slice(storage: &'a Storage, ids: &'a [Id<CapturedSpanInner>]) -> Self {
22        Self {
23            storage,
24            ids_iter: IdsIter::Slice(ids.iter()),
25        }
26    }
27
28    pub(crate) fn from_arena(storage: &'a Storage) -> Self {
29        Self {
30            storage,
31            ids_iter: IdsIter::Arena(storage.spans.iter()),
32        }
33    }
34}
35
36impl<'a> Iterator for CapturedSpans<'a> {
37    type Item = CapturedSpan<'a>;
38
39    fn next(&mut self) -> Option<Self::Item> {
40        match &mut self.ids_iter {
41            IdsIter::Arena(arena) => {
42                let (_, inner) = arena.next()?;
43                Some(CapturedSpan {
44                    inner,
45                    storage: self.storage,
46                })
47            }
48            IdsIter::Slice(slice) => {
49                let id = *slice.next()?;
50                Some(self.storage.span(id))
51            }
52        }
53    }
54
55    fn size_hint(&self) -> (usize, Option<usize>) {
56        match &self.ids_iter {
57            IdsIter::Arena(arena) => arena.size_hint(),
58            IdsIter::Slice(slice) => slice.size_hint(),
59        }
60    }
61}
62
63impl DoubleEndedIterator for CapturedSpans<'_> {
64    fn next_back(&mut self) -> Option<Self::Item> {
65        match &mut self.ids_iter {
66            IdsIter::Arena(arena) => {
67                let (_, inner) = arena.next_back()?;
68                Some(CapturedSpan {
69                    inner,
70                    storage: self.storage,
71                })
72            }
73            IdsIter::Slice(slice) => {
74                let id = *slice.next_back()?;
75                Some(self.storage.span(id))
76            }
77        }
78    }
79}
80
81impl ExactSizeIterator for CapturedSpans<'_> {
82    fn len(&self) -> usize {
83        match &self.ids_iter {
84            IdsIter::Arena(arena) => arena.len(),
85            IdsIter::Slice(slice) => slice.len(),
86        }
87    }
88}
89
90/// Iterator over [`CapturedEvent`]s returned from [`Storage::all_events()`] etc.
91#[derive(Debug)]
92pub struct CapturedEvents<'a> {
93    storage: &'a Storage,
94    ids_iter: IdsIter<'a, CapturedEventInner>,
95}
96
97impl<'a> CapturedEvents<'a> {
98    pub(crate) fn from_slice(storage: &'a Storage, ids: &'a [Id<CapturedEventInner>]) -> Self {
99        Self {
100            storage,
101            ids_iter: IdsIter::Slice(ids.iter()),
102        }
103    }
104
105    pub(crate) fn from_arena(storage: &'a Storage) -> Self {
106        Self {
107            storage,
108            ids_iter: IdsIter::Arena(storage.events.iter()),
109        }
110    }
111}
112
113impl<'a> Iterator for CapturedEvents<'a> {
114    type Item = CapturedEvent<'a>;
115
116    fn next(&mut self) -> Option<Self::Item> {
117        match &mut self.ids_iter {
118            IdsIter::Arena(arena) => {
119                let (_, inner) = arena.next()?;
120                Some(CapturedEvent {
121                    inner,
122                    storage: self.storage,
123                })
124            }
125            IdsIter::Slice(slice) => {
126                let id = *slice.next()?;
127                Some(self.storage.event(id))
128            }
129        }
130    }
131
132    fn size_hint(&self) -> (usize, Option<usize>) {
133        match &self.ids_iter {
134            IdsIter::Arena(arena) => arena.size_hint(),
135            IdsIter::Slice(slice) => slice.size_hint(),
136        }
137    }
138}
139
140impl DoubleEndedIterator for CapturedEvents<'_> {
141    fn next_back(&mut self) -> Option<Self::Item> {
142        match &mut self.ids_iter {
143            IdsIter::Arena(arena) => {
144                let (_, inner) = arena.next_back()?;
145                Some(CapturedEvent {
146                    inner,
147                    storage: self.storage,
148                })
149            }
150            IdsIter::Slice(slice) => {
151                let id = *slice.next_back()?;
152                Some(self.storage.event(id))
153            }
154        }
155    }
156}
157
158impl ExactSizeIterator for CapturedEvents<'_> {
159    fn len(&self) -> usize {
160        match &self.ids_iter {
161            IdsIter::Arena(arena) => arena.len(),
162            IdsIter::Slice(slice) => slice.len(),
163        }
164    }
165}
166
167/// Iterator over descendant [`CapturedSpan`]s of a span.
168/// Returned by [`CapturedSpan::descendants()`].
169#[derive(Debug)]
170pub struct DescendantSpans<'a> {
171    storage: &'a Storage,
172    layers: Vec<&'a [Id<CapturedSpanInner>]>,
173}
174
175impl<'a> DescendantSpans<'a> {
176    pub(crate) fn new(root: &CapturedSpan<'a>) -> Self {
177        Self {
178            storage: root.storage,
179            layers: vec![&root.inner.child_ids],
180        }
181    }
182}
183
184impl<'a> Iterator for DescendantSpans<'a> {
185    type Item = CapturedSpan<'a>;
186
187    fn next(&mut self) -> Option<Self::Item> {
188        loop {
189            let last_layer = self.layers.last_mut()?;
190            if let Some((&head, tail)) = last_layer.split_first() {
191                let span = self.storage.span(head);
192                *last_layer = tail;
193                if !span.inner.child_ids.is_empty() {
194                    self.layers.push(&span.inner.child_ids);
195                }
196                break Some(span);
197            }
198            // The last layer is empty at this point.
199            self.layers.pop();
200        }
201    }
202}
203
204/// Iterator over the descendant [events](CapturedEvent) of a [`CapturedSpan`].
205/// Returned by [`CapturedSpan::descendant_events()`].
206#[derive(Debug)]
207pub struct DescendantEvents<'a> {
208    inner: FlatMap<
209        DescendantSpans<'a>,
210        CapturedEvents<'a>,
211        fn(CapturedSpan<'a>) -> CapturedEvents<'a>,
212    >,
213}
214
215impl<'a> DescendantEvents<'a> {
216    pub(crate) fn new(root: &CapturedSpan<'a>) -> Self {
217        Self {
218            inner: root.descendants().flat_map(|span| span.events()),
219        }
220    }
221}
222
223impl<'a> Iterator for DescendantEvents<'a> {
224    type Item = CapturedEvent<'a>;
225
226    fn next(&mut self) -> Option<Self::Item> {
227        self.inner.next()
228    }
229
230    fn size_hint(&self) -> (usize, Option<usize>) {
231        self.inner.size_hint()
232    }
233}
tracing_capture/iter.rs

tracing_capture/
iter.rs
