From ec9f17070c5e4fc263c98d70409f8aad6e5a31b1 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Fri, 8 Mar 2024 13:26:55 -0500
Subject: [PATCH] Implement add-assign and subtract-assign

---
 src/abstract_tree/assignment.rs | 74 ++++++++++++++++++++++++++++++++-
 src/abstract_tree/math.rs       | 20 +--------
 src/value.rs                    | 20 +++++++++
 3 files changed, 95 insertions(+), 19 deletions(-)

diff --git a/src/abstract_tree/assignment.rs b/src/abstract_tree/assignment.rs
index 40af64c..6c8d66d 100644
--- a/src/abstract_tree/assignment.rs
+++ b/src/abstract_tree/assignment.rs
@@ -62,7 +62,33 @@ impl<'src> AbstractTree for Assignment<'src> {
             r#break => return Ok(r#break),
         };
 
-        context.set_value(self.identifier, value)?;
+        match self.operator {
+            AssignmentOperator::Assign => {
+                context.set_value(self.identifier, value)?;
+            }
+            AssignmentOperator::AddAssign => {
+                if let Some(previous_value) = context.get_value(&self.identifier)? {
+                    let new_value = previous_value.add(&value)?;
+
+                    context.set_value(self.identifier, new_value)?;
+                } else {
+                    return Err(RuntimeError::ValidationFailure(
+                        ValidationError::VariableNotFound(self.identifier),
+                    ));
+                }
+            }
+            AssignmentOperator::SubAssign => {
+                if let Some(previous_value) = context.get_value(&self.identifier)? {
+                    let new_value = previous_value.subtract(&value)?;
+
+                    context.set_value(self.identifier, new_value)?;
+                } else {
+                    return Err(RuntimeError::ValidationFailure(
+                        ValidationError::VariableNotFound(self.identifier),
+                    ));
+                }
+            }
+        }
 
         Ok(Action::None)
     }
@@ -97,6 +123,52 @@ mod tests {
         )
     }
 
+    #[test]
+    fn add_assign_value() {
+        let context = Context::new();
+
+        context
+            .set_value(Identifier::new("foobar"), Value::integer(1))
+            .unwrap();
+
+        Assignment::new(
+            Identifier::new("foobar"),
+            None,
+            AssignmentOperator::AddAssign,
+            Statement::Expression(Expression::Value(ValueNode::Integer(41))),
+        )
+        .run(&context)
+        .unwrap();
+
+        assert_eq!(
+            context.get_value(&Identifier::new("foobar")),
+            Ok(Some(Value::integer(42)))
+        )
+    }
+
+    #[test]
+    fn subtract_assign_value() {
+        let context = Context::new();
+
+        context
+            .set_value(Identifier::new("foobar"), Value::integer(43))
+            .unwrap();
+
+        Assignment::new(
+            Identifier::new("foobar"),
+            None,
+            AssignmentOperator::SubAssign,
+            Statement::Expression(Expression::Value(ValueNode::Integer(1))),
+        )
+        .run(&context)
+        .unwrap();
+
+        assert_eq!(
+            context.get_value(&Identifier::new("foobar")),
+            Ok(Some(Value::integer(42)))
+        )
+    }
+
     #[test]
     fn type_check() {
         let validation = Assignment::new(
diff --git a/src/abstract_tree/math.rs b/src/abstract_tree/math.rs
index a8e7b0f..565a8db 100644
--- a/src/abstract_tree/math.rs
+++ b/src/abstract_tree/math.rs
@@ -54,29 +54,13 @@ impl<'src> AbstractTree for Math<'src> {
                 let left_value = left.run(_context)?.as_return_value()?;
                 let right_value = right.run(_context)?.as_return_value()?;
 
-                if let (ValueInner::Integer(left), ValueInner::Integer(right)) =
-                    (left_value.inner().as_ref(), right_value.inner().as_ref())
-                {
-                    Value::integer(left + right)
-                } else {
-                    return Err(RuntimeError::ValidationFailure(
-                        ValidationError::ExpectedIntegerOrFloat,
-                    ));
-                }
+                left_value.add(&right_value)?
             }
             Math::Subtract(left, right) => {
                 let left_value = left.run(_context)?.as_return_value()?;
                 let right_value = right.run(_context)?.as_return_value()?;
 
-                if let (ValueInner::Integer(left), ValueInner::Integer(right)) =
-                    (left_value.inner().as_ref(), right_value.inner().as_ref())
-                {
-                    Value::integer(left - right)
-                } else {
-                    return Err(RuntimeError::ValidationFailure(
-                        ValidationError::ExpectedIntegerOrFloat,
-                    ));
-                }
+                left_value.subtract(&right_value)?
             }
             Math::Multiply(left, right) => {
                 let left_value = left.run(_context)?.as_return_value()?;
diff --git a/src/value.rs b/src/value.rs
index b65be9c..c0385e3 100644
--- a/src/value.rs
+++ b/src/value.rs
@@ -128,6 +128,26 @@ impl Value {
     pub fn is_none(&self) -> bool {
         self == get_none()
     }
+
+    pub fn add(&self, other: &Self) -> Result<Value, ValidationError> {
+        if let (ValueInner::Integer(left), ValueInner::Integer(right)) =
+            (self.inner().as_ref(), other.inner().as_ref())
+        {
+            Ok(Value::integer(left + right))
+        } else {
+            Err(ValidationError::ExpectedIntegerOrFloat)
+        }
+    }
+
+    pub fn subtract(&self, other: &Self) -> Result<Value, ValidationError> {
+        if let (ValueInner::Integer(left), ValueInner::Integer(right)) =
+            (self.inner().as_ref(), other.inner().as_ref())
+        {
+            Ok(Value::integer(left - right))
+        } else {
+            Err(ValidationError::ExpectedIntegerOrFloat)
+        }
+    }
 }
 
 impl Display for Value {