From 9a4196fb2a274dbc82d61c277bb30b88c5864ee7 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Sat, 16 Dec 2023 19:40:14 -0500
Subject: [PATCH] Fix list add-assign type checking

---
 src/abstract_tree/assignment.rs | 26 ++++++++++----------------
 src/error.rs                    |  8 ++++++++
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/src/abstract_tree/assignment.rs b/src/abstract_tree/assignment.rs
index 6c6d0b8..15003bc 100644
--- a/src/abstract_tree/assignment.rs
+++ b/src/abstract_tree/assignment.rs
@@ -26,6 +26,7 @@ impl AbstractTree for Assignment {
 
         let identifier_node = node.child(0).unwrap();
         let identifier = Identifier::from_syntax_node(source, identifier_node, context)?;
+        let identifier_type = identifier.expected_type(context)?;
 
         let type_node = node.child(1);
         let type_definition = if let Some(type_node) = type_node {
@@ -74,12 +75,7 @@ impl AbstractTree for Assignment {
                         .map_err(|error| error.at_node(statement_node, source))?;
                 }
                 AssignmentOperator::PlusEqual => {
-                    let identifier_type = identifier.expected_type(context)?;
-
                     if let Type::List(item_type) = type_definition.inner() {
-                        item_type
-                            .check(&identifier_type)
-                            .map_err(|error| error.at_node(identifier_node, source))?;
                         item_type
                             .check(&statement_type)
                             .map_err(|error| error.at_node(statement_node, source))?;
@@ -88,14 +84,18 @@ impl AbstractTree for Assignment {
                             .inner()
                             .check(&identifier_type)
                             .map_err(|error| error.at_node(identifier_node, source))?;
-                        type_definition
-                            .inner()
-                            .check(&statement_type)
-                            .map_err(|error| error.at_node(statement_node, source))?;
                     }
                 }
                 AssignmentOperator::MinusEqual => todo!(),
             }
+        } else {
+            if let Type::List(item_type) = identifier_type {
+                println!("{item_type} {statement_type}");
+
+                item_type
+                    .check(&statement_type)
+                    .map_err(|error| error.at_node(statement_node, source))?;
+            }
         }
 
         Ok(Assignment {
@@ -185,12 +185,6 @@ mod tests {
             ",
         );
 
-        assert_eq!(
-            Err(Error::TypeCheck {
-                expected: Type::String,
-                actual: Type::Integer
-            }),
-            result
-        )
+        assert!(result.unwrap_err().is_type_check_error())
     }
 }
diff --git a/src/error.rs b/src/error.rs
index 89fedff..7341506 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -213,6 +213,14 @@ impl Error {
             })
         }
     }
+
+    pub fn is_type_check_error(&self) -> bool {
+        match self {
+            Error::TypeCheck { .. } => true,
+            Error::WithContext { error, .. } => error.is_type_check_error(),
+            _ => false,
+        }
+    }
 }
 
 impl<T> From<PoisonError<T>> for Error {