add new --flatten parameter to the ast command (#14400)

# Description

By request, this PR introduces a new `--flatten` parameter to the ast
command for generating a more readable version of the AST output. This
enhancement improves usability by allowing users to easily visualize the
structure of the AST.


![image](https://github.com/user-attachments/assets/a66644ef-5fff-4d3d-a334-4e9f80edb39d)

```nushell
❯ ast 'ls | sort-by type name -i' --flatten --json
[
  {
    "content": "ls",
    "shape": "shape_internalcall",
    "span": {
      "start": 0,
      "end": 2
    }
  },
  {
    "content": "|",
    "shape": "shape_pipe",
    "span": {
      "start": 3,
      "end": 4
    }
  },
  {
    "content": "sort-by",
    "shape": "shape_internalcall",
    "span": {
      "start": 5,
      "end": 12
    }
  },
  {
    "content": "type",
    "shape": "shape_string",
    "span": {
      "start": 13,
      "end": 17
    }
  },
  {
    "content": "name",
    "shape": "shape_string",
    "span": {
      "start": 18,
      "end": 22
    }
  },
  {
    "content": "-i",
    "shape": "shape_flag",
    "span": {
      "start": 23,
      "end": 25
    }
  }
]
❯ ast 'ls | sort-by type name -i' --flatten --json --minify
[{"content":"ls","shape":"shape_internalcall","span":{"start":0,"end":2}},{"content":"|","shape":"shape_pipe","span":{"start":3,"end":4}},{"content":"sort-by","shape":"shape_internalcall","span":{"start":5,"end":12}},{"content":"type","shape":"shape_string","span":{"start":13,"end":17}},{"content":"name","shape":"shape_string","span":{"start":18,"end":22}},{"content":"-i","shape":"shape_flag","span":{"start":23,"end":25}}]
```
# User-Facing Changes
<!-- List of all changes that impact the user experience here. This
helps us keep track of breaking changes. -->

# 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` to
check that you're using the standard code style
- `cargo test --workspace` to check that all tests pass (on Windows make
sure to [enable developer
mode](https://learn.microsoft.com/en-us/windows/apps/get-started/developer-mode-features-and-debugging))
- `cargo run -- -c "use toolkit.nu; toolkit test stdlib"` to run the
tests for the standard library

> **Note**
> from `nushell` you can also use the `toolkit` as follows
> ```bash
> use toolkit.nu # or use an `env_change` hook to activate it
automatically
> toolkit check pr
> ```
-->

# 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.
-->

crates/nu-command/src/debug/ast.rs

@@ -1,6 +1,7 @@
 use nu_engine::command_prelude::*;
-use nu_parser::parse;
+use nu_parser::{flatten_block, parse};
-use nu_protocol::engine::StateWorkingSet;
+use nu_protocol::{engine::StateWorkingSet, record};
+use serde_json::{json, Value as JsonValue};
 
 #[derive(Clone)]
 pub struct Ast;
@@ -16,18 +17,120 @@ impl Command for Ast {
 
     fn signature(&self) -> Signature {
         Signature::build("ast")
-            .input_output_types(vec![(Type::String, Type::record())])
+            .input_output_types(vec![
+                (Type::Nothing, Type::table()),
+                (Type::Nothing, Type::record()),
+                (Type::Nothing, Type::String),
+            ])
             .required(
                 "pipeline",
                 SyntaxShape::String,
                 "The pipeline to print the ast for.",
             )
-            .switch("json", "serialize to json", Some('j'))
+            .switch("json", "Serialize to json", Some('j'))
-            .switch("minify", "minify the nuon or json output", Some('m'))
+            .switch("minify", "Minify the nuon or json output", Some('m'))
+            .switch("flatten", "An easier to read version of the ast", Some('f'))
             .allow_variants_without_examples(true)
             .category(Category::Debug)
     }
 
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Print the ast of a string",
+                example: "ast 'hello'",
+                result: None,
+            },
+            Example {
+                description: "Print the ast of a pipeline",
+                example: "ast 'ls | where name =~ README'",
+                result: None,
+            },
+            Example {
+                description: "Print the ast of a pipeline with an error",
+                example: "ast 'for x in 1..10 { echo $x '",
+                result: None,
+            },
+            Example {
+                description:
+                    "Print the ast of a pipeline with an error, as json, in a nushell table",
+                example: "ast 'for x in 1..10 { echo $x ' --json | get block | from json",
+                result: None,
+            },
+            Example {
+                description: "Print the ast of a pipeline with an error, as json, minified",
+                example: "ast 'for x in 1..10 { echo $x ' --json --minify",
+                result: None,
+            },
+            Example {
+                description: "Print the ast of a string flattened",
+                example: r#"ast "'hello'" --flatten"#,
+                result: Some(Value::test_list(vec![Value::test_record(record! {
+                    "content" => Value::test_string("'hello'"),
+                    "shape" => Value::test_string("shape_string"),
+                    "span" => Value::test_record(record! {
+                        "start" => Value::test_int(0),
+                        "end" => Value::test_int(7),}),
+                })])),
+            },
+            Example {
+                description: "Print the ast of a string flattened, as json, minified",
+                example: r#"ast "'hello'" --flatten --json --minify"#,
+                result: Some(Value::test_string(
+                    r#"[{"content":"'hello'","shape":"shape_string","span":{"start":0,"end":7}}]"#,
+                )),
+            },
+            Example {
+                description: "Print the ast of a pipeline flattened",
+                example: r#"ast 'ls | sort-by type name -i' --flatten"#,
+                result: Some(Value::test_list(vec![
+                    Value::test_record(record! {
+                        "content" => Value::test_string("ls"),
+                        "shape" => Value::test_string("shape_external"),
+                        "span" => Value::test_record(record! {
+                            "start" => Value::test_int(0),
+                            "end" => Value::test_int(2),}),
+                    }),
+                    Value::test_record(record! {
+                        "content" => Value::test_string("|"),
+                        "shape" => Value::test_string("shape_pipe"),
+                        "span" => Value::test_record(record! {
+                            "start" => Value::test_int(3),
+                            "end" => Value::test_int(4),}),
+                    }),
+                    Value::test_record(record! {
+                        "content" => Value::test_string("sort-by"),
+                        "shape" => Value::test_string("shape_internalcall"),
+                        "span" => Value::test_record(record! {
+                            "start" => Value::test_int(5),
+                            "end" => Value::test_int(12),}),
+                    }),
+                    Value::test_record(record! {
+                        "content" => Value::test_string("type"),
+                        "shape" => Value::test_string("shape_string"),
+                        "span" => Value::test_record(record! {
+                            "start" => Value::test_int(13),
+                            "end" => Value::test_int(17),}),
+                    }),
+                    Value::test_record(record! {
+                        "content" => Value::test_string("name"),
+                        "shape" => Value::test_string("shape_string"),
+                        "span" => Value::test_record(record! {
+                            "start" => Value::test_int(18),
+                            "end" => Value::test_int(22),}),
+                    }),
+                    Value::test_record(record! {
+                        "content" => Value::test_string("-i"),
+                        "shape" => Value::test_string("shape_flag"),
+                        "span" => Value::test_record(record! {
+                            "start" => Value::test_int(23),
+                            "end" => Value::test_int(25),}),
+                    }),
+                ])),
+            },
+        ]
+    }
+
     fn run(
         &self,
         engine_state: &EngineState,
@@ -38,19 +141,81 @@ impl Command for Ast {
         let pipeline: Spanned<String> = call.req(engine_state, stack, 0)?;
         let to_json = call.has_flag(engine_state, stack, "json")?;
         let minify = call.has_flag(engine_state, stack, "minify")?;
+        let flatten = call.has_flag(engine_state, stack, "flatten")?;
+
         let mut working_set = StateWorkingSet::new(engine_state);
-        let block_output = parse(&mut working_set, None, pipeline.item.as_bytes(), false);
+        let offset = working_set.next_span_start();
+        let parsed_block = parse(&mut working_set, None, pipeline.item.as_bytes(), false);
+
+        if flatten {
+            let flat = flatten_block(&working_set, &parsed_block);
+            if to_json {
+                let mut json_val: JsonValue = json!([]);
+                for (span, shape) in flat {
+                    let content =
+                        String::from_utf8_lossy(working_set.get_span_contents(span)).to_string();
+
+                    let json = json!(
+                        {
+                            "content": content,
+                            "shape": shape.to_string(),
+                            "span": {
+                                "start": span.start.checked_sub(offset),
+                                "end": span.end.checked_sub(offset),
+                            },
+                        }
+                    );
+                    json_merge(&mut json_val, &json);
+                }
+                let json_string = if minify {
+                    if let Ok(json_str) = serde_json::to_string(&json_val) {
+                        json_str
+                    } else {
+                        "{}".to_string()
+                    }
+                } else if let Ok(json_str) = serde_json::to_string_pretty(&json_val) {
+                    json_str
+                } else {
+                    "{}".to_string()
+                };
+
+                Ok(Value::string(json_string, pipeline.span).into_pipeline_data())
+            } else {
+                // let mut rec: Record = Record::new();
+                let mut rec = vec![];
+                for (span, shape) in flat {
+                    let content =
+                        String::from_utf8_lossy(working_set.get_span_contents(span)).to_string();
+                    let each_rec = record! {
+                        "content" => Value::test_string(content),
+                        "shape" => Value::test_string(shape.to_string()),
+                        "span" => Value::test_record(record!{
+                            "start" => Value::test_int(match span.start.checked_sub(offset) {
+                                Some(start) => start as i64,
+                                None => 0
+                            }),
+                            "end" => Value::test_int(match span.end.checked_sub(offset) {
+                                Some(end) => end as i64,
+                                None => 0
+                            }),
+                        }),
+                    };
+                    rec.push(Value::test_record(each_rec));
+                }
+                Ok(Value::list(rec, pipeline.span).into_pipeline_data())
+            }
+        } else {
             let error_output = working_set.parse_errors.first();
-        let block_span = match &block_output.span {
+            let block_span = match &parsed_block.span {
                 Some(span) => span,
                 None => &pipeline.span,
             };
             if to_json {
                 // Get the block as json
                 let serde_block_str = if minify {
-                serde_json::to_string(&*block_output)
+                    serde_json::to_string(&*parsed_block)
                 } else {
-                serde_json::to_string_pretty(&*block_output)
+                    serde_json::to_string_pretty(&*parsed_block)
                 };
                 let block_json = match serde_block_str {
                     Ok(json) => json,
@@ -59,7 +224,7 @@ impl Command for Ast {
                         from_type: "block".to_string(),
                         span: *block_span,
                         help: Some(format!(
-                        "Error: {e}\nCan't convert {block_output:?} to string"
+                            "Error: {e}\nCan't convert {parsed_block:?} to string"
                         )),
                     })?,
                 };
@@ -94,9 +259,9 @@ impl Command for Ast {
             } else {
                 let block_value = Value::string(
                     if minify {
-                    format!("{block_output:?}")
+                        format!("{parsed_block:?}")
                     } else {
-                    format!("{block_output:#?}")
+                        format!("{parsed_block:#?}")
                     },
                     pipeline.span,
                 );
@@ -118,36 +283,25 @@ impl Command for Ast {
                 Ok(output_record.into_pipeline_data())
             }
         }
+    }
+}
 
-    fn examples(&self) -> Vec<Example> {
+fn json_merge(a: &mut JsonValue, b: &JsonValue) {
-        vec![
+    match (a, b) {
-            Example {
+        (JsonValue::Object(ref mut a), JsonValue::Object(b)) => {
-                description: "Print the ast of a string",
+            for (k, v) in b {
-                example: "ast 'hello'",
+                json_merge(a.entry(k).or_insert(JsonValue::Null), v);
-                result: None,
+            }
-            },
+        }
-            Example {
+        (JsonValue::Array(ref mut a), JsonValue::Array(b)) => {
-                description: "Print the ast of a pipeline",
+            a.extend(b.clone());
-                example: "ast 'ls | where name =~ README'",
+        }
-                result: None,
+        (JsonValue::Array(ref mut a), JsonValue::Object(b)) => {
-            },
+            a.extend([JsonValue::Object(b.clone())]);
-            Example {
+        }
-                description: "Print the ast of a pipeline with an error",
+        (a, b) => {
-                example: "ast 'for x in 1..10 { echo $x '",
+            *a = b.clone();
-                result: None,
+        }
-            },
-            Example {
-                description:
-                    "Print the ast of a pipeline with an error, as json, in a nushell table",
-                example: "ast 'for x in 1..10 { echo $x ' --json | get block | from json",
-                result: None,
-            },
-            Example {
-                description: "Print the ast of a pipeline with an error, as json, minified",
-                example: "ast 'for x in 1..10 { echo $x ' --json --minify",
-                result: None,
-            },
-        ]
     }
 }