Update #4202: Add shift operator bshl and bshr for integers (#5928)

* Update #4202: Add shift operator bshl and bshr for integers

* Add more tests

crates/nu-command/src/database/values/dsl/expression.rs

@@ -118,6 +118,8 @@ impl CustomValue for ExprDb {
             Operator::In
             | Operator::NotIn
             | Operator::Pow
+            | Operator::ShiftLeft
+            | Operator::ShiftRight
             | Operator::StartsWith
             | Operator::EndsWith => Err(ShellError::UnsupportedOperator(operator, op)),
         }?;

crates/nu-engine/src/eval.rs

@@ -438,6 +438,14 @@ pub fn eval_expression(
                     let rhs = eval_expression(engine_state, stack, rhs)?;
                     lhs.ends_with(op_span, &rhs, expr.span)
                 }
+                Operator::ShiftRight => {
+                    let rhs = eval_expression(engine_state, stack, rhs)?;
+                    lhs.bshr(op_span, &rhs, expr.span)
+                }
+                Operator::ShiftLeft => {
+                    let rhs = eval_expression(engine_state, stack, rhs)?;
+                    lhs.bshl(op_span, &rhs, expr.span)
+                }
             }
         }
         Expr::Subexpression(block_id) => {

crates/nu-parser/src/parser.rs

@@ -4106,6 +4106,8 @@ pub fn parse_operator(
         b"in" => Operator::In,
         b"not-in" => Operator::NotIn,
         b"mod" => Operator::Modulo,
+        b"bshl" => Operator::ShiftLeft,
+        b"bshr" => Operator::ShiftRight,
         b"starts-with" => Operator::StartsWith,
         b"ends-with" => Operator::EndsWith,
         b"&&" | b"and" => Operator::And,

crates/nu-parser/src/type_check.rs

@@ -492,6 +492,25 @@ pub fn math_result_type(
                     )
                 }
             },
+            Operator::ShiftLeft | Operator::ShiftRight => match (&lhs.ty, &rhs.ty) {
+                (Type::Int, Type::Int) => (Type::Int, None),
+
+                (Type::Any, _) => (Type::Any, None),
+                (_, Type::Any) => (Type::Any, None),
+                _ => {
+                    *op = Expression::garbage(op.span);
+                    (
+                        Type::Any,
+                        Some(ParseError::UnsupportedOperation(
+                            op.span,
+                            lhs.span,
+                            lhs.ty.clone(),
+                            rhs.span,
+                            rhs.ty.clone(),
+                        )),
+                    )
+                }
+            },
         },
         _ => {
             *op = Expression::garbage(op.span);

crates/nu-protocol/src/ast/expression.rs

@@ -35,6 +35,7 @@ impl Expression {
                     | Operator::Modulo
                     | Operator::FloorDivision => 95,
                     Operator::Plus | Operator::Minus => 90,
+                    Operator::ShiftLeft | Operator::ShiftRight => 85,
                     Operator::NotRegexMatch
                     | Operator::RegexMatch
                     | Operator::StartsWith

crates/nu-protocol/src/ast/operator.rs

@@ -26,6 +26,8 @@ pub enum Operator {
     Pow,
     StartsWith,
     EndsWith,
+    ShiftLeft,
+    ShiftRight,
 }
 
 impl Display for Operator {
@@ -48,6 +50,8 @@ impl Display for Operator {
             Operator::And => write!(f, "&&"),
             Operator::Or => write!(f, "||"),
             Operator::Pow => write!(f, "**"),
+            Operator::ShiftLeft => write!(f, "bshl"),
+            Operator::ShiftRight => write!(f, "bshr"),
             Operator::LessThanOrEqual => write!(f, "<="),
             Operator::GreaterThanOrEqual => write!(f, ">="),
             Operator::StartsWith => write!(f, "starts-with"),

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

@@ -2240,6 +2240,44 @@ impl Value {
         }
     }
 
+    pub fn bshl(&self, op: Span, rhs: &Value, span: Span) -> Result<Value, ShellError> {
+        match (self, rhs) {
+            (Value::Int { val: lhs, .. }, Value::Int { val: rhs, .. }) => Ok(Value::Int {
+                span,
+                val: *lhs << rhs,
+            }),
+            (Value::CustomValue { val: lhs, span }, rhs) => {
+                lhs.operation(*span, Operator::ShiftLeft, op, rhs)
+            }
+            _ => Err(ShellError::OperatorMismatch {
+                op_span: op,
+                lhs_ty: self.get_type(),
+                lhs_span: self.span()?,
+                rhs_ty: rhs.get_type(),
+                rhs_span: rhs.span()?,
+            }),
+        }
+    }
+
+    pub fn bshr(&self, op: Span, rhs: &Value, span: Span) -> Result<Value, ShellError> {
+        match (self, rhs) {
+            (Value::Int { val: lhs, .. }, Value::Int { val: rhs, .. }) => Ok(Value::Int {
+                span,
+                val: *lhs >> rhs,
+            }),
+            (Value::CustomValue { val: lhs, span }, rhs) => {
+                lhs.operation(*span, Operator::ShiftRight, op, rhs)
+            }
+            _ => Err(ShellError::OperatorMismatch {
+                op_span: op,
+                lhs_ty: self.get_type(),
+                lhs_span: self.span()?,
+                rhs_ty: rhs.get_type(),
+                rhs_span: rhs.span()?,
+            }),
+        }
+    }
+
     pub fn modulo(&self, op: Span, rhs: &Value, span: Span) -> Result<Value, ShellError> {
         match (self, rhs) {
             (Value::Int { val: lhs, .. }, Value::Int { val: rhs, .. }) => {

src/tests/test_math.rs

@@ -25,6 +25,26 @@ fn modulo2() -> TestResult {
     run_test("5.25 mod 2", "1.25")
 }
 
+#[test]
+fn bshr() -> TestResult {
+    run_test("16 bshr 1", "8")
+}
+
+#[test]
+fn bshl() -> TestResult {
+    run_test("5 bshl 1", "10")
+}
+
+#[test]
+fn bshl_add() -> TestResult {
+    run_test("2 bshl 1 + 2", "16")
+}
+
+#[test]
+fn sub_bshr() -> TestResult {
+    run_test("10 - 2 bshr 2", "2")
+}
+
 #[test]
 fn and() -> TestResult {
     run_test("true && false", "false")