Skip to main content

nu_protocol/debugger/
profiler.rs

1//! Nushell Profiler
2//!
3//! Profiler implements the Debugger trait and can be used via the `debug profile` command for
4//! profiling Nushell code.
5
6use 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/// Stores profiling information about one pipeline element
25#[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/// Whether [`Profiler`] should report duration as [`Value::Duration`]
53#[derive(Debug, Clone, Copy)]
54pub enum DurationMode {
55    Milliseconds,
56    Value,
57}
58
59/// Options for [`Profiler`]
60#[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/// Basic profiler, used in `debug profile`
74#[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(&registers[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
379// Find a file name and a line number (indexed from 1) of a span
380fn 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            // count the number of lines between file start and the searched span start
384            let chunk =
385                engine_state.get_span_contents(Span::new(file.covered_span.start, span.start));
386            let nlines = chunk.lines().count();
387            // account for leading part of current line being counted as a separate line
388            let line_num = if chunk.last() == Some(&b'\n') {
389                nlines + 1
390            } else {
391                nlines
392            };
393
394            // first line has no previous line, clamp up to `1`
395            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}