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

@@ -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<Example> {
+        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<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        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 {})
+    }
+}

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<PipelineData, ShellError> {
+        let separator: Option<Vec<u8>> = 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<Example> {
+        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 {})
+    }
+}

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::*;

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;

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<u8>,
+    end: bool,
+    column_paths: Option<Vec<CellPath>>,
+    all: bool,
+}
+
+impl BytesArgument for Arguments {
+    fn take_column_paths(&mut self) -> Option<Vec<CellPath>> {
+        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<PipelineData, ShellError> {
+        let column_paths: Vec<CellPath> = 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::<Spanned<Vec<u8>>>(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<u8> = 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<Example> {
+        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 {})
+    }
+}

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::<Vec<u8>>(engine_state, stack, 0)?;
-        if find.is_empty() {
+        let find = call.req::<Spanned<Vec<u8>>>(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::<Vec<u8>>(engine_state, stack, 1)?,
             column_paths,
             all: call.has_flag("all"),

crates/nu-command/src/default_context.rs

@@ -218,6 +218,9 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
             BytesAdd,
             BytesAt,
             BytesIndexOf,
+            BytesCollect,
+            BytesRemove,
+            BytesBuild,
         }
 
         // FileSystem

crates/nu-protocol/src/value/from_value.rs

@@ -404,6 +404,27 @@ impl FromValue for Vec<u8> {
     }
 }
 
+impl FromValue for Spanned<Vec<u8>> {
+    fn from_value(v: &Value) -> Result<Self, ShellError> {
+        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<PathBuf> {
     fn from_value(v: &Value) -> Result<Self, ShellError> {
         match v {