From 9e3c64aa845ae99c87800501e7d63e2395c748ea Mon Sep 17 00:00:00 2001 From: WindSoilder Date: Mon, 11 Jul 2022 19:26:00 +0800 Subject: [PATCH] Add bytes collect, bytes remove, bytes build cmd (#6008) * add bytes collect * index_of support searching from end * add bytes remove * make bytes replace work better for empty pattern * add bytes build * remove comment * tweak words Co-authored-by: Darren Schroeder <343840+fdncred@users.noreply.github.com> --- crates/nu-command/src/bytes/build_.rs | 81 +++++++++ crates/nu-command/src/bytes/collect.rs | 127 +++++++++++++ crates/nu-command/src/bytes/index_of.rs | 73 +++++--- crates/nu-command/src/bytes/mod.rs | 6 + crates/nu-command/src/bytes/remove.rs | 196 +++++++++++++++++++++ crates/nu-command/src/bytes/replace.rs | 10 +- crates/nu-command/src/default_context.rs | 3 + crates/nu-protocol/src/value/from_value.rs | 21 +++ 8 files changed, 492 insertions(+), 25 deletions(-) create mode 100644 crates/nu-command/src/bytes/build_.rs create mode 100644 crates/nu-command/src/bytes/collect.rs create mode 100644 crates/nu-command/src/bytes/remove.rs diff --git a/crates/nu-command/src/bytes/build_.rs b/crates/nu-command/src/bytes/build_.rs new file mode 100644 index 000000000..4b4f0b623 --- /dev/null +++ b/crates/nu-command/src/bytes/build_.rs @@ -0,0 +1,81 @@ +use nu_engine::eval_expression; +use nu_protocol::ast::Call; +use nu_protocol::engine::{Command, EngineState, Stack}; +use nu_protocol::{ + Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, SyntaxShape, + Value, +}; + +#[derive(Clone)] +pub struct BytesBuild; + +impl Command for BytesBuild { + fn name(&self) -> &str { + "bytes build" + } + + fn usage(&self) -> &str { + "Create bytes from the arguments." + } + + fn search_terms(&self) -> Vec<&str> { + vec!["concatenate", "join"] + } + + fn signature(&self) -> nu_protocol::Signature { + Signature::build("bytes build") + .rest("rest", SyntaxShape::Any, "list of bytes") + .category(Category::Bytes) + } + + fn examples(&self) -> Vec { + vec![Example { + example: "bytes build 0x[01 02] 0x[03] 0x[04]", + description: "Builds binary data from 0x[01 02], 0x[03], 0x[04]", + result: Some(Value::Binary { + val: vec![0x01, 0x02, 0x03, 0x04], + span: Span::test_data(), + }), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + _input: PipelineData, + ) -> Result { + let mut output = vec![]; + for expr in call.positional_iter() { + let val = eval_expression(engine_state, stack, expr)?; + match val { + Value::Binary { mut val, .. } => output.append(&mut val), + other => { + return Err(ShellError::UnsupportedInput( + "only support expression which yields to binary data".to_string(), + other.span().unwrap_or(call.head), + )) + } + } + } + + Ok(Value::Binary { + val: output, + span: call.head, + } + .into_pipeline_data()) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(BytesBuild {}) + } +} diff --git a/crates/nu-command/src/bytes/collect.rs b/crates/nu-command/src/bytes/collect.rs new file mode 100644 index 000000000..34ddc4470 --- /dev/null +++ b/crates/nu-command/src/bytes/collect.rs @@ -0,0 +1,127 @@ +use nu_engine::CallExt; +use nu_protocol::ast::Call; +use nu_protocol::engine::{Command, EngineState, Stack}; +use nu_protocol::{ + Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, SyntaxShape, + Value, +}; + +#[derive(Clone, Copy)] +pub struct BytesCollect; + +impl Command for BytesCollect { + fn name(&self) -> &str { + "bytes collect" + } + + fn signature(&self) -> Signature { + Signature::build("bytes collect") + .optional( + "separator", + SyntaxShape::Binary, + "optional separator to use when creating binary", + ) + .category(Category::Bytes) + } + + fn usage(&self) -> &str { + "Concatenate multiple binary into a single binary, with an optional separator between each" + } + + fn search_terms(&self) -> Vec<&str> { + vec!["join", "concatenate"] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let separator: Option> = call.opt(engine_state, stack, 0)?; + // input should be a list of binary data. + let mut output_binary = vec![]; + for value in input { + match value { + Value::Binary { mut val, .. } => { + output_binary.append(&mut val); + // manually concat + // TODO: make use of std::slice::Join when it's available in stable. + if let Some(sep) = &separator { + let mut work_sep = sep.clone(); + output_binary.append(&mut work_sep) + } + } + other => { + return Err(ShellError::UnsupportedInput( + format!( + "The element type is {}, this command only works with bytes.", + other.get_type() + ), + other.span().unwrap_or(call.head), + )) + } + } + } + + match separator { + None => Ok(Value::Binary { + val: output_binary, + span: call.head, + } + .into_pipeline_data()), + Some(sep) => { + if output_binary.is_empty() { + Ok(Value::Binary { + val: output_binary, + span: call.head, + } + .into_pipeline_data()) + } else { + // have push one extra separator in previous step, pop them out. + for _ in sep { + let _ = output_binary.pop(); + } + Ok(Value::Binary { + val: output_binary, + span: call.head, + } + .into_pipeline_data()) + } + } + } + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Create a byte array from input", + example: "[0x[11] 0x[13 15]] | bytes collect", + result: Some(Value::Binary { + val: vec![0x11, 0x13, 0x15], + span: Span::test_data(), + }), + }, + Example { + description: "Create a byte array from input with a separator", + example: "[0x[11] 0x[33] 0x[44]] | bytes collect 0x[01]", + result: Some(Value::Binary { + val: vec![0x11, 0x01, 0x33, 0x01, 0x44], + span: Span::test_data(), + }), + }, + ] + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(BytesCollect {}) + } +} diff --git a/crates/nu-command/src/bytes/index_of.rs b/crates/nu-command/src/bytes/index_of.rs index dc6fa677e..335e3b1ee 100644 --- a/crates/nu-command/src/bytes/index_of.rs +++ b/crates/nu-command/src/bytes/index_of.rs @@ -95,6 +95,14 @@ impl Command for BytesIndexOf { span: Span::test_data(), }), }, + Example { + description: "Returns all matched index, searching from end", + example: " 0x[33 44 55 10 01 33 44 33 44] | bytes index-of -a -e 0x[33 44]", + result: Some(Value::List { + vals: vec![Value::test_int(7), Value::test_int(5), Value::test_int(0)], + span: Span::test_data(), + }), + }, Example { description: "Returns index of pattern for specific column", example: r#" [[ColA ColB ColC]; [0x[11 12 13] 0x[14 15 16] 0x[17 18 19]]] | bytes index-of 0x[11] ColA ColC"#, @@ -119,27 +127,8 @@ impl Command for BytesIndexOf { } fn index_of(input: &[u8], arg: &Arguments, span: Span) -> Value { - // currently, `--all` flag doesn't support finding from end. if arg.all { - let mut result = vec![]; - // doing find stuff. - let (mut left, mut right) = (0, arg.pattern.len()); - let input_len = input.len(); - let pattern_len = arg.pattern.len(); - while right <= input_len { - if input[left..right] == arg.pattern { - result.push(Value::Int { - val: left as i64, - span, - }); - left += pattern_len; - right += pattern_len; - } else { - left += 1; - right += 1; - } - } - Value::List { vals: result, span } + search_all_index(input, &arg.pattern, arg.end, span) } else { let mut iter = input.windows(arg.pattern.len()); @@ -164,6 +153,50 @@ fn index_of(input: &[u8], arg: &Arguments, span: Span) -> Value { } } +fn search_all_index(input: &[u8], pattern: &[u8], from_end: bool, span: Span) -> Value { + let mut result = vec![]; + if from_end { + let (mut left, mut right) = ( + input.len() as isize - pattern.len() as isize, + input.len() as isize, + ); + while left >= 0 { + if &input[left as usize..right as usize] == pattern { + result.push(Value::Int { + val: left as i64, + span, + }); + left -= pattern.len() as isize; + right -= pattern.len() as isize; + } else { + left -= 1; + right -= 1; + } + } + Value::List { vals: result, span } + } else { + // doing find stuff. + let (mut left, mut right) = (0, pattern.len()); + let input_len = input.len(); + let pattern_len = pattern.len(); + while right <= input_len { + if &input[left..right] == pattern { + result.push(Value::Int { + val: left as i64, + span, + }); + left += pattern_len; + right += pattern_len; + } else { + left += 1; + right += 1; + } + } + + Value::List { vals: result, span } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/nu-command/src/bytes/mod.rs b/crates/nu-command/src/bytes/mod.rs index 3aef3a4b0..a8dca8180 100644 --- a/crates/nu-command/src/bytes/mod.rs +++ b/crates/nu-command/src/bytes/mod.rs @@ -1,9 +1,12 @@ mod add; mod at; +mod build_; mod bytes_; +mod collect; mod ends_with; mod index_of; mod length; +mod remove; mod replace; mod reverse; mod starts_with; @@ -15,10 +18,13 @@ use std::sync::Arc; pub use add::BytesAdd; pub use at::BytesAt; +pub use build_::BytesBuild; pub use bytes_::Bytes; +pub use collect::BytesCollect; pub use ends_with::BytesEndsWith; pub use index_of::BytesIndexOf; pub use length::BytesLen; +pub use remove::BytesRemove; pub use replace::BytesReplace; pub use reverse::BytesReverse; pub use starts_with::BytesStartsWith; diff --git a/crates/nu-command/src/bytes/remove.rs b/crates/nu-command/src/bytes/remove.rs new file mode 100644 index 000000000..e7d66ca5e --- /dev/null +++ b/crates/nu-command/src/bytes/remove.rs @@ -0,0 +1,196 @@ +use super::{operate, BytesArgument}; +use nu_engine::CallExt; +use nu_protocol::{ + ast::{Call, CellPath}, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, Value, +}; + +struct Arguments { + pattern: Vec, + end: bool, + column_paths: Option>, + all: bool, +} + +impl BytesArgument for Arguments { + fn take_column_paths(&mut self) -> Option> { + self.column_paths.take() + } +} + +#[derive(Clone)] +pub struct BytesRemove; + +impl Command for BytesRemove { + fn name(&self) -> &str { + "bytes remove" + } + + fn signature(&self) -> Signature { + Signature::build("bytes remove") + .required("pattern", SyntaxShape::Binary, "the pattern to find") + .rest( + "rest", + SyntaxShape::CellPath, + "optionally remove bytes by column paths", + ) + .switch("end", "remove from end of binary", Some('e')) + .switch("all", "remove occurrences of finding binary", Some('a')) + .category(Category::Bytes) + } + + fn usage(&self) -> &str { + "remove bytes" + } + + fn search_terms(&self) -> Vec<&str> { + vec!["search", "shift", "switch"] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let column_paths: Vec = call.rest(engine_state, stack, 1)?; + let column_paths = if column_paths.is_empty() { + None + } else { + Some(column_paths) + }; + let pattern_to_remove = call.req::>>(engine_state, stack, 0)?; + if pattern_to_remove.item.is_empty() { + return Err(ShellError::UnsupportedInput( + "the pattern to remove cannot be empty".to_string(), + pattern_to_remove.span, + )); + } + + let pattern_to_remove: Vec = pattern_to_remove.item; + let arg = Arguments { + pattern: pattern_to_remove, + end: call.has_flag("end"), + column_paths, + all: call.has_flag("all"), + }; + + operate(remove, arg, input, call.head, engine_state.ctrlc.clone()) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Remove contents", + example: "0x[10 AA FF AA FF] | bytes remove 0x[10 AA]", + result: Some(Value::Binary { + val: vec![0xFF, 0xAA, 0xFF], + span: Span::test_data(), + }), + }, + Example { + description: "Remove all occurrences of find binary", + example: "0x[10 AA 10 BB 10] | bytes remove -a 0x[10]", + result: Some(Value::Binary { + val: vec![0xAA, 0xBB], + span: Span::test_data(), + }), + }, + Example { + description: "Remove occurrences of find binary from end", + example: "0x[10 AA 10 BB CC AA 10] | bytes remove -e 0x[10]", + result: Some(Value::Binary { + val: vec![0x10, 0xAA, 0x10, 0xBB, 0xCC, 0xAA], + span: Span::test_data(), + }), + }, + Example { + description: "Remove all occurrences of find binary in table", + example: "[[ColA ColB ColC]; [0x[11 12 13] 0x[14 15 16] 0x[17 18 19]]] | bytes remove 0x[11] ColA ColC", + result: Some(Value::List { + vals: vec![Value::Record { + cols: vec!["ColA".to_string(), "ColB".to_string(), "ColC".to_string()], + vals: vec![ + Value::Binary { + val: vec![0x12, 0x13], + span: Span::test_data(), + }, + Value::Binary { + val: vec![0x14, 0x15, 0x16], + span: Span::test_data(), + }, + Value::Binary { + val: vec![0x17, 0x18, 0x19], + span: Span::test_data(), + }, + ], + span: Span::test_data(), + }], + span: Span::test_data(), + }), + }, + ] + } +} + +fn remove(input: &[u8], arg: &Arguments, span: Span) -> Value { + let mut result = vec![]; + let remove_all = arg.all; + let input_len = input.len(); + let pattern_len = arg.pattern.len(); + + // Note: + // remove_all from start and end will generate the same result. + // so we'll put `remove_all` relative logic into else clouse. + if arg.end && !remove_all { + let (mut left, mut right) = ( + input.len() as isize - arg.pattern.len() as isize, + input.len() as isize, + ); + while left >= 0 && input[left as usize..right as usize] != arg.pattern { + result.push(input[right as usize - 1]); + left -= 1; + right -= 1; + } + // append the remaining thing to result, this can be happeneed when + // we have something to remove and remove_all is False. + let mut remain = input[..left as usize].iter().copied().rev().collect(); + result.append(&mut remain); + result = result.into_iter().rev().collect(); + Value::Binary { val: result, span } + } else { + let (mut left, mut right) = (0, arg.pattern.len()); + while right <= input_len { + if input[left..right] == arg.pattern { + left += pattern_len; + right += pattern_len; + if !remove_all { + break; + } + } else { + result.push(input[left]); + left += 1; + right += 1; + } + } + // append the remaing thing to result, this can happened when + // we have something to remove and remove_all is False. + let mut remain = input[left..].to_vec(); + result.append(&mut remain); + Value::Binary { val: result, span } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(BytesRemove {}) + } +} diff --git a/crates/nu-command/src/bytes/replace.rs b/crates/nu-command/src/bytes/replace.rs index 69702915b..fde42567a 100644 --- a/crates/nu-command/src/bytes/replace.rs +++ b/crates/nu-command/src/bytes/replace.rs @@ -3,7 +3,7 @@ use nu_engine::CallExt; use nu_protocol::{ ast::{Call, CellPath}, engine::{Command, EngineState, Stack}, - Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, + Category, Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, Value, }; struct Arguments { @@ -61,16 +61,16 @@ impl Command for BytesReplace { } else { Some(column_paths) }; - let find = call.req::>(engine_state, stack, 0)?; - if find.is_empty() { + let find = call.req::>>(engine_state, stack, 0)?; + if find.item.is_empty() { return Err(ShellError::UnsupportedInput( "the pattern to find cannot be empty".to_string(), - call.head, + find.span, )); } let arg = Arguments { - find, + find: find.item, replace: call.req::>(engine_state, stack, 1)?, column_paths, all: call.has_flag("all"), diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index ea5b326ab..ac81cae3f 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -218,6 +218,9 @@ pub fn create_default_context(cwd: impl AsRef) -> EngineState { BytesAdd, BytesAt, BytesIndexOf, + BytesCollect, + BytesRemove, + BytesBuild, } // FileSystem diff --git a/crates/nu-protocol/src/value/from_value.rs b/crates/nu-protocol/src/value/from_value.rs index 66b7b89a5..c2bd16208 100644 --- a/crates/nu-protocol/src/value/from_value.rs +++ b/crates/nu-protocol/src/value/from_value.rs @@ -404,6 +404,27 @@ impl FromValue for Vec { } } +impl FromValue for Spanned> { + fn from_value(v: &Value) -> Result { + match v { + Value::Binary { val, span } => Ok(Spanned { + item: val.clone(), + span: *span, + }), + Value::String { val, span } => Ok(Spanned { + item: val.bytes().collect(), + span: *span, + }), + v => Err(ShellError::CantConvert( + "binary data".into(), + v.get_type().to_string(), + v.span()?, + None, + )), + } + } +} + impl FromValue for Spanned { fn from_value(v: &Value) -> Result { match v {