diff --git a/crates/nu-command/src/commands/nth.rs b/crates/nu-command/src/commands/nth.rs index 8f7c053e50..9801aa4d8c 100644 --- a/crates/nu-command/src/commands/nth.rs +++ b/crates/nu-command/src/commands/nth.rs @@ -1,7 +1,7 @@ use crate::prelude::*; use nu_engine::WholeStreamCommand; use nu_errors::ShellError; -use nu_protocol::{ReturnSuccess, Signature, SyntaxShape, Value}; +use nu_protocol::{Signature, SyntaxShape, Value}; use nu_source::Tagged; #[derive(Deserialize)] @@ -33,7 +33,7 @@ impl WholeStreamCommand for Nth { "Return or skip only the selected rows." } - fn run_with_actions(&self, args: CommandArgs) -> Result { + fn run(&self, args: CommandArgs) -> Result { nth(args) } @@ -58,7 +58,7 @@ impl WholeStreamCommand for Nth { } } -fn nth(args: CommandArgs) -> Result { +fn nth(args: CommandArgs) -> Result { let ( NthArgs { row_number, @@ -68,22 +68,60 @@ fn nth(args: CommandArgs) -> Result { input, ) = args.process()?; - let row_numbers = vec![vec![row_number], and_rows] - .into_iter() - .flatten() - .map(|x| x.item) - .collect::>(); + let mut rows: Vec<_> = and_rows.into_iter().map(|x| x.item as usize).collect(); + rows.push(row_number.item as usize); + rows.sort_unstable(); - Ok(input - .enumerate() - .filter_map(move |(idx, item)| { - if row_numbers.contains(&(idx as u64)) ^ skip { - Some(ReturnSuccess::value(item)) + Ok(NthIterator { + input, + rows, + skip, + current: 0, + } + .to_output_stream()) +} + +struct NthIterator { + input: InputStream, + rows: Vec, + skip: bool, + current: usize, +} + +impl Iterator for NthIterator { + type Item = Value; + + fn next(&mut self) -> Option { + loop { + if !self.skip { + if let Some(row) = self.rows.get(0) { + if self.current == *row { + self.rows.remove(0); + self.current += 1; + return self.input.next(); + } else { + self.current += 1; + let _ = self.input.next(); + continue; + } + } else { + return None; + } + } else if let Some(row) = self.rows.get(0) { + if self.current == *row { + self.rows.remove(0); + self.current += 1; + let _ = self.input.next(); + continue; + } else { + self.current += 1; + return self.input.next(); + } } else { - None + return self.input.next(); } - }) - .to_action_stream()) + } + } } #[cfg(test)]