1use crate::{
7 PipelineData, PipelineExecutionData, ShellError, Span, Value,
8 ast::{Block, Expr, PipelineElement},
9 debugger::Debugger,
10 engine::EngineState,
11 ir::IrBlock,
12 record,
13 shell_error::generic::GenericError,
14};
15use nu_utils::time::Instant;
16use std::{
17 borrow::{Borrow, Cow},
18 io::BufRead,
19};
20
21#[derive(Debug, Clone, Copy)]
22struct ElementId(usize);
23
24#[derive(Debug, Clone)]
26struct ElementInfo {
27 start: Instant,
28 duration_ns: i64,
29 depth: i64,
30 element_span: Span,
31 element_output: Option<Value>,
32 expr: Option<String>,
33 instruction: Option<(usize, String)>,
34 children: Vec<ElementId>,
35}
36
37impl ElementInfo {
38 pub fn new(depth: i64, element_span: Span) -> Self {
39 ElementInfo {
40 start: Instant::now(),
41 duration_ns: 0,
42 depth,
43 element_span,
44 element_output: None,
45 expr: None,
46 instruction: None,
47 children: vec![],
48 }
49 }
50}
51
52#[derive(Debug, Clone, Copy)]
54pub enum DurationMode {
55 Milliseconds,
56 Value,
57}
58
59#[derive(Debug, Clone)]
61pub struct ProfilerOptions {
62 pub max_depth: i64,
63 pub collect_spans: bool,
64 pub collect_source: bool,
65 pub collect_expanded_source: bool,
66 pub collect_values: bool,
67 pub collect_exprs: bool,
68 pub collect_instructions: bool,
69 pub collect_lines: bool,
70 pub duration_mode: DurationMode,
71}
72
73#[derive(Debug, Clone)]
75pub struct Profiler {
76 depth: i64,
77 opts: ProfilerOptions,
78 elements: Vec<ElementInfo>,
79 element_stack: Vec<ElementId>,
80}
81
82impl Profiler {
83 #[allow(clippy::too_many_arguments)]
84 pub fn new(opts: ProfilerOptions, span: Span) -> Self {
85 let first = ElementInfo {
86 start: Instant::now(),
87 duration_ns: 0,
88 depth: 0,
89 element_span: span,
90 element_output: opts.collect_values.then(|| Value::nothing(span)),
91 expr: opts.collect_exprs.then(|| "call".to_string()),
92 instruction: opts
93 .collect_instructions
94 .then(|| (0, "<start>".to_string())),
95 children: vec![],
96 };
97
98 Profiler {
99 depth: 0,
100 opts,
101 elements: vec![first],
102 element_stack: vec![ElementId(0)],
103 }
104 }
105
106 fn last_element_id(&self) -> Option<ElementId> {
107 self.element_stack.last().copied()
108 }
109
110 fn last_element_mut(&mut self) -> Option<&mut ElementInfo> {
111 self.last_element_id()
112 .and_then(|id| self.elements.get_mut(id.0))
113 }
114}
115
116impl Debugger for Profiler {
117 fn activate(&mut self) {
118 let Some(root_element) = self.last_element_mut() else {
119 eprintln!("Profiler Error: Missing root element.");
120 return;
121 };
122
123 root_element.start = Instant::now();
124 }
125
126 fn deactivate(&mut self) {
127 let Some(root_element) = self.last_element_mut() else {
128 eprintln!("Profiler Error: Missing root element.");
129 return;
130 };
131
132 root_element.duration_ns = root_element.start.elapsed().as_nanos() as i64;
133 }
134
135 fn enter_block(&mut self, _engine_state: &EngineState, _block: &Block) {
136 self.depth += 1;
137 }
138
139 fn leave_block(&mut self, _engine_state: &EngineState, _block: &Block) {
140 self.depth -= 1;
141 }
142
143 fn enter_element(&mut self, engine_state: &EngineState, element: &PipelineElement) {
144 if self.depth > self.opts.max_depth {
145 return;
146 }
147
148 let Some(parent_id) = self.last_element_id() else {
149 eprintln!("Profiler Error: Missing parent element ID.");
150 return;
151 };
152
153 let expr_opt = self
154 .opts
155 .collect_exprs
156 .then(|| expr_to_string(engine_state, &element.expr.expr));
157
158 let new_id = ElementId(self.elements.len());
159
160 let mut new_element = ElementInfo::new(self.depth, element.expr.span);
161 new_element.expr = expr_opt;
162
163 self.elements.push(new_element);
164
165 let Some(parent) = self.elements.get_mut(parent_id.0) else {
166 eprintln!("Profiler Error: Missing parent element.");
167 return;
168 };
169
170 parent.children.push(new_id);
171 self.element_stack.push(new_id);
172 }
173
174 fn leave_element(
175 &mut self,
176 _engine_state: &EngineState,
177 element: &PipelineElement,
178 result: &Result<PipelineData, ShellError>,
179 ) {
180 if self.depth > self.opts.max_depth {
181 return;
182 }
183
184 let element_span = element.expr.span;
185
186 let out_opt = self.opts.collect_values.then(|| match result {
187 Ok(pipeline_data) => match pipeline_data {
188 PipelineData::Value(val, ..) => val.clone(),
189 PipelineData::ListStream(..) => Value::string("list stream", element_span),
190 PipelineData::ByteStream(..) => Value::string("byte stream", element_span),
191 _ => Value::nothing(element_span),
192 },
193 Err(e) => Value::error(e.clone(), element_span),
194 });
195
196 let Some(last_element) = self.last_element_mut() else {
197 eprintln!("Profiler Error: Missing last element.");
198 return;
199 };
200
201 last_element.duration_ns = last_element.start.elapsed().as_nanos() as i64;
202 last_element.element_output = out_opt;
203
204 self.element_stack.pop();
205 }
206
207 fn enter_instruction(
208 &mut self,
209 engine_state: &EngineState,
210 ir_block: &IrBlock,
211 instruction_index: usize,
212 _registers: &[PipelineExecutionData],
213 ) {
214 if self.depth > self.opts.max_depth {
215 return;
216 }
217
218 let Some(parent_id) = self.last_element_id() else {
219 eprintln!("Profiler Error: Missing parent element ID.");
220 return;
221 };
222
223 let instruction = &ir_block.instructions[instruction_index];
224 let span = ir_block.spans[instruction_index];
225
226 let instruction_opt = self.opts.collect_instructions.then(|| {
227 (
228 instruction_index,
229 instruction
230 .display(engine_state, &ir_block.data)
231 .to_string(),
232 )
233 });
234
235 let new_id = ElementId(self.elements.len());
236
237 let mut new_element = ElementInfo::new(self.depth, span);
238 new_element.instruction = instruction_opt;
239
240 self.elements.push(new_element);
241
242 let Some(parent) = self.elements.get_mut(parent_id.0) else {
243 eprintln!("Profiler Error: Missing parent element.");
244 return;
245 };
246
247 parent.children.push(new_id);
248 self.element_stack.push(new_id);
249 }
250
251 fn leave_instruction(
252 &mut self,
253 _engine_state: &EngineState,
254 ir_block: &IrBlock,
255 instruction_index: usize,
256 registers: &[PipelineExecutionData],
257 error: Option<&ShellError>,
258 ) {
259 if self.depth > self.opts.max_depth {
260 return;
261 }
262
263 let instruction = &ir_block.instructions[instruction_index];
264 let span = ir_block.spans[instruction_index];
265
266 let out_opt = self
267 .opts
268 .collect_values
269 .then(|| {
270 error
271 .map(Err)
272 .or_else(|| {
273 instruction
274 .output_register()
275 .map(|register| Ok(®isters[register.get() as usize]))
276 })
277 .map(|result| format_result(result.map(|r| &r.body), span))
278 })
279 .flatten();
280
281 let Some(last_element) = self.last_element_mut() else {
282 eprintln!("Profiler Error: Missing last element.");
283 return;
284 };
285
286 last_element.duration_ns = last_element.start.elapsed().as_nanos() as i64;
287 last_element.element_output = out_opt;
288
289 self.element_stack.pop();
290 }
291
292 fn report(&self, engine_state: &EngineState, profiler_span: Span) -> Result<Value, ShellError> {
293 Ok(Value::list(
294 collect_data(
295 engine_state,
296 self,
297 ElementId(0),
298 ElementId(0),
299 profiler_span,
300 )?,
301 profiler_span,
302 ))
303 }
304}
305
306fn profiler_error(msg: impl Into<Cow<'static, str>>, span: Span) -> ShellError {
307 ShellError::Generic(GenericError::new("Profiler Error", msg, span))
308}
309
310fn expr_to_string(engine_state: &EngineState, expr: &Expr) -> String {
311 match expr {
312 Expr::AttributeBlock(ab) => expr_to_string(engine_state, &ab.item.expr),
313 Expr::Binary(_) => "binary".to_string(),
314 Expr::BinaryOp(_, _, _) => "binary operation".to_string(),
315 Expr::Block(_) => "block".to_string(),
316 Expr::Bool(_) => "bool".to_string(),
317 Expr::Call(call) => {
318 let decl = engine_state.get_decl(call.decl_id);
319 if decl.name() == "collect" && call.head == Span::new(0, 0) {
320 "call (implicit collect)"
321 } else {
322 "call"
323 }
324 .to_string()
325 }
326 Expr::CellPath(_) => "cell path".to_string(),
327 Expr::Closure(_) => "closure".to_string(),
328 Expr::DateTime(_) => "datetime".to_string(),
329 Expr::Directory(_, _) => "directory".to_string(),
330 Expr::ExternalCall(_, _) => "external call".to_string(),
331 Expr::Filepath(_, _) => "filepath".to_string(),
332 Expr::Float(_) => "float".to_string(),
333 Expr::FullCellPath(full_cell_path) => {
334 let head = expr_to_string(engine_state, &full_cell_path.head.expr);
335 format!("full cell path ({head})")
336 }
337 Expr::Garbage => "garbage".to_string(),
338 Expr::GlobPattern(_, _) => "glob pattern".to_string(),
339 Expr::ImportPattern(_) => "import pattern".to_string(),
340 Expr::Int(_) => "int".to_string(),
341 Expr::Keyword(_) => "keyword".to_string(),
342 Expr::List(_) => "list".to_string(),
343 Expr::MatchBlock(_) => "match block".to_string(),
344 Expr::Nothing => "nothing".to_string(),
345 Expr::Operator(_) => "operator".to_string(),
346 Expr::Overlay(_) => "overlay".to_string(),
347 Expr::Range(_) => "range".to_string(),
348 Expr::Record(_) => "record".to_string(),
349 Expr::RowCondition(_) => "row condition".to_string(),
350 Expr::Signature(_) => "signature".to_string(),
351 Expr::String(_) | Expr::RawString(_) => "string".to_string(),
352 Expr::StringInterpolation(_) => "string interpolation".to_string(),
353 Expr::GlobInterpolation(_, _) => "glob interpolation".to_string(),
354 Expr::Collect(_, _) => "collect".to_string(),
355 Expr::Subexpression(_) => "subexpression".to_string(),
356 Expr::Table(_) => "table".to_string(),
357 Expr::UnaryNot(_) => "unary not".to_string(),
358 Expr::ValueWithUnit(_) => "value with unit".to_string(),
359 Expr::Var(_) => "var".to_string(),
360 Expr::VarDecl(_) => "var decl".to_string(),
361 }
362}
363
364fn format_result(
365 result: Result<&PipelineData, impl Borrow<ShellError>>,
366 element_span: Span,
367) -> Value {
368 match result {
369 Ok(pipeline_data) => match pipeline_data {
370 PipelineData::Value(val, ..) => val.clone(),
371 PipelineData::ListStream(..) => Value::string("list stream", element_span),
372 PipelineData::ByteStream(..) => Value::string("byte stream", element_span),
373 _ => Value::nothing(element_span),
374 },
375 Err(e) => Value::error(e.borrow().clone(), element_span),
376 }
377}
378
379fn find_file_of_span(engine_state: &EngineState, span: Span) -> Option<(&str, usize)> {
381 for file in engine_state.files() {
382 if file.covered_span.contains_span(span) {
383 let chunk =
385 engine_state.get_span_contents(Span::new(file.covered_span.start, span.start));
386 let nlines = chunk.lines().count();
387 let line_num = if chunk.last() == Some(&b'\n') {
389 nlines + 1
390 } else {
391 nlines
392 };
393
394 let line_num = usize::max(line_num, 1);
396
397 return Some((&file.name, line_num));
398 }
399 }
400
401 None
402}
403
404fn collect_data(
405 engine_state: &EngineState,
406 profiler: &Profiler,
407 element_id: ElementId,
408 parent_id: ElementId,
409 profiler_span: Span,
410) -> Result<Vec<Value>, ShellError> {
411 let element = &profiler.elements[element_id.0];
412
413 let mut row = record! {
414 "depth" => Value::int(element.depth, profiler_span),
415 "id" => Value::int(element_id.0 as i64, profiler_span),
416 "parent_id" => Value::int(parent_id.0 as i64, profiler_span),
417 };
418
419 if profiler.opts.collect_lines {
420 if let Some((fname, line_num)) = find_file_of_span(engine_state, element.element_span) {
421 row.push("file", Value::string(fname, profiler_span));
422 row.push("line", Value::int(line_num as i64, profiler_span));
423 } else {
424 row.push("file", Value::nothing(profiler_span));
425 row.push("line", Value::nothing(profiler_span));
426 }
427 }
428
429 if profiler.opts.collect_spans {
430 let span_start = i64::try_from(element.element_span.start)
431 .map_err(|_| profiler_error("error converting span start to i64", profiler_span))?;
432 let span_end = i64::try_from(element.element_span.end)
433 .map_err(|_| profiler_error("error converting span end to i64", profiler_span))?;
434
435 row.push(
436 "span",
437 Value::record(
438 record! {
439 "start" => Value::int(span_start, profiler_span),
440 "end" => Value::int(span_end, profiler_span),
441 },
442 profiler_span,
443 ),
444 );
445 }
446
447 if profiler.opts.collect_source {
448 let val = String::from_utf8_lossy(engine_state.get_span_contents(element.element_span));
449 let val = val.trim();
450 let nlines = val.lines().count();
451
452 let fragment = if profiler.opts.collect_expanded_source {
453 val.to_string()
454 } else {
455 let mut first_line = val.lines().next().unwrap_or("").to_string();
456
457 if nlines > 1 {
458 first_line.push_str(" ...");
459 }
460
461 first_line
462 };
463
464 row.push("source", Value::string(fragment, profiler_span));
465 }
466
467 if let Some(expr_string) = &element.expr {
468 row.push("expr", Value::string(expr_string.clone(), profiler_span));
469 }
470
471 if let Some((instruction_index, instruction)) = &element.instruction {
472 row.push(
473 "pc",
474 (*instruction_index)
475 .try_into()
476 .map(|index| Value::int(index, profiler_span))
477 .unwrap_or(Value::nothing(profiler_span)),
478 );
479 row.push("instruction", Value::string(instruction, profiler_span));
480 }
481
482 if let Some(val) = &element.element_output {
483 row.push("output", val.clone());
484 }
485
486 match profiler.opts.duration_mode {
487 DurationMode::Milliseconds => {
488 let val = Value::float(element.duration_ns as f64 / 1000.0 / 1000.0, profiler_span);
489 row.push("duration_ms", val);
490 }
491 DurationMode::Value => {
492 let val = Value::duration(element.duration_ns, profiler_span);
493 row.push("duration", val);
494 }
495 };
496
497 let mut rows = vec![Value::record(row, profiler_span)];
498
499 for child in &element.children {
500 let child_rows = collect_data(engine_state, profiler, *child, element_id, profiler_span)?;
501 rows.extend(child_rows);
502 }
503
504 Ok(rows)
505}