Add bits shl and bits shr command (#6202)

* Add `bits shift-left` and `bits shift-right` command

* update bits shift error tips

* some code refactor

* update shift right

* some code refactor for bits shift commands

* rename bits shift commands align with bits operators

* update search term

* Update crates/nu-command/src/bits/not.rs

Co-authored-by: Stefan Holderbach <sholderbach@users.noreply.github.com>

* Update crates/nu-command/src/bits/shift_left.rs

Co-authored-by: Stefan Holderbach <sholderbach@users.noreply.github.com>

* Update crates/nu-command/src/bits/shift_right.rs

Co-authored-by: Stefan Holderbach <sholderbach@users.noreply.github.com>

* ci skip

* change default number-bytes for bits shift

* fix bits not tests

* fix bits tests

Co-authored-by: Stefan Holderbach <sholderbach@users.noreply.github.com>

Cargo.lock

@@ -2620,6 +2620,7 @@ dependencies = [
  "nu-test-support",
  "nu-utils",
  "num 0.4.0",
+ "num-traits",
  "pathdiff",
  "polars",
  "powierza-coefficient",

crates/nu-command/Cargo.toml

@@ -58,6 +58,7 @@ meval = "0.2.0"
 mime = "0.3.16"
 notify = "4.0.17"
 num = { version = "0.4.0", optional = true }
+num-traits = "0.2.14"
 pathdiff = "0.2.1"
 powierza-coefficient = "1.0.1"
 quick-xml = "0.23.0"

crates/nu-command/src/bits/mod.rs

@@ -2,10 +2,94 @@ mod and;
 mod bits_;
 mod not;
 mod or;
+mod shift_left;
+mod shift_right;
 mod xor;
 
+use nu_protocol::Spanned;
+
 pub use and::SubCommand as BitsAnd;
 pub use bits_::Bits;
 pub use not::SubCommand as BitsNot;
 pub use or::SubCommand as BitsOr;
+pub use shift_left::SubCommand as BitsShiftLeft;
+pub use shift_right::SubCommand as BitsShiftRight;
 pub use xor::SubCommand as BitsXor;
+
+#[derive(Clone, Copy)]
+enum NumberBytes {
+    One,
+    Two,
+    Four,
+    Eight,
+    Auto,
+    Invalid,
+}
+
+#[derive(Clone, Copy)]
+enum InputNumType {
+    One,
+    Two,
+    Four,
+    Eight,
+    SignedOne,
+    SignedTwo,
+    SignedFour,
+    SignedEight,
+}
+
+fn get_number_bytes(number_bytes: &Option<Spanned<String>>) -> NumberBytes {
+    match number_bytes.as_ref() {
+        None => NumberBytes::Eight,
+        Some(size) => match size.item.as_str() {
+            "1" => NumberBytes::One,
+            "2" => NumberBytes::Two,
+            "4" => NumberBytes::Four,
+            "8" => NumberBytes::Eight,
+            "auto" => NumberBytes::Auto,
+            _ => NumberBytes::Invalid,
+        },
+    }
+}
+
+fn get_input_num_type(val: i64, signed: bool, number_size: NumberBytes) -> InputNumType {
+    if signed || val < 0 {
+        match number_size {
+            NumberBytes::One => InputNumType::SignedOne,
+            NumberBytes::Two => InputNumType::SignedTwo,
+            NumberBytes::Four => InputNumType::SignedFour,
+            NumberBytes::Eight => InputNumType::SignedEight,
+            NumberBytes::Auto => {
+                if val <= 0x7F && val >= -(2i64.pow(7)) {
+                    InputNumType::SignedOne
+                } else if val <= 0x7FFF && val >= -(2i64.pow(15)) {
+                    InputNumType::SignedTwo
+                } else if val <= 0x7FFFFFFF && val >= -(2i64.pow(31)) {
+                    InputNumType::SignedFour
+                } else {
+                    InputNumType::SignedEight
+                }
+            }
+            NumberBytes::Invalid => InputNumType::SignedFour,
+        }
+    } else {
+        match number_size {
+            NumberBytes::One => InputNumType::One,
+            NumberBytes::Two => InputNumType::Two,
+            NumberBytes::Four => InputNumType::Four,
+            NumberBytes::Eight => InputNumType::Eight,
+            NumberBytes::Auto => {
+                if val <= 0xFF {
+                    InputNumType::One
+                } else if val <= 0xFFFF {
+                    InputNumType::Two
+                } else if val <= 0xFFFFFFFF {
+                    InputNumType::Four
+                } else {
+                    InputNumType::Eight
+                }
+            }
+            NumberBytes::Invalid => InputNumType::Four,
+        }
+    }
+}

crates/nu-command/src/bits/not.rs

@@ -1,3 +1,4 @@
+use super::{get_number_bytes, NumberBytes};
 use nu_engine::CallExt;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
@@ -8,15 +9,6 @@ use nu_protocol::{
 #[derive(Clone)]
 pub struct SubCommand;
 
-#[derive(Clone, Copy)]
-enum NumberBytes {
-    One,
-    Two,
-    Four,
-    Eight,
-    Auto,
-}
-
 impl Command for SubCommand {
     fn name(&self) -> &str {
         "bits not"
@@ -57,25 +49,18 @@ impl Command for SubCommand {
         let signed = call.has_flag("signed");
         let number_bytes: Option<Spanned<String>> =
             call.get_flag(engine_state, stack, "number-bytes")?;
-        let number_bytes = match number_bytes.as_ref() {
+        let bytes_len = get_number_bytes(&number_bytes);
-            None => NumberBytes::Auto,
+        if let NumberBytes::Invalid = bytes_len {
-            Some(size) => match size.item.as_str() {
+            if let Some(val) = number_bytes {
-                "1" => NumberBytes::One,
-                "2" => NumberBytes::Two,
-                "4" => NumberBytes::Four,
-                "8" => NumberBytes::Eight,
-                "auto" => NumberBytes::Auto,
-                _ => {
                 return Err(ShellError::UnsupportedInput(
-                        "the size of number is invalid".to_string(),
+                    "Only 1, 2, 4, 8, or 'auto' bytes are supported as word sizes".to_string(),
-                        size.span,
+                    val.span,
-                    ))
+                ));
+            }
         }
-            },
-        };
 
         input.map(
-            move |value| operate(value, head, signed, number_bytes),
+            move |value| operate(value, head, signed, bytes_len),
             engine_state.ctrlc.clone(),
         )
     }
@@ -83,20 +68,20 @@ impl Command for SubCommand {
     fn examples(&self) -> Vec<Example> {
         vec![
             Example {
-                description: "Apply the logical negation to a list of numbers",
+                description: "Apply logical negation to a list of numbers",
                 example: "[4 3 2] | bits not",
                 result: Some(Value::List {
                     vals: vec![
-                        Value::test_int(251),
+                        Value::test_int(140737488355323),
-                        Value::test_int(252),
+                        Value::test_int(140737488355324),
-                        Value::test_int(253),
+                        Value::test_int(140737488355325),
                     ],
                     span: Span::test_data(),
                 }),
             },
             Example {
                 description:
-                    "Apply the logical negation to a list of numbers, treat input as 2 bytes number",
+                    "Apply logical negation to a list of numbers, treat input as 2 bytes number",
                 example: "[4 3 2] | bits not -n 2",
                 result: Some(Value::List {
                     vals: vec![
@@ -109,7 +94,7 @@ impl Command for SubCommand {
             },
             Example {
                 description:
-                    "Apply the logical negation to a list of numbers, treat input as signed number",
+                    "Apply logical negation to a list of numbers, treat input as signed number",
                 example: "[4 3 2] | bits not -s",
                 result: Some(Value::List {
                     vals: vec![
@@ -147,6 +132,8 @@ fn operate(value: Value, head: Span, signed: bool, number_size: NumberBytes) ->
                             !val & 0x7F_FF_FF_FF_FF_FF
                         }
                     }
+                    // This case shouldn't happen here, as it's handled before
+                    Invalid => 0,
                 };
                 Value::Int { val: out_val, span }
             }

crates/nu-command/src/bits/shift_left.rs

@@ -0,0 +1,188 @@
+use super::{get_input_num_type, get_number_bytes, InputNumType, NumberBytes};
+use nu_engine::CallExt;
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{
+    Category, Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, Value,
+};
+use num_traits::CheckedShl;
+use std::fmt::Display;
+
+#[derive(Clone)]
+pub struct SubCommand;
+
+impl Command for SubCommand {
+    fn name(&self) -> &str {
+        "bits shl"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("bits shl")
+            .required("bits", SyntaxShape::Int, "number of bits to shift left")
+            .switch(
+                "signed",
+                "always treat input number as a signed number",
+                Some('s'),
+            )
+            .named(
+                "number-bytes",
+                SyntaxShape::String,
+                "the word size in number of bytes, it can be 1, 2, 4, 8, auto, default value `8`",
+                Some('n'),
+            )
+            .category(Category::Bits)
+    }
+
+    fn usage(&self) -> &str {
+        "Bitwise shift left for integers"
+    }
+
+    fn search_terms(&self) -> Vec<&str> {
+        vec!["shift left"]
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        let head = call.head;
+        let bits: usize = call.req(engine_state, stack, 0)?;
+        let signed = call.has_flag("signed");
+        let number_bytes: Option<Spanned<String>> =
+            call.get_flag(engine_state, stack, "number-bytes")?;
+        let bytes_len = get_number_bytes(&number_bytes);
+        if let NumberBytes::Invalid = bytes_len {
+            if let Some(val) = number_bytes {
+                return Err(ShellError::UnsupportedInput(
+                    "Only 1, 2, 4, 8, or 'auto' bytes are supported as word sizes".to_string(),
+                    val.span,
+                ));
+            }
+        }
+
+        input.map(
+            move |value| operate(value, bits, head, signed, bytes_len),
+            engine_state.ctrlc.clone(),
+        )
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Shift left a number by 7 bits",
+                example: "2 | bits shl 7",
+                result: Some(Value::Int {
+                    val: 256,
+                    span: Span::test_data(),
+                }),
+            },
+            Example {
+                description: "Shift left a number with 1 byte by 7 bits",
+                example: "2 | bits shl 7 -n 1",
+                result: Some(Value::Int {
+                    val: 0,
+                    span: Span::test_data(),
+                }),
+            },
+            Example {
+                description: "Shift left a signed number by 1 bit",
+                example: "0x7F | bits shl 1 -s",
+                result: Some(Value::Int {
+                    val: 254,
+                    span: Span::test_data(),
+                }),
+            },
+            Example {
+                description: "Shift left a list of numbers",
+                example: "[5 3 2] | bits shl 2",
+                result: Some(Value::List {
+                    vals: vec![Value::test_int(20), Value::test_int(12), Value::test_int(8)],
+                    span: Span::test_data(),
+                }),
+            },
+        ]
+    }
+}
+
+fn get_shift_left<T: CheckedShl + Display + Copy>(val: T, bits: u32, span: Span) -> Value
+where
+    i64: std::convert::TryFrom<T>,
+{
+    match val.checked_shl(bits) {
+        Some(val) => {
+            let shift_result = i64::try_from(val);
+            match shift_result {
+                Ok(val) => Value::Int { val, span },
+                Err(_) => Value::Error {
+                    error: ShellError::GenericError(
+                        "Shift left result beyond the range of 64 bit signed number".to_string(),
+                        format!(
+                            "{} of the specified number of bytes shift left {} bits exceed limit",
+                            val, bits
+                        ),
+                        Some(span),
+                        None,
+                        Vec::new(),
+                    ),
+                },
+            }
+        }
+        None => Value::Error {
+            error: ShellError::GenericError(
+                "Shift left failed".to_string(),
+                format!(
+                    "{} shift left {} bits failed, you may shift too many bits",
+                    val, bits
+                ),
+                Some(span),
+                None,
+                Vec::new(),
+            ),
+        },
+    }
+}
+
+fn operate(value: Value, bits: usize, head: Span, signed: bool, number_size: NumberBytes) -> Value {
+    match value {
+        Value::Int { val, span } => {
+            use InputNumType::*;
+            // let bits = (((bits % 64) + 64) % 64) as u32;
+            let bits = bits as u32;
+            let input_type = get_input_num_type(val, signed, number_size);
+            match input_type {
+                One => get_shift_left(val as u8, bits, span),
+                Two => get_shift_left(val as u16, bits, span),
+                Four => get_shift_left(val as u32, bits, span),
+                Eight => get_shift_left(val as u64, bits, span),
+                SignedOne => get_shift_left(val as i8, bits, span),
+                SignedTwo => get_shift_left(val as i16, bits, span),
+                SignedFour => get_shift_left(val as i32, bits, span),
+                SignedEight => get_shift_left(val as i64, bits, span),
+            }
+        }
+        other => Value::Error {
+            error: ShellError::UnsupportedInput(
+                format!(
+                    "Only integer values are supported, input type: {:?}",
+                    other.get_type()
+                ),
+                other.span().unwrap_or(head),
+            ),
+        },
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(SubCommand {})
+    }
+}

crates/nu-command/src/bits/shift_right.rs

@@ -0,0 +1,172 @@
+use super::{get_input_num_type, get_number_bytes, InputNumType, NumberBytes};
+use nu_engine::CallExt;
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{
+    Category, Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, Value,
+};
+use num_traits::CheckedShr;
+use std::fmt::Display;
+
+#[derive(Clone)]
+pub struct SubCommand;
+
+impl Command for SubCommand {
+    fn name(&self) -> &str {
+        "bits shr"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("bits shr")
+            .required("bits", SyntaxShape::Int, "number of bits to shift right")
+            .switch(
+                "signed",
+                "always treat input number as a signed number",
+                Some('s'),
+            )
+            .named(
+                "number-bytes",
+                SyntaxShape::String,
+                "the word size in number of bytes, it can be 1, 2, 4, 8, auto, default value `8`",
+                Some('n'),
+            )
+            .category(Category::Bits)
+    }
+
+    fn usage(&self) -> &str {
+        "Bitwise shift right for integers"
+    }
+
+    fn search_terms(&self) -> Vec<&str> {
+        vec!["shift right"]
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        let head = call.head;
+        let bits: usize = call.req(engine_state, stack, 0)?;
+        let signed = call.has_flag("signed");
+        let number_bytes: Option<Spanned<String>> =
+            call.get_flag(engine_state, stack, "number-bytes")?;
+        let bytes_len = get_number_bytes(&number_bytes);
+        if let NumberBytes::Invalid = bytes_len {
+            if let Some(val) = number_bytes {
+                return Err(ShellError::UnsupportedInput(
+                    "Only 1, 2, 4, 8, or 'auto' bytes are supported as word sizes".to_string(),
+                    val.span,
+                ));
+            }
+        }
+
+        input.map(
+            move |value| operate(value, bits, head, signed, bytes_len),
+            engine_state.ctrlc.clone(),
+        )
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Shift right a number with 2 bits",
+                example: "8 | bits shr 2",
+                result: Some(Value::Int {
+                    val: 2,
+                    span: Span::test_data(),
+                }),
+            },
+            Example {
+                description: "Shift right a list of numbers",
+                example: "[15 35 2] | bits shr 2",
+                result: Some(Value::List {
+                    vals: vec![Value::test_int(3), Value::test_int(8), Value::test_int(0)],
+                    span: Span::test_data(),
+                }),
+            },
+        ]
+    }
+}
+
+fn get_shift_right<T: CheckedShr + Display + Copy>(val: T, bits: u32, span: Span) -> Value
+where
+    i64: std::convert::TryFrom<T>,
+{
+    match val.checked_shr(bits) {
+        Some(val) => {
+            let shift_result = i64::try_from(val);
+            match shift_result {
+                Ok(val) => Value::Int { val, span },
+                Err(_) => Value::Error {
+                    error: ShellError::GenericError(
+                        "Shift right result beyond the range of 64 bit signed number".to_string(),
+                        format!(
+                            "{} of the specified number of bytes shift right {} bits exceed limit",
+                            val, bits
+                        ),
+                        Some(span),
+                        None,
+                        Vec::new(),
+                    ),
+                },
+            }
+        }
+        None => Value::Error {
+            error: ShellError::GenericError(
+                "Shift right failed".to_string(),
+                format!(
+                    "{} shift right {} bits failed, you may shift too many bits",
+                    val, bits
+                ),
+                Some(span),
+                None,
+                Vec::new(),
+            ),
+        },
+    }
+}
+
+fn operate(value: Value, bits: usize, head: Span, signed: bool, number_size: NumberBytes) -> Value {
+    match value {
+        Value::Int { val, span } => {
+            use InputNumType::*;
+            // let bits = (((bits % 64) + 64) % 64) as u32;
+            let bits = bits as u32;
+            let input_type = get_input_num_type(val, signed, number_size);
+            match input_type {
+                One => get_shift_right(val as u8, bits, span),
+                Two => get_shift_right(val as u16, bits, span),
+                Four => get_shift_right(val as u32, bits, span),
+                Eight => get_shift_right(val as u64, bits, span),
+                SignedOne => get_shift_right(val as i8, bits, span),
+                SignedTwo => get_shift_right(val as i16, bits, span),
+                SignedFour => get_shift_right(val as i32, bits, span),
+                SignedEight => get_shift_right(val as i64, bits, span),
+            }
+        }
+        other => Value::Error {
+            error: ShellError::UnsupportedInput(
+                format!(
+                    "Only integer values are supported, input type: {:?}",
+                    other.get_type()
+                ),
+                other.span().unwrap_or(head),
+            ),
+        },
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(SubCommand {})
+    }
+}

crates/nu-command/src/bytes/add.rs

@@ -47,7 +47,7 @@ impl Command for BytesAdd {
     }
 
     fn usage(&self) -> &str {
-        "add specified bytes to the input"
+        "Add specified bytes to the input"
     }
 
     fn search_terms(&self) -> Vec<&str> {

crates/nu-command/src/bytes/remove.rs

@@ -41,7 +41,7 @@ impl Command for BytesRemove {
     }
 
     fn usage(&self) -> &str {
-        "remove bytes"
+        "Remove bytes"
     }
 
     fn search_terms(&self) -> Vec<&str> {

crates/nu-command/src/default_context.rs

@@ -214,6 +214,8 @@ pub fn create_default_context() -> EngineState {
             BitsNot,
             BitsOr,
             BitsXor,
+            BitsShiftLeft,
+            BitsShiftRight,
         }
 
         // Bytes

crates/nu-command/src/strings/encode_decode/decode_base64.rs

@@ -35,7 +35,7 @@ impl Command for DecodeBase64 {
     }
 
     fn usage(&self) -> &str {
-        "base64 decode a value"
+        "Base64 decode a value"
     }
 
     fn extra_usage(&self) -> &str {

crates/nu-command/src/strings/encode_decode/encode_base64.rs

@@ -31,7 +31,7 @@ impl Command for EncodeBase64 {
     }
 
     fn usage(&self) -> &str {
-        "base64 encode a value"
+        "Base64 encode a value"
     }
 
     fn examples(&self) -> Vec<Example> {

src/tests/test_bits.rs

@@ -44,3 +44,33 @@ fn bits_xor_negative() -> TestResult {
 fn bits_xor_list() -> TestResult {
     run_test("[1 2 3 8 9 10] | bits xor 2 | str collect '.'", "3.0.1.10.11.8")
 }
+
+#[test]
+fn bits_shift_left() -> TestResult {
+    run_test("2 | bits shl 3", "16")
+}
+
+#[test]
+fn bits_shift_left_negative() -> TestResult {
+    run_test("-3 | bits shl 5", "-96")
+}
+
+#[test]
+fn bits_shift_left_list() -> TestResult {
+    run_test("[1 2 7 32 9 10] | bits shl 3 | str collect '.'", "8.16.56.256.72.80")
+}
+
+#[test]
+fn bits_shift_right() -> TestResult {
+    run_test("8 | bits shr 2", "2")
+}
+
+#[test]
+fn bits_shift_right_negative() -> TestResult {
+    run_test("-32 | bits shr 2", "-8")
+}
+
+#[test]
+fn bits_shift_right_list() -> TestResult {
+    run_test("[12 98 7 64 900 10] | bits shr 3 | str collect '.'", "1.12.0.8.112.1")
+}