From 99873dda8c9e6e4033e083efd25d5b443fd1c335 Mon Sep 17 00:00:00 2001
From: Devyn Cairns <devyn.cairns@gmail.com>
Date: Thu, 11 Jul 2024 03:38:51 -0700
Subject: [PATCH] Add more argument types to `view ir`

---
 crates/nu-command/src/debug/view_ir.rs | 101 +++++++++++++++++++++++--
 1 file changed, 94 insertions(+), 7 deletions(-)

diff --git a/crates/nu-command/src/debug/view_ir.rs b/crates/nu-command/src/debug/view_ir.rs
index df4f6cad6b..35c99125c6 100644
--- a/crates/nu-command/src/debug/view_ir.rs
+++ b/crates/nu-command/src/debug/view_ir.rs
@@ -1,5 +1,4 @@
 use nu_engine::command_prelude::*;
-use nu_protocol::engine::Closure;
 
 #[derive(Clone)]
 pub struct ViewIr;
@@ -12,15 +11,20 @@ impl Command for ViewIr {
     fn signature(&self) -> Signature {
         Signature::new(self.name())
             .required(
-                "closure",
-                SyntaxShape::Closure(None),
-                "The closure to see compiled code for.",
+                "target",
+                SyntaxShape::Any,
+                "The name or block to view compiled code for.",
             )
             .switch(
                 "json",
                 "Dump the raw block data as JSON (unstable).",
                 Some('j'),
             )
+            .switch(
+                "decl-id",
+                "Integer is a declaration ID rather than a block ID.",
+                Some('d'),
+            )
             .input_output_type(Type::Nothing, Type::String)
     }
 
@@ -28,6 +32,20 @@ impl Command for ViewIr {
         "View the compiled IR code for a block of code."
     }
 
+    fn extra_usage(&self) -> &str {
+        "
+The target can be a closure, the name of a custom command, or an internal block
+ID. Closure literals within IR dumps often reference the block by ID (e.g.
+`closure(3231)`), so this provides an easy way to read the IR of any embedded
+closures.
+
+The --decl-id option is provided to use a declaration ID instead, which can be
+found on `call` instructions. This is sometimes better than using the name, as
+the declaration may not be in scope.
+"
+        .trim()
+    }
+
     fn run(
         &self,
         engine_state: &EngineState,
@@ -35,10 +53,79 @@ impl Command for ViewIr {
         call: &Call,
         _input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
-        let closure: Closure = call.req(engine_state, stack, 0)?;
+        let target: Value = call.req(engine_state, stack, 0)?;
         let json = call.has_flag(engine_state, stack, "json")?;
+        let is_decl_id = call.has_flag(engine_state, stack, "decl-id")?;
+
+        let block_id = match target {
+            Value::Closure { ref val, .. } => val.block_id,
+            // Decl by name
+            Value::String { ref val, .. } => {
+                if let Some(decl_id) = engine_state.find_decl(val.as_bytes(), &[]) {
+                    let decl = engine_state.get_decl(decl_id);
+                    decl.block_id().ok_or_else(|| ShellError::GenericError {
+                        error: format!("Can't view IR for `{val}`"),
+                        msg: "not a custom command".into(),
+                        span: Some(target.span()),
+                        help: Some("internal commands don't have Nushell source code".into()),
+                        inner: vec![],
+                    })?
+                } else {
+                    return Err(ShellError::GenericError {
+                        error: format!("Can't view IR for `{val}`"),
+                        msg: "can't find a command with this name".into(),
+                        span: Some(target.span()),
+                        help: None,
+                        inner: vec![],
+                    });
+                }
+            }
+            // Decl by ID - IR dump always shows name of decl, but sometimes it isn't in scope
+            Value::Int { val, .. } if is_decl_id => {
+                let decl_id = val
+                    .try_into()
+                    .ok()
+                    .filter(|id| *id < engine_state.num_decls())
+                    .ok_or_else(|| ShellError::IncorrectValue {
+                        msg: "not a valid decl id".into(),
+                        val_span: target.span(),
+                        call_span: call.head,
+                    })?;
+                let decl = engine_state.get_decl(decl_id);
+                decl.block_id().ok_or_else(|| ShellError::GenericError {
+                    error: format!("Can't view IR for `{}`", decl.name()),
+                    msg: "not a custom command".into(),
+                    span: Some(target.span()),
+                    help: Some("internal commands don't have Nushell source code".into()),
+                    inner: vec![],
+                })?
+            }
+            // Block by ID - often shows up in IR
+            Value::Int { val, .. } => val.try_into().map_err(|_| ShellError::IncorrectValue {
+                msg: "not a valid block id".into(),
+                val_span: target.span(),
+                call_span: call.head,
+            })?,
+            // Pass through errors
+            Value::Error { error, .. } => return Err(*error),
+            _ => {
+                return Err(ShellError::TypeMismatch {
+                    err_message: "expected closure, string, or int".into(),
+                    span: call.head,
+                })
+            }
+        };
+
+        let Some(block) = engine_state.try_get_block(block_id) else {
+            return Err(ShellError::GenericError {
+                error: format!("Unknown block ID: {block_id}"),
+                msg: "ensure the block ID is correct and try again".into(),
+                span: Some(target.span()),
+                help: None,
+                inner: vec![],
+            });
+        };
 
-        let block = engine_state.get_block(closure.block_id);
         let ir_block = block
             .ir_block
             .as_ref()
@@ -62,7 +149,7 @@ impl Command for ViewIr {
                 .collect::<Vec<_>>();
 
             serde_json::to_string_pretty(&serde_json::json!({
-                "block_id": closure.block_id,
+                "block_id": block_id,
                 "span": block.span,
                 "ir_block": ir_block,
                 "formatted_instructions": formatted_instructions,