dynamecs_analyze/
span_tree.rs

1use crate::SpanPath;
2use itertools::izip;
3use std::fmt::{Debug, Display, Formatter};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct SpanTree<Payload> {
7    // Stored in depth-first order
8    tree_depth_first: Vec<SpanPath>,
9    payloads: Vec<Payload>,
10    // TODO: Precompute children indices so that we can just skip directly to
11    // relevant indices
12}
13
14#[derive(Debug, Clone)]
15pub struct SpanTreeError {
16    message: String,
17}
18
19impl Display for SpanTreeError {
20    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
21        write!(f, "{}", self.message)
22    }
23}
24
25impl std::error::Error for SpanTreeError {}
26
27impl SpanTreeError {
28    fn message(message: impl Into<String>) -> Self {
29        Self {
30            message: message.into(),
31        }
32    }
33}
34
35impl<Payload> SpanTree<Payload> {
36    pub fn root(&self) -> Option<SpanTreeNode<Payload>> {
37        debug_assert_eq!(self.tree_depth_first.len(), self.payloads.len());
38        (!self.tree_depth_first.is_empty()).then(|| SpanTreeNode {
39            tree_depth_first: &self.tree_depth_first,
40            payloads: &self.payloads,
41            index: 0,
42        })
43    }
44
45    pub fn try_from_depth_first_ordering(paths: Vec<SpanPath>, payloads: Vec<Payload>) -> Result<Self, SpanTreeError> {
46        if let Some((root, others)) = paths.split_first() {
47            let mut stack = Vec::new();
48            for name in root.span_names() {
49                stack.push(name.as_str());
50            }
51
52            for path in others {
53                let num_common_names = izip!(&stack, path.span_names())
54                    .take_while(|(&stack_name, path_name)| stack_name == path_name.as_str())
55                    .count();
56
57                stack.truncate(num_common_names);
58                if num_common_names < root.depth() {
59                    return Err(SpanTreeError::message(
60                        "first path is not an ancestor of all other nodes",
61                    ));
62                }
63
64                if path.depth() > num_common_names + 1 {
65                    return Err(SpanTreeError::message("a non-root node is missing its parent"));
66                } else if path.depth() == num_common_names + 1 {
67                    stack.push(path.span_name().unwrap());
68                } else if path.depth() == num_common_names {
69                    return Err(SpanTreeError::message("duplicate paths detected"));
70                } else {
71                    unreachable!(
72                        "by definition, path depth cannot be smaller \
73                              than the number of common span names"
74                    )
75                }
76            }
77        }
78
79        assert_eq!(paths.len(), payloads.len());
80        Ok(Self {
81            tree_depth_first: paths,
82            payloads,
83        })
84    }
85
86    /// Return an identical tree in which the payload associated with each node
87    /// is transformed by the provided transformation function.
88    pub fn transform_payloads<Payload2>(
89        &mut self,
90        transform: impl FnMut(SpanTreeNode<Payload>) -> Payload2,
91    ) -> SpanTree<Payload2> {
92        let new_payloads: Vec<_> = (0..self.tree_depth_first.len())
93            .map(|i| SpanTreeNode {
94                tree_depth_first: &self.tree_depth_first,
95                payloads: &self.payloads,
96                index: i,
97            })
98            .map(transform)
99            .collect();
100
101        SpanTree {
102            tree_depth_first: self.tree_depth_first.clone(),
103            payloads: new_payloads,
104        }
105    }
106}
107
108pub struct SpanTreeNode<'a, Payload> {
109    tree_depth_first: &'a [SpanPath],
110    payloads: &'a [Payload],
111    index: usize,
112}
113
114impl<'a, Payload> Debug for SpanTreeNode<'a, Payload>
115where
116    Payload: Debug,
117{
118    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
119        let Self {
120            tree_depth_first,
121            payloads,
122            index,
123        } = self;
124        f.debug_struct("SpanTreeNode")
125            .field("path", &tree_depth_first[*index])
126            .field("payload", &payloads[*index])
127            .finish()
128    }
129}
130
131impl<'a, Payload> SpanTreeNode<'a, Payload> {
132    pub fn payload(&self) -> &Payload {
133        &self.payloads[self.index]
134    }
135
136    pub fn path(&self) -> SpanPath {
137        self.tree_depth_first[self.index].clone()
138    }
139
140    pub fn count_children(&self) -> usize {
141        self.visit_children().count()
142    }
143
144    pub fn root(&self) -> SpanTreeNode<'a, Payload> {
145        SpanTreeNode { index: 0, ..*self }
146    }
147
148    pub fn parent(&self) -> Option<SpanTreeNode<'a, Payload>> {
149        self.path().parent().and_then(|parent_path| {
150            self.tree_depth_first[..self.index]
151                .binary_search_by_key(&parent_path.span_names(), |path| path.span_names())
152                .ok()
153                .map(|index| SpanTreeNode {
154                    tree_depth_first: self.tree_depth_first,
155                    payloads: self.payloads,
156                    index,
157                })
158        })
159    }
160
161    pub fn visit_children(&self) -> impl Iterator<Item = SpanTreeNode<'a, Payload>> {
162        // This is just for type inference, to make sure that we get the 'a lifetime
163        // and not something tied to 'self
164        let tree_depth_first: &'a [SpanPath] = self.tree_depth_first;
165        let payloads: &'a [Payload] = self.payloads;
166
167        // TODO: Fix this. It's a temporary workaround for the fact that we cannot move
168        // in the same SpanPath to two different closures, since it's not Copy.
169        // Might want to split SpanPath into SpanPathBuf and SpanPath or something like that
170        let self_path1 = self.path();
171        let self_path2 = self_path1.clone();
172
173        // TODO: Use exponential search to avoid accidental complexity explosion for
174        // very large trees? (It seems unlikely that anyone will have a tree large enough
175        // to make a significant difference though)
176        self.tree_depth_first
177            .iter()
178            .enumerate()
179            // Start at the first potential child
180            .skip(self.index + 1)
181            .take_while(move |(_, maybe_child)| self_path1.is_ancestor_of(maybe_child))
182            .filter(move |(_, descendant)| self_path2.is_parent_of(descendant))
183            .map(move |(child_index, _)| SpanTreeNode {
184                tree_depth_first,
185                payloads,
186                index: child_index,
187            })
188    }
189}
dynamecs_analyze/span_tree.rs

dynamecs_analyze/
span_tree.rs
