From d8850b2d3c4315fc6bb7710e3663b1c975b22bd0 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Mon, 8 Jan 2024 10:25:01 -0500
Subject: [PATCH] Fix assignment type check bug

---
 src/abstract_tree/assignment.rs | 30 ++++++++++++++++++++++--------
 src/abstract_tree/value_node.rs | 10 +---------
 2 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/src/abstract_tree/assignment.rs b/src/abstract_tree/assignment.rs
index 27978b0..18a740b 100644
--- a/src/abstract_tree/assignment.rs
+++ b/src/abstract_tree/assignment.rs
@@ -42,13 +42,13 @@ impl AbstractTree for Assignment {
         let variable_key = identifier.inner().clone();
         let variable_type = if let Some(definition) = &type_definition {
             definition.inner().clone()
-        } else if let Some((_, r#type)) = context.variables()?.get(identifier.inner()) {
-            r#type.clone()
         } else {
             statement_type
         };
 
-        context.set(variable_key, Value::none(), Some(variable_type))?;
+        if let AssignmentOperator::Equal = operator {
+            context.set(variable_key, Value::none(), Some(variable_type))?;
+        }
 
         Ok(Assignment {
             identifier,
@@ -60,19 +60,29 @@ impl AbstractTree for Assignment {
     }
 
     fn check_type(&self, source: &str, context: &Map) -> Result<()> {
-        let statement_type = self.statement.expected_type(context)?;
+        let variables = context.variables()?;
+        let established_type = variables
+            .get(self.identifier.inner())
+            .map(|(_value, _type)| _type);
+        let actual_type = self.statement.expected_type(context)?;
 
         if let Some(type_definition) = &self.type_definition {
             match self.operator {
                 AssignmentOperator::Equal => {
+                    if let Some(r#type) = established_type {
+                        r#type.check(&actual_type).map_err(|error| {
+                            error.at_source_position(source, self.syntax_position)
+                        })?;
+                    }
+
                     type_definition
                         .inner()
-                        .check(&statement_type)
+                        .check(&actual_type)
                         .map_err(|error| error.at_source_position(source, self.syntax_position))?;
                 }
                 AssignmentOperator::PlusEqual => {
                     if let Type::List(item_type) = type_definition.inner() {
-                        item_type.check(&statement_type)?;
+                        item_type.check(&actual_type)?;
                     } else {
                         type_definition
                             .inner()
@@ -86,10 +96,14 @@ impl AbstractTree for Assignment {
             }
         } else {
             match self.operator {
-                AssignmentOperator::Equal => {}
+                AssignmentOperator::Equal => {
+                    if let Some(r#type) = established_type {
+                        r#type.check(&actual_type)?;
+                    }
+                }
                 AssignmentOperator::PlusEqual => {
                     if let Type::List(item_type) = self.identifier.expected_type(context)? {
-                        item_type.check(&statement_type).map_err(|error| {
+                        item_type.check(&actual_type).map_err(|error| {
                             error.at_source_position(source, self.syntax_position)
                         })?;
                     }
diff --git a/src/abstract_tree/value_node.rs b/src/abstract_tree/value_node.rs
index e474a99..fa9c0ef 100644
--- a/src/abstract_tree/value_node.rs
+++ b/src/abstract_tree/value_node.rs
@@ -190,20 +190,12 @@ impl AbstractTree for ValueNode {
 
     fn check_type(&self, _source: &str, _context: &Map) -> Result<()> {
         match self {
-            ValueNode::Boolean(_) => todo!(),
-            ValueNode::Float(_) => todo!(),
             ValueNode::Function(function) => {
                 if let Function::ContextDefined(function_node) = function {
                     function_node.check_type(_source, _context)?;
                 }
             }
-            ValueNode::Integer(_) => todo!(),
-            ValueNode::String(_) => todo!(),
-            ValueNode::List(_) => todo!(),
-            ValueNode::Option(_) => todo!(),
-            ValueNode::Map(_) => todo!(),
-            ValueNode::BuiltInValue(_) => todo!(),
-            ValueNode::Structure(_) => todo!(),
+            _ => {}
         }
 
         Ok(())