1use std::borrow::Cow;
6use std::collections::HashMap;
7use std::fmt::Display;
8use std::fmt::Formatter;
9
10use crate::collector::SpanId;
11use crate::collector::SpanRecord;
12use crate::collector::SpanSet;
13use crate::util::CollectToken;
14use crate::util::RawSpans;
15
16type TreeChildren = HashMap<
17 SpanId,
18 (
19 Cow<'static, str>,
20 Vec<SpanId>,
21 Vec<(Cow<'static, str>, Cow<'static, str>)>,
22 ),
23>;
24
25#[derive(Debug, PartialOrd, PartialEq, Ord, Eq)]
26pub struct Tree {
27 name: Cow<'static, str>,
28 children: Vec<Tree>,
29 properties: Vec<(Cow<'static, str>, Cow<'static, str>)>,
30}
31
32impl Display for Tree {
33 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
34 self.fmt_with_depth(f, 0)
35 }
36}
37
38impl Tree {
39 fn fmt_with_depth(&self, f: &mut Formatter<'_>, depth: usize) -> std::fmt::Result {
40 writeln!(
41 f,
42 "{:indent$}{} {:?}",
43 "",
44 self.name,
45 self.properties,
46 indent = depth * 4
47 )?;
48 for child in &self.children {
49 child.fmt_with_depth(f, depth + 1)?;
50 }
51 Ok(())
52 }
53}
54
55impl Tree {
56 pub fn sort(&mut self) {
57 for child in &mut self.children {
58 child.sort();
59 }
60 self.children.as_mut_slice().sort_unstable();
61 }
62
63 pub fn from_raw_spans(raw_spans: RawSpans) -> Vec<Tree> {
64 let mut children: TreeChildren = HashMap::new();
65
66 let spans = raw_spans.into_inner();
67 children.insert(SpanId::default(), ("".into(), vec![], vec![]));
68 for span in &spans {
69 children.insert(
70 span.id,
71 (span.name.clone(), vec![], span.properties.clone()),
72 );
73 }
74 for span in &spans {
75 children
76 .get_mut(&span.parent_id)
77 .as_mut()
78 .unwrap()
79 .1
80 .push(span.id);
81 }
82
83 let mut t = Self::build_tree(SpanId::default(), &mut children);
84 t.sort();
85 t.children
86 }
87
88 pub fn from_span_sets(span_sets: &[(SpanSet, CollectToken)]) -> Vec<(usize, Tree)> {
90 let mut collect = HashMap::<
91 usize,
92 HashMap<
93 SpanId,
94 (
95 Cow<'static, str>,
96 Vec<SpanId>,
97 Vec<(Cow<'static, str>, Cow<'static, str>)>,
98 ),
99 >,
100 >::new();
101
102 for (span_set, token) in span_sets {
103 for item in token.iter() {
104 collect
105 .entry(item.collect_id)
106 .or_default()
107 .insert(SpanId::default(), ("".into(), vec![], vec![]));
108 match span_set {
109 SpanSet::Span(span) => {
110 collect.entry(item.collect_id).or_default().insert(
111 span.id,
112 (span.name.clone(), vec![], span.properties.clone()),
113 );
114 }
115 SpanSet::LocalSpansInner(spans) => {
116 for span in spans.spans.iter() {
117 collect.entry(item.collect_id).or_default().insert(
118 span.id,
119 (span.name.clone(), vec![], span.properties.clone()),
120 );
121 }
122 }
123 SpanSet::SharedLocalSpans(spans) => {
124 for span in spans.spans.iter() {
125 collect.entry(item.collect_id).or_default().insert(
126 span.id,
127 (span.name.clone(), vec![], span.properties.clone()),
128 );
129 }
130 }
131 }
132 }
133 }
134
135 for (span_set, token) in span_sets {
136 for item in token.iter() {
137 match span_set {
138 SpanSet::Span(span) => {
139 let parent_id = if span.parent_id == SpanId::default() {
140 item.parent_id
141 } else {
142 span.parent_id
143 };
144 collect
145 .get_mut(&item.collect_id)
146 .as_mut()
147 .unwrap()
148 .get_mut(&parent_id)
149 .as_mut()
150 .unwrap()
151 .1
152 .push(span.id);
153 }
154 SpanSet::LocalSpansInner(spans) => {
155 for span in spans.spans.iter() {
156 let parent_id = if span.parent_id == SpanId::default() {
157 item.parent_id
158 } else {
159 span.parent_id
160 };
161 collect
162 .get_mut(&item.collect_id)
163 .as_mut()
164 .unwrap()
165 .get_mut(&parent_id)
166 .as_mut()
167 .unwrap()
168 .1
169 .push(span.id);
170 }
171 }
172 SpanSet::SharedLocalSpans(spans) => {
173 for span in spans.spans.iter() {
174 let parent_id = if span.parent_id == SpanId::default() {
175 item.parent_id
176 } else {
177 span.parent_id
178 };
179 collect
180 .get_mut(&item.collect_id)
181 .as_mut()
182 .unwrap()
183 .get_mut(&parent_id)
184 .as_mut()
185 .unwrap()
186 .1
187 .push(span.id);
188 }
189 }
190 }
191 }
192 }
193
194 let mut res = collect
195 .into_iter()
196 .map(|(id, mut children)| {
197 let mut tree = Self::build_tree(SpanId::default(), &mut children);
198 tree.sort();
199 assert_eq!(tree.children.len(), 1);
200 (id, tree.children.pop().unwrap())
201 })
202 .collect::<Vec<(usize, Tree)>>();
203 res.sort_unstable();
204 res
205 }
206
207 pub fn from_span_records(span_records: Vec<SpanRecord>) -> Tree {
208 let mut children: TreeChildren = HashMap::new();
209
210 children.insert(SpanId::default(), ("".into(), vec![], vec![]));
211 for span in &span_records {
212 children.insert(
213 span.span_id,
214 (span.name.clone(), vec![], span.properties.clone()),
215 );
216 }
217 for span in &span_records {
218 children
219 .get_mut(&span.parent_id)
220 .as_mut()
221 .unwrap()
222 .1
223 .push(span.span_id);
224 }
225
226 let mut t = Self::build_tree(SpanId::default(), &mut children);
227 t.sort();
228 assert_eq!(t.children.len(), 1);
229 t.children.remove(0)
230 }
231
232 fn build_tree(id: SpanId, raw: &mut TreeChildren) -> Tree {
233 let (name, children, properties) = raw.get(&id).cloned().unwrap();
234 Tree {
235 name,
236 children: children
237 .into_iter()
238 .map(|id| Self::build_tree(id, raw))
239 .collect(),
240 properties,
241 }
242 }
243}
244
245pub fn tree_str_from_raw_spans(raw_spans: RawSpans) -> String {
246 Tree::from_raw_spans(raw_spans)
247 .iter()
248 .map(|t| format!("\n{}", t))
249 .collect::<Vec<_>>()
250 .join("")
251}
252
253pub fn tree_str_from_span_sets(span_sets: &[(SpanSet, CollectToken)]) -> String {
254 Tree::from_span_sets(span_sets)
255 .iter()
256 .map(|(id, t)| format!("\n#{}\n{}", id, t))
257 .collect::<Vec<_>>()
258 .join("")
259}
260
261pub fn tree_str_from_span_records(span_records: Vec<SpanRecord>) -> String {
262 format!("\n{}", Tree::from_span_records(span_records))
263}