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 { + 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 { + 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 {