From aada1c72d6afadb526ffbfc200e00623eeedf341 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Sat, 9 Dec 2023 18:50:17 -0500
Subject: [PATCH] Fix list type checking

---
 src/abstract_tree/assignment.rs | 46 +++++++++++++++++++--------------
 src/value/map.rs                |  9 +++++--
 2 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/src/abstract_tree/assignment.rs b/src/abstract_tree/assignment.rs
index e0e7668..f9e0c0f 100644
--- a/src/abstract_tree/assignment.rs
+++ b/src/abstract_tree/assignment.rs
@@ -58,18 +58,6 @@ impl AbstractTree for Assignment {
         let statement_node = node.child(child_count - 1).unwrap();
         let statement = Statement::from_syntax_node(source, statement_node, context)?;
 
-        if let Some((_previous_value, previous_type)) = context.variables()?.get(identifier.inner())
-        {
-            let type_check = previous_type.check(&statement.expected_type(context)?);
-
-            if let Err(error) = type_check {
-                return Err(error.with_context(
-                    statement_node.start_position(),
-                    source[statement_node.byte_range()].to_string(),
-                ));
-            }
-        }
-
         if let Some(type_definition) = &type_definition {
             let statement_type = statement.expected_type(context)?;
 
@@ -81,8 +69,6 @@ impl AbstractTree for Assignment {
                     let identifier_type = identifier.expected_type(context)?;
 
                     if let Type::List(item_type) = type_definition.inner() {
-                        println!("{item_type}");
-
                         item_type.check(&identifier_type)?;
                         item_type.check(&statement_type)?;
                     } else {
@@ -92,6 +78,22 @@ impl AbstractTree for Assignment {
                 }
                 AssignmentOperator::MinusEqual => todo!(),
             }
+        } else if let Some((_previous_value, previous_type)) =
+            context.variables()?.get(identifier.inner())
+        {
+            let statement_type = statement.expected_type(context)?;
+            let type_check = if let Type::List(item_type) = previous_type {
+                item_type.check(&statement_type)
+            } else {
+                previous_type.check(&statement_type)
+            };
+
+            if let Err(error) = type_check {
+                return Err(error.with_context(
+                    statement_node.start_position(),
+                    source[statement_node.byte_range()].to_string(),
+                ));
+            }
         }
 
         Ok(Assignment {
@@ -126,7 +128,11 @@ impl AbstractTree for Assignment {
             AssignmentOperator::Equal => value,
         };
 
-        context.set(key.clone(), new_value)?;
+        if let Some(type_defintion) = &self.type_definition {
+            context.set(key.clone(), new_value, Some(type_defintion.inner().clone()))?;
+        } else {
+            context.set(key.clone(), new_value, None)?;
+        }
 
         Ok(Value::Empty)
     }
@@ -175,11 +181,13 @@ mod tests {
             x <[str]> = []
             x += 1
             ",
-        );
+        )
+        .unwrap_err();
 
-        if let Err(Error::TypeCheck { .. }) = test {
-        } else {
-            panic!()
+        match test {
+            Error::WithContext { .. } => {}
+            Error::TypeCheck { .. } => {}
+            _ => panic!(),
         }
     }
 }
diff --git a/src/value/map.rs b/src/value/map.rs
index 8dd8d68..1efe1d3 100644
--- a/src/value/map.rs
+++ b/src/value/map.rs
@@ -41,8 +41,13 @@ impl Map {
         Ok(self.variables.read()?)
     }
 
-    pub fn set(&self, key: String, value: Value) -> Result<Option<(Value, Type)>> {
-        let value_type = value.r#type();
+    pub fn set(
+        &self,
+        key: String,
+        value: Value,
+        r#type: Option<Type>,
+    ) -> Result<Option<(Value, Type)>> {
+        let value_type = r#type.unwrap_or(value.r#type());
         let previous = self
             .variables
             .write()?