From a4b8d4a09869fb2565eb9ded67206d4d54ace3e0 Mon Sep 17 00:00:00 2001
From: Jonathan Turner <jonathandturner@users.noreply.github.com>
Date: Sat, 23 Jan 2021 10:28:32 +1300
Subject: [PATCH] Add rest support to blocks (#2962)

---
 crates/nu-engine/src/whole_stream_command.rs | 56 ++++++++++++++++----
 tests/shell/pipeline/commands/internal.rs    | 36 +++++++++++++
 2 files changed, 83 insertions(+), 9 deletions(-)

diff --git a/crates/nu-engine/src/whole_stream_command.rs b/crates/nu-engine/src/whole_stream_command.rs
index ac1fccb4a7..1dea22261a 100644
--- a/crates/nu-engine/src/whole_stream_command.rs
+++ b/crates/nu-engine/src/whole_stream_command.rs
@@ -8,7 +8,7 @@ use nu_errors::ShellError;
 use nu_parser::ParserScope;
 use nu_protocol::hir::Block;
 use nu_protocol::{ReturnSuccess, Signature, UntaggedValue};
-use nu_source::{b, DebugDocBuilder, PrettyDebugWithSource, Tag};
+use nu_source::{b, DebugDocBuilder, PrettyDebugWithSource, Span, Tag};
 use nu_stream::{OutputStream, ToOutputStream};
 use std::sync::Arc;
 
@@ -67,16 +67,54 @@ impl WholeStreamCommand for Block {
         let input = args.input;
         ctx.scope.enter_scope();
         if let Some(args) = evaluated.args.positional {
-            // FIXME: do not do this
-            for arg in args.into_iter().zip(self.params.positional.iter()) {
-                let name = arg.1 .0.name();
-
-                if name.starts_with('$') {
-                    ctx.scope.add_var(name, arg.0);
-                } else {
-                    ctx.scope.add_var(format!("${}", name), arg.0);
+            let mut args_iter = args.into_iter().peekable();
+            let mut params_iter = self.params.positional.iter();
+            loop {
+                match (args_iter.peek(), params_iter.next()) {
+                    (Some(_), Some(param)) => {
+                        let name = param.0.name();
+                        // we just checked the peek above, so this should be infallible
+                        if let Some(arg) = args_iter.next() {
+                            if name.starts_with('$') {
+                                ctx.scope.add_var(name.to_string(), arg);
+                            } else {
+                                ctx.scope.add_var(format!("${}", name), arg);
+                            }
+                        }
+                    }
+                    (Some(arg), None) => {
+                        if block.params.rest_positional.is_none() {
+                            ctx.scope.exit_scope();
+                            return Err(ShellError::labeled_error(
+                                "Unexpected argument to command",
+                                "unexpected argument",
+                                arg.tag.span,
+                            ));
+                        } else {
+                            break;
+                        }
+                    }
+                    _ => break,
                 }
             }
+            if block.params.rest_positional.is_some() {
+                let elements: Vec<_> = args_iter.collect();
+                let start = if let Some(first) = elements.first() {
+                    first.tag.span.start()
+                } else {
+                    0
+                };
+                let end = if let Some(last) = elements.last() {
+                    last.tag.span.end()
+                } else {
+                    0
+                };
+
+                ctx.scope.add_var(
+                    "$rest",
+                    UntaggedValue::Table(elements).into_value(Span::new(start, end)),
+                );
+            }
         }
         if let Some(args) = evaluated.args.named {
             for named in &block.params.named {
diff --git a/tests/shell/pipeline/commands/internal.rs b/tests/shell/pipeline/commands/internal.rs
index c6168f5a69..5b5e58c58f 100644
--- a/tests/shell/pipeline/commands/internal.rs
+++ b/tests/shell/pipeline/commands/internal.rs
@@ -428,6 +428,42 @@ fn run_broken_inner_custom_command() {
     assert!(actual.err.contains("not found"));
 }
 
+#[test]
+fn run_custom_command_with_rest() {
+    let actual = nu!(
+        cwd: ".",
+        r#"
+            def rest-me [...rest: string] { echo $rest.1 $rest.0}; rest-me "hello" "world" | to json
+        "#
+    );
+
+    assert_eq!(actual.out, r#"["world","hello"]"#);
+}
+
+#[test]
+fn run_custom_command_with_rest_and_arg() {
+    let actual = nu!(
+        cwd: ".",
+        r#"
+            def rest-me-with-arg [name: string, ...rest: string] { echo $rest.1 $rest.0 $name}; rest-me-with-arg "hello" "world" "yay" | to json
+        "#
+    );
+
+    assert_eq!(actual.out, r#"["yay","world","hello"]"#);
+}
+
+#[test]
+fn run_custom_command_with_rest_and_flag() {
+    let actual = nu!(
+        cwd: ".",
+        r#"
+            def rest-me-with-flag [--name: string, ...rest: string] { echo $rest.1 $rest.0 $name}; rest-me-with-flag "hello" "world" --name "yay" | to json
+        "#
+    );
+
+    assert_eq!(actual.out, r#"["world","hello","yay"]"#);
+}
+
 #[test]
 fn set_variable() {
     let actual = nu!(