From c1105e945e40f3cb24f429fae8d215f024f3a48b Mon Sep 17 00:00:00 2001
From: JT <547158+jntrnr@users.noreply.github.com>
Date: Sat, 12 Nov 2022 07:50:43 +1300
Subject: [PATCH] Add additional assignment operators (#7102)

---
 crates/nu-command/tests/commands/mut_.rs | 48 ++++++++++++++++++++++++
 crates/nu-engine/src/eval.rs             | 22 ++++++++++-
 crates/nu-parser/src/parser.rs           |  4 ++
 crates/nu-parser/src/type_check.rs       |  4 +-
 crates/nu-protocol/src/ast/expression.rs |  2 +-
 crates/nu-protocol/src/ast/operator.rs   |  8 ++++
 6 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/crates/nu-command/tests/commands/mut_.rs b/crates/nu-command/tests/commands/mut_.rs
index b6a72e737..0034f6d99 100644
--- a/crates/nu-command/tests/commands/mut_.rs
+++ b/crates/nu-command/tests/commands/mut_.rs
@@ -47,3 +47,51 @@ fn mut_a_field() {
 
     assert_eq!(actual.out, "456");
 }
+
+#[test]
+fn mut_add_assign() {
+    let actual = nu!(
+        cwd: ".", pipeline(
+        r#"
+        mut y = 3; $y += 2; $y
+        "#
+    ));
+
+    assert_eq!(actual.out, "5");
+}
+
+#[test]
+fn mut_minus_assign() {
+    let actual = nu!(
+        cwd: ".", pipeline(
+        r#"
+        mut y = 3; $y -= 2; $y
+        "#
+    ));
+
+    assert_eq!(actual.out, "1");
+}
+
+#[test]
+fn mut_multiply_assign() {
+    let actual = nu!(
+        cwd: ".", pipeline(
+        r#"
+        mut y = 3; $y *= 2; $y
+        "#
+    ));
+
+    assert_eq!(actual.out, "6");
+}
+
+#[test]
+fn mut_divide_assign() {
+    let actual = nu!(
+        cwd: ".", pipeline(
+        r#"
+        mut y = 8; $y /= 2; $y
+        "#
+    ));
+
+    assert_eq!(actual.out, "4");
+}
diff --git a/crates/nu-engine/src/eval.rs b/crates/nu-engine/src/eval.rs
index 97bfd6b5e..43ec854e6 100644
--- a/crates/nu-engine/src/eval.rs
+++ b/crates/nu-engine/src/eval.rs
@@ -434,9 +434,29 @@ pub fn eval_expression(
                         Bits::ShiftRight => lhs.bit_shr(op_span, &rhs, expr.span),
                     }
                 }
-                Operator::Assignment(Assignment::Assign) => {
+                Operator::Assignment(assignment) => {
                     let rhs = eval_expression(engine_state, stack, rhs)?;
 
+                    let rhs = match assignment {
+                        Assignment::Assign => rhs,
+                        Assignment::PlusAssign => {
+                            let lhs = eval_expression(engine_state, stack, lhs)?;
+                            lhs.add(op_span, &rhs, op_span)?
+                        }
+                        Assignment::MinusAssign => {
+                            let lhs = eval_expression(engine_state, stack, lhs)?;
+                            lhs.sub(op_span, &rhs, op_span)?
+                        }
+                        Assignment::MultiplyAssign => {
+                            let lhs = eval_expression(engine_state, stack, lhs)?;
+                            lhs.mul(op_span, &rhs, op_span)?
+                        }
+                        Assignment::DivideAssign => {
+                            let lhs = eval_expression(engine_state, stack, lhs)?;
+                            lhs.div(op_span, &rhs, op_span)?
+                        }
+                    };
+
                     match &lhs.expr {
                         Expr::Var(var_id) | Expr::VarDecl(var_id) => {
                             let var_info = engine_state.get_var(*var_id);
diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs
index 7f5ef20aa..f17e10b16 100644
--- a/crates/nu-parser/src/parser.rs
+++ b/crates/nu-parser/src/parser.rs
@@ -4420,6 +4420,10 @@ pub fn parse_operator(
 
     let operator = match contents {
         b"=" => Operator::Assignment(Assignment::Assign),
+        b"+=" => Operator::Assignment(Assignment::PlusAssign),
+        b"-=" => Operator::Assignment(Assignment::MinusAssign),
+        b"*=" => Operator::Assignment(Assignment::MultiplyAssign),
+        b"/=" => Operator::Assignment(Assignment::DivideAssign),
         b"==" => Operator::Comparison(Comparison::Equal),
         b"!=" => Operator::Comparison(Comparison::NotEqual),
         b"<" => Operator::Comparison(Comparison::LessThan),
diff --git a/crates/nu-parser/src/type_check.rs b/crates/nu-parser/src/type_check.rs
index 0daf0fa93..75978db1f 100644
--- a/crates/nu-parser/src/type_check.rs
+++ b/crates/nu-parser/src/type_check.rs
@@ -1,6 +1,6 @@
 use crate::ParseError;
 use nu_protocol::{
-    ast::{Assignment, Bits, Boolean, Comparison, Expr, Expression, Math, Operator},
+    ast::{Bits, Boolean, Comparison, Expr, Expression, Math, Operator},
     engine::StateWorkingSet,
     Type,
 };
@@ -555,7 +555,7 @@ pub fn math_result_type(
                     )
                 }
             },
-            Operator::Assignment(Assignment::Assign) => match (&lhs.ty, &rhs.ty) {
+            Operator::Assignment(_) => match (&lhs.ty, &rhs.ty) {
                 (x, y) if x == y => (Type::Nothing, None),
                 (Type::Any, _) => (Type::Nothing, None),
                 (_, Type::Any) => (Type::Nothing, None),
diff --git a/crates/nu-protocol/src/ast/expression.rs b/crates/nu-protocol/src/ast/expression.rs
index 5e30d3a17..c4e6e7752 100644
--- a/crates/nu-protocol/src/ast/expression.rs
+++ b/crates/nu-protocol/src/ast/expression.rs
@@ -55,7 +55,7 @@ impl Expression {
                     Operator::Bits(Bits::BitOr) => 60,
                     Operator::Boolean(Boolean::And) => 50,
                     Operator::Boolean(Boolean::Or) => 40,
-                    Operator::Assignment(Assignment::Assign) => 10,
+                    Operator::Assignment(_) => 10,
                 }
             }
             _ => 0,
diff --git a/crates/nu-protocol/src/ast/operator.rs b/crates/nu-protocol/src/ast/operator.rs
index 644e7364f..9c0286dd0 100644
--- a/crates/nu-protocol/src/ast/operator.rs
+++ b/crates/nu-protocol/src/ast/operator.rs
@@ -49,6 +49,10 @@ pub enum Bits {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub enum Assignment {
     Assign,
+    PlusAssign,
+    MinusAssign,
+    MultiplyAssign,
+    DivideAssign,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -64,6 +68,10 @@ impl Display for Operator {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Operator::Assignment(Assignment::Assign) => write!(f, "="),
+            Operator::Assignment(Assignment::PlusAssign) => write!(f, "+="),
+            Operator::Assignment(Assignment::MinusAssign) => write!(f, "-="),
+            Operator::Assignment(Assignment::MultiplyAssign) => write!(f, "*="),
+            Operator::Assignment(Assignment::DivideAssign) => write!(f, "/="),
             Operator::Comparison(Comparison::Equal) => write!(f, "=="),
             Operator::Comparison(Comparison::NotEqual) => write!(f, "!="),
             Operator::Comparison(Comparison::LessThan) => write!(f, "<"),