Benchmark each pipeline element (#7854)

# Description

Adds a `profile` command that profiles each pipeline element of a block
and can also recursively step into child blocks.

# Limitations
* It is implemented using pipeline metadata which currently get lost in
some circumstances (e.g.,
https://github.com/nushell/nushell/issues/4501). This means that the
profiler will lose data coming from subexpressions. This issue will
hopefully be solved in the future.
* It also does not step into individual loop iteration which I'm not
sure why but maybe that's a good thing.

# User-Facing Changes

Shouldn't change any existing behavior.

# Tests + Formatting

Don't forget to add tests that cover your changes.

Make sure you've run and fixed any issues with these commands:

- `cargo fmt --all -- --check` to check standard code formatting (`cargo
fmt --all` applies these changes)
- `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used -A
clippy::needless_collect` to check that you're using the standard code
style
- `cargo test --workspace` to check that all tests pass

# After Submitting

If your PR had any user-facing changes, update [the
documentation](https://github.com/nushell/nushell.github.io) after the
PR is merged, if necessary. This will help us keep the docs up to date.

---------

Co-authored-by: Darren Schroeder <343840+fdncred@users.noreply.github.com>

crates/nu-command/src/core_commands/metadata.rs

@@ -75,7 +75,7 @@ impl Command for Metadata {
             None => {
                 let mut cols = vec![];
                 let mut vals = vec![];
-                if let Some(x) = &input.metadata() {
+                if let Some(x) = input.metadata().as_deref() {
                     match x {
                         PipelineMetadata {
                             data_source: DataSource::Ls,
@@ -89,6 +89,12 @@ impl Command for Metadata {
                             cols.push("source".into());
                             vals.push(Value::string("into html --list", head))
                         }
+                        PipelineMetadata {
+                            data_source: DataSource::Profiling(values),
+                        } => {
+                            cols.push("profiling".into());
+                            vals.push(Value::list(values.clone(), head))
+                        }
                     }
                 }
 
@@ -118,7 +124,11 @@ impl Command for Metadata {
     }
 }
 
-fn build_metadata_record(arg: &Value, metadata: &Option<PipelineMetadata>, head: Span) -> Value {
+fn build_metadata_record(
+    arg: &Value,
+    metadata: &Option<Box<PipelineMetadata>>,
+    head: Span,
+) -> Value {
     let mut cols = vec![];
     let mut vals = vec![];
 
@@ -140,7 +150,7 @@ fn build_metadata_record(arg: &Value, metadata: &Option<PipelineMetadata>, head:
         });
     }
 
-    if let Some(x) = &metadata {
+    if let Some(x) = metadata.as_deref() {
         match x {
             PipelineMetadata {
                 data_source: DataSource::Ls,
@@ -154,6 +164,12 @@ fn build_metadata_record(arg: &Value, metadata: &Option<PipelineMetadata>, head:
                 cols.push("source".into());
                 vals.push(Value::string("into html --list", head))
             }
+            PipelineMetadata {
+                data_source: DataSource::Profiling(values),
+            } => {
+                cols.push("profiling".into());
+                vals.push(Value::list(values.clone(), head))
+            }
         }
     }
 

crates/nu-command/src/default_context.rs

@@ -475,6 +475,7 @@ pub fn create_default_context() -> EngineState {
         // Experimental
         bind_command! {
             IsAdmin,
+            Profile,
             View,
             ViewFiles,
             ViewSource,

crates/nu-command/src/experimental/mod.rs

@@ -1,10 +1,12 @@
 mod is_admin;
+mod profile;
 mod view;
 mod view_files;
 mod view_source;
 mod view_span;
 
 pub use is_admin::IsAdmin;
+pub use profile::Profile;
 pub use view::View;
 pub use view_files::ViewFiles;
 pub use view_source::ViewSource;

crates/nu-command/src/experimental/profile.rs

@@ -0,0 +1,113 @@
+use nu_engine::{eval_block, CallExt};
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Closure, Command, EngineState, ProfilingConfig, Stack};
+use nu_protocol::{
+    Category, DataSource, Example, IntoPipelineData, PipelineData, PipelineMetadata, Signature,
+    Spanned, SyntaxShape, Type, Value,
+};
+
+#[derive(Clone)]
+pub struct Profile;
+
+impl Command for Profile {
+    fn name(&self) -> &str {
+        "profile"
+    }
+
+    fn usage(&self) -> &str {
+        "Profile each pipeline element in a closure."
+    }
+
+    fn extra_usage(&self) -> &str {
+        r#"The command collects run time of every pipeline element, recursively stepping into child closures
+until a maximum depth. Optionally, it also collects the source code and intermediate values.
+
+Current known limitations are:
+* profiling data from subexpressions is not tracked
+* it does not step into loop iterations"#
+    }
+
+    fn signature(&self) -> nu_protocol::Signature {
+        Signature::build("profile")
+            .required(
+                "closure",
+                SyntaxShape::Closure(Some(vec![SyntaxShape::Any])),
+                "the closure to run",
+            )
+            .switch("source", "Collect source code in the report", None)
+            .switch("values", "Collect values in the report", None)
+            .named(
+                "max-depth",
+                SyntaxShape::Int,
+                "How many levels of blocks to step into (default: 1)",
+                Some('d'),
+            )
+            .input_output_types(vec![(Type::Any, Type::Table(vec![]))])
+            .allow_variants_without_examples(true)
+            .category(Category::Debug)
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        let capture_block: Spanned<Closure> = call.req(engine_state, stack, 0)?;
+        let block = engine_state.get_block(capture_block.item.block_id);
+
+        let redirect_stdout = call.redirect_stdout;
+        let redirect_stderr = call.redirect_stderr;
+
+        let mut stack = stack.captures_to_stack(&capture_block.item.captures);
+
+        let input_val = input.into_value(call.head);
+
+        if let Some(var) = block.signature.get_positional(0) {
+            if let Some(var_id) = &var.var_id {
+                stack.add_var(*var_id, input_val.clone());
+            }
+        }
+
+        stack.profiling_config = ProfilingConfig::new(
+            call.get_flag::<i64>(engine_state, &mut stack, "max-depth")?
+                .unwrap_or(1),
+            call.has_flag("source"),
+            call.has_flag("values"),
+        );
+
+        let profiling_metadata = Box::new(PipelineMetadata {
+            data_source: DataSource::Profiling(vec![]),
+        });
+
+        let result = if let Some(PipelineMetadata {
+            data_source: DataSource::Profiling(values),
+        }) = eval_block(
+            engine_state,
+            &mut stack,
+            block,
+            input_val.into_pipeline_data_with_metadata(profiling_metadata),
+            redirect_stdout,
+            redirect_stderr,
+        )?
+        .metadata()
+        .map(|m| *m)
+        {
+            Value::list(values, call.head)
+        } else {
+            Value::nothing(call.head)
+        };
+
+        Ok(result.into_pipeline_data())
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description:
+                "Profile some code, stepping into the `spam` command and collecting source.",
+            example: r#"def spam [] { "spam" }; profile { spam | str length } -d 2 --source"#,
+            result: None,
+        }]
+    }
+}

crates/nu-command/src/filesystem/ls.rs

@@ -264,9 +264,9 @@ impl Command for Ls {
                 _ => Some(Value::Nothing { span: call_span }),
             })
             .into_pipeline_data_with_metadata(
-                PipelineMetadata {
+                Box::new(PipelineMetadata {
                     data_source: DataSource::Ls,
-                },
+                }),
                 engine_state.ctrlc.clone(),
             ))
     }

crates/nu-command/src/filters/uniq.rs

@@ -264,7 +264,7 @@ pub fn uniq(
     call: &Call,
     input: Vec<Value>,
     item_mapper: Box<dyn Fn(ItemMapperState) -> ValueCounter>,
-    metadata: Option<PipelineMetadata>,
+    metadata: Option<Box<PipelineMetadata>>,
 ) -> Result<PipelineData, ShellError> {
     let ctrlc = engine_state.ctrlc.clone();
     let head = call.head;

crates/nu-command/src/formats/to/html.rs

@@ -320,9 +320,9 @@ fn to_html(
             vals: result,
             span: head,
         }
-        .into_pipeline_data_with_metadata(PipelineMetadata {
+        .into_pipeline_data_with_metadata(Box::new(PipelineMetadata {
             data_source: DataSource::HtmlThemes,
-        }));
+        })));
     } else {
         let theme_span = match &theme {
             Some(v) => v.span,

crates/nu-command/src/viewers/table.rs

@@ -612,9 +612,9 @@ fn handle_row_stream(
     call: &Call,
     row_offset: usize,
     ctrlc: Option<Arc<AtomicBool>>,
-    metadata: Option<PipelineMetadata>,
+    metadata: Option<Box<PipelineMetadata>>,
 ) -> Result<PipelineData, ShellError> {
-    let stream = match metadata {
+    let stream = match metadata.as_deref() {
         // First, `ls` sources:
         Some(PipelineMetadata {
             data_source: DataSource::Ls,