From 4a42f51580c79b3624408ff7bb267b5201062af0 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Wed, 17 Jan 2024 14:45:34 -0500
Subject: [PATCH] Reimplement type setting for type check system

---
 src/abstract_tree/assignment.rs    | 11 ++++++++++-
 src/abstract_tree/function_node.rs |  4 ++++
 src/abstract_tree/type.rs          |  2 ++
 src/value/map.rs                   |  8 ++++++++
 tests/types.rs                     |  2 ++
 5 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/abstract_tree/assignment.rs b/src/abstract_tree/assignment.rs
index 7aab566..2daa5b2 100644
--- a/src/abstract_tree/assignment.rs
+++ b/src/abstract_tree/assignment.rs
@@ -37,6 +37,13 @@ impl AbstractTree for Assignment {
         let statement_node = syntax_node.child(child_count - 1).unwrap();
         let statement = Statement::from_syntax(statement_node, source, context)?;
 
+        if let AssignmentOperator::Equal = operator {
+            context.set_type(
+                identifier.inner().clone(),
+                statement.expected_type(context)?,
+            )?;
+        }
+
         Ok(Assignment {
             identifier,
             type_definition,
@@ -87,7 +94,9 @@ impl AbstractTree for Assignment {
             }
         }
 
-        self.statement.check_type(source, context)?;
+        self.statement
+            .check_type(source, context)
+            .map_err(|error| error.at_source_position(source, self.syntax_position))?;
 
         Ok(())
     }
diff --git a/src/abstract_tree/function_node.rs b/src/abstract_tree/function_node.rs
index 824a095..3d9e683 100644
--- a/src/abstract_tree/function_node.rs
+++ b/src/abstract_tree/function_node.rs
@@ -115,6 +115,10 @@ impl AbstractTree for FunctionNode {
 
         function_context.clone_complex_values_from(context)?;
 
+        for (parameter, parameter_type) in parameters.iter().zip(parameter_types.iter()) {
+            function_context.set_type(parameter.inner().clone(), parameter_type.clone())?;
+        }
+
         let body_node = node.child(child_count - 1).unwrap();
         let body = Block::from_syntax(body_node, source, &function_context)?;
 
diff --git a/src/abstract_tree/type.rs b/src/abstract_tree/type.rs
index b504384..9326cb5 100644
--- a/src/abstract_tree/type.rs
+++ b/src/abstract_tree/type.rs
@@ -41,6 +41,8 @@ impl Type {
     }
 
     pub fn check(&self, other: &Type) -> Result<()> {
+        log::info!("Checking type {self} against {other}.");
+
         match (self, other) {
             (Type::Any, _)
             | (_, Type::Any)
diff --git a/src/value/map.rs b/src/value/map.rs
index 6a94607..0754080 100644
--- a/src/value/map.rs
+++ b/src/value/map.rs
@@ -80,6 +80,14 @@ impl Map {
         Ok(previous)
     }
 
+    pub fn set_type(&self, key: String, r#type: Type) -> Result<Option<(Value, Type)>> {
+        log::info!("Setting type {key} = {}", r#type);
+
+        let previous = self.variables.write()?.insert(key, (Value::none(), r#type));
+
+        Ok(previous)
+    }
+
     pub fn unset_all(&self) -> Result<()> {
         for (_key, (value, r#_type)) in self.variables.write()?.iter_mut() {
             *value = Value::none();
diff --git a/tests/types.rs b/tests/types.rs
index 8dabad4..b1e88a6 100644
--- a/tests/types.rs
+++ b/tests/types.rs
@@ -12,6 +12,8 @@ fn simple_type_check() {
 
 #[test]
 fn argument_count_check() {
+    env_logger::builder().is_test(true).try_init().unwrap();
+
     let source = "
             foo = (x <int>) <int> {
                 x