diff --git a/crates/nu-cli/src/errors.rs b/crates/nu-cli/src/errors.rs index d35f0c9ede..a75d06ae99 100644 --- a/crates/nu-cli/src/errors.rs +++ b/crates/nu-cli/src/errors.rs @@ -312,6 +312,13 @@ pub fn report_shell_error( .with_labels(vec![Label::primary(diag_file_id, diag_range) .with_message(format!("can't convert to {}", s))]) } + ShellError::CannotCreateRange(span) => { + let (diag_file_id, diag_range) = convert_span_to_diag(working_set, span)?; + Diagnostic::error() + .with_message("Can't convert range to countable values") + .with_labels(vec![Label::primary(diag_file_id, diag_range) + .with_message("can't convert to countable values")]) + } ShellError::DivisionByZero(span) => { let (diag_file_id, diag_range) = convert_span_to_diag(working_set, span)?; diff --git a/crates/nu-command/src/each.rs b/crates/nu-command/src/each.rs index c1e2a5cb80..fbb24506d7 100644 --- a/crates/nu-command/src/each.rs +++ b/crates/nu-command/src/each.rs @@ -30,8 +30,8 @@ impl Command for Each { let context = context.clone(); match input { - Value::List { val, .. } => Ok(Value::List { - val: val + Value::Range { val, .. } => Ok(Value::ValueStream { + stream: val .into_iter() .map(move |x| { let engine_state = context.engine_state.borrow(); @@ -46,10 +46,32 @@ impl Command for Each { match eval_block(&state, block, Value::nothing()) { Ok(v) => v, - Err(err) => Value::Error { err }, + Err(error) => Value::Error { error }, } }) - .collect(), + .into_value_stream(), + span: call.head, + }), + Value::List { val, .. } => Ok(Value::ValueStream { + stream: val + .into_iter() + .map(move |x| { + let engine_state = context.engine_state.borrow(); + let block = engine_state.get_block(block_id); + + let state = context.enter_scope(); + if let Some(var) = block.signature.required_positional.first() { + if let Some(var_id) = &var.var_id { + state.add_var(*var_id, x); + } + } + + match eval_block(&state, block, Value::nothing()) { + Ok(v) => v, + Err(error) => Value::Error { error }, + } + }) + .into_value_stream(), span: call.head, }), Value::ValueStream { stream, .. } => Ok(Value::ValueStream { @@ -67,7 +89,7 @@ impl Command for Each { match eval_block(&state, block, Value::nothing()) { Ok(v) => v, - Err(err) => Value::Error { err }, + Err(error) => Value::Error { error }, } }) .into_value_stream(), diff --git a/crates/nu-protocol/src/shell_error.rs b/crates/nu-protocol/src/shell_error.rs index 78f672921d..92eb306434 100644 --- a/crates/nu-protocol/src/shell_error.rs +++ b/crates/nu-protocol/src/shell_error.rs @@ -16,4 +16,5 @@ pub enum ShellError { VariableNotFoundAtRuntime(Span), CantConvert(String, Span), DivisionByZero(Span), + CannotCreateRange(Span), } diff --git a/crates/nu-protocol/src/value.rs b/crates/nu-protocol/src/value.rs index 74a78aff63..09c3e57806 100644 --- a/crates/nu-protocol/src/value.rs +++ b/crates/nu-protocol/src/value.rs @@ -111,6 +111,120 @@ pub struct Range { pub inclusion: RangeInclusion, } +impl IntoIterator for Range { + type Item = Value; + + type IntoIter = RangeIterator; + + fn into_iter(self) -> Self::IntoIter { + let span = self.from.span(); + + RangeIterator::new(self, span) + } +} + +pub struct RangeIterator { + curr: Value, + end: Value, + span: Span, + is_end_inclusive: bool, + moves_up: bool, + one: Value, + negative_one: Value, + done: bool, +} + +impl RangeIterator { + pub fn new(range: Range, span: Span) -> RangeIterator { + let start = match range.from { + Value::Nothing { .. } => Value::Int { val: 0, span }, + x => x, + }; + + let end = match range.to { + Value::Nothing { .. } => Value::Int { + val: i64::MAX, + span, + }, + x => x, + }; + + RangeIterator { + moves_up: matches!(start.lte(span, &end), Ok(Value::Bool { val: true, .. })), + curr: start, + end, + span, + is_end_inclusive: matches!(range.inclusion, RangeInclusion::Inclusive), + done: false, + one: Value::Int { val: 1, span }, + negative_one: Value::Int { val: -1, span }, + } + } +} + +impl Iterator for RangeIterator { + type Item = Value; + fn next(&mut self) -> Option { + use std::cmp::Ordering; + if self.done { + return None; + } + + let ordering = if matches!(self.end, Value::Nothing { .. }) { + Ordering::Less + } else { + match (&self.curr, &self.end) { + (Value::Int { val: x, .. }, Value::Int { val: y, .. }) => x.cmp(y), + // (Value::Float { val: x, .. }, Value::Float { val: y, .. }) => x.cmp(y), + // (Value::Float { val: x, .. }, Value::Int { val: y, .. }) => x.cmp(y), + // (Value::Int { val: x, .. }, Value::Float { val: y, .. }) => x.cmp(y), + _ => { + self.done = true; + return Some(Value::Error { + error: ShellError::CannotCreateRange(self.span), + }); + } + } + }; + + if self.moves_up + && (ordering == Ordering::Less || self.is_end_inclusive && ordering == Ordering::Equal) + { + let next_value = self.curr.add(self.span, &self.one); + + let mut next = match next_value { + Ok(result) => result, + + Err(error) => { + self.done = true; + return Some(Value::Error { error }); + } + }; + std::mem::swap(&mut self.curr, &mut next); + + Some(next) + } else if !self.moves_up + && (ordering == Ordering::Greater + || self.is_end_inclusive && ordering == Ordering::Equal) + { + let next_value = self.curr.add(self.span, &self.negative_one); + + let mut next = match next_value { + Ok(result) => result, + Err(error) => { + self.done = true; + return Some(Value::Error { error }); + } + }; + std::mem::swap(&mut self.curr, &mut next); + + Some(next) + } else { + None + } + } +} + #[derive(Debug, Clone)] pub enum Value { Bool { @@ -159,7 +273,7 @@ pub enum Value { span: Span, }, Error { - err: ShellError, + error: ShellError, }, } @@ -270,7 +384,7 @@ impl Value { } => stream.into_string(headers), Value::Block { val, .. } => format!("", val), Value::Nothing { .. } => String::new(), - Value::Error { err } => format!("{:?}", err), + Value::Error { error } => format!("{:?}", error), } } diff --git a/src/tests.rs b/src/tests.rs index 820af046a9..13b1aa2a92 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -206,10 +206,20 @@ fn alias_2() -> TestResult { #[test] fn block_param1() -> TestResult { - run_test("[3] | each { $it + 10 }", "13") + run_test("[3] | each { $it + 10 }", "[13]") } #[test] fn block_param2() -> TestResult { - run_test("[3] | each { |y| $y + 10 }", "13") + run_test("[3] | each { |y| $y + 10 }", "[13]") +} + +#[test] +fn range_iteration1() -> TestResult { + run_test("1..4 | each { |y| $y + 10 }", "[11, 12, 13, 14]") +} + +#[test] +fn range_iteration2() -> TestResult { + run_test("4..1 | each { |y| $y + 100 }", "[104, 103, 102, 101]") }