Clean up

src/abstract_tree/assignment.rs

@@ -1,7 +1,9 @@
 use serde::{Deserialize, Serialize};
 use tree_sitter::Node;
 
-use crate::{AbstractTree, Error, Identifier, Map, Result, Statement, Type, TypeDefinition, Value};
+use crate::{
+    AbstractTree, Error, Function, Identifier, Map, Result, Statement, Type, TypeDefinition, Value,
+};
 
 #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, PartialOrd, Ord)]
 pub struct Assignment {
@@ -56,6 +58,45 @@ impl AbstractTree for Assignment {
         let statement_node = node.child_by_field_name("statement").unwrap();
         let statement = Statement::from_syntax_node(source, statement_node, context)?;
 
+        if let Some(type_definition) = &type_definition {
+            let statement_type = statement.expected_type(context)?;
+
+            match operator {
+                AssignmentOperator::Equal => {
+                    type_definition.abstract_check(
+                        &statement_type,
+                        context,
+                        statement_node,
+                        source,
+                    )?;
+                }
+                AssignmentOperator::PlusEqual => {
+                    let identifier_type = identifier.expected_type(context)?;
+
+                    type_definition.abstract_check(
+                        &identifier_type,
+                        context,
+                        type_node.unwrap(),
+                        source,
+                    )?;
+
+                    let type_definition = if let Type::List(item_type) = type_definition.inner() {
+                        TypeDefinition::new(item_type.as_ref().clone())
+                    } else {
+                        type_definition.clone()
+                    };
+
+                    type_definition.abstract_check(
+                        &statement_type,
+                        context,
+                        statement_node,
+                        source,
+                    )?;
+                }
+                AssignmentOperator::MinusEqual => todo!(),
+            }
+        }
+
         Ok(Assignment {
             identifier,
             type_definition,
@@ -72,13 +113,19 @@ impl AbstractTree for Assignment {
             AssignmentOperator::PlusEqual => {
                 if let Some(mut previous_value) = context.variables()?.get(key).cloned() {
                     if let Ok(list) = previous_value.as_list() {
-                        let first_item_type = if let Some(first) = list.items().first() {
+                        let item_type = if let Some(type_defintion) = &self.type_definition {
+                            if let Type::List(item_type) = type_defintion.inner() {
+                                TypeDefinition::new(item_type.as_ref().clone())
+                            } else {
+                                TypeDefinition::new(Type::Empty)
+                            }
+                        } else if let Some(first) = list.items().first() {
                             first.r#type(context)?
                         } else {
                             TypeDefinition::new(Type::Any)
                         };
 
-                        first_item_type.check(&value.r#type(context)?, context)?;
+                        item_type.runtime_check(&value.r#type(context)?, context)?;
                     }
 
                     previous_value += value;
@@ -95,15 +142,24 @@ impl AbstractTree for Assignment {
                     return Err(Error::VariableIdentifierNotFound(key.clone()));
                 }
             }
-            AssignmentOperator::Equal => {
-                if let Some(type_definition) = &self.type_definition {
-                    let new_value_type = value.r#type(context)?;
+            AssignmentOperator::Equal => value,
+        };
 
-                    type_definition.check(&new_value_type, context)?;
-                }
+        let new_value = if let Some(type_definition) = &self.type_definition {
+            let new_value_type = new_value.r#type(context)?;
 
-                value
+            type_definition.runtime_check(&new_value_type, context)?;
+
+            if let Value::Function(function) = new_value {
+                Value::Function(Function::new(
+                    function.parameters().clone(),
+                    function.body().clone(),
+                ))
+            } else {
+                new_value
             }
+        } else {
+            new_value
         };
 
         context.variables_mut()?.insert(key.clone(), new_value);
@@ -115,3 +171,49 @@ impl AbstractTree for Assignment {
         Ok(TypeDefinition::new(Type::Empty))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{evaluate, List, Value};
+
+    #[test]
+    fn simple_assignment() {
+        let test = evaluate("x = 1 x").unwrap();
+
+        assert_eq!(Value::Integer(1), test);
+    }
+
+    #[test]
+    fn simple_assignment_with_type() {
+        let test = evaluate("x <int> = 1 x").unwrap();
+
+        assert_eq!(Value::Integer(1), test);
+    }
+
+    #[test]
+    fn list_add_assign() {
+        let test = evaluate(
+            "
+            x <list int> = []
+            x += 1
+            x
+        ",
+        )
+        .unwrap();
+
+        assert_eq!(Value::List(List::with_items(vec![Value::Integer(1)])), test);
+    }
+
+    #[test]
+    fn function_assignment() {
+        let test = evaluate(
+            "
+            foobar <fn str -> str> = |text| { 'hi' }
+            foobar
+        ",
+        )
+        .unwrap();
+
+        assert_eq!(Value::String("hi".to_string()), test);
+    }
+}

src/abstract_tree/type_defintion.rs

@@ -23,7 +23,23 @@ impl TypeDefinition {
         self.r#type
     }
 
-    pub fn check(&self, other: &TypeDefinition, context: &Map) -> Result<()> {
+    pub fn abstract_check(
+        &self,
+        other: &TypeDefinition,
+        context: &Map,
+        node: Node,
+        source: &str,
+    ) -> Result<()> {
+        self.runtime_check(other, context)
+            .map_err(|_| Error::TypeCheck {
+                expected: self.clone(),
+                actual: other.clone(),
+                location: node.start_position(),
+                source: source[node.byte_range()].to_string(),
+            })
+    }
+
+    pub fn runtime_check(&self, other: &TypeDefinition, context: &Map) -> Result<()> {
         match (&self.r#type, &other.r#type) {
             (Type::Any, _)
             | (_, Type::Any)
@@ -43,7 +59,7 @@ impl TypeDefinition {
                 let self_defintion = TypeDefinition::new(self_item_type.as_ref().clone());
                 let other_definition = &TypeDefinition::new(other_item_type.as_ref().clone());
 
-                self_defintion.check(other_definition, context)
+                self_defintion.runtime_check(other_definition, context)
             }
             _ => Err(Error::RuntimeTypeCheck {
                 expected: self.clone(),
@@ -85,8 +101,8 @@ pub enum Type {
     Empty,
     Float,
     Function {
-        parameter_types: Vec<TypeDefinition>,
-        return_type: Box<TypeDefinition>,
+        parameter_types: Vec<Type>,
+        return_type: Box<Type>,
     },
     Integer,
     List(Box<Type>),
@@ -115,7 +131,7 @@ impl AbstractTree for Type {
                     let parameter_type =
                         Type::from_syntax_node(source, parameter_type_node, context)?;
 
-                    parameter_types.push(TypeDefinition::new(parameter_type));
+                    parameter_types.push(parameter_type);
                 }
 
                 let return_type_node = node.child(child_count - 1).unwrap();
@@ -123,7 +139,7 @@ impl AbstractTree for Type {
 
                 Type::Function {
                     parameter_types,
-                    return_type: Box::new(TypeDefinition::new(return_type)),
+                    return_type: Box::new(return_type),
                 }
             }
             "int" => Type::Integer,

src/value/function.rs

@@ -12,6 +12,10 @@ pub struct Function {
 }
 
 impl Function {
+    pub fn new(parameters: Vec<Identifier>, body: Block) -> Self {
+        Self { parameters, body }
+    }
+
     pub fn parameters(&self) -> &Vec<Identifier> {
         &self.parameters
     }
@@ -44,13 +48,11 @@ impl AbstractTree for Function {
     }
 
     fn run(&self, source: &str, context: &Map) -> Result<Value> {
-        let return_value = self.body.run(source, context)?;
-
-        Ok(return_value)
+        self.body.run(source, context)
     }
 
     fn expected_type(&self, context: &Map) -> Result<TypeDefinition> {
-        self.body.expected_type(context)
+        Value::Function(self.clone()).r#type(context)
     }
 }
 

src/value/mod.rs

@@ -69,13 +69,10 @@ impl Value {
             Value::Map(_) => Type::Map,
             Value::Table(_) => Type::Table,
             Value::Function(function) => {
-                let parameter_types = Vec::new();
-
-                for identifier in function.parameters() {
-                    let _type = identifier.expected_type(context)?;
-                }
-
-                let return_type = function.body().expected_type(context)?;
+                let parameters = function.parameters();
+                let parameter_types = vec![Type::Any; parameters.len()];
+                let body = function.body();
+                let return_type = body.expected_type(context)?.take_inner();
 
                 Type::Function {
                     parameter_types,