Add array type inference

grammar.ebnf

@@ -38,7 +38,7 @@ expression =
     function_call |
     function_call_failed |
     identifier |
-    list |
+    array |
     null |
     number |
     parentheses |
@@ -119,10 +119,10 @@ boolean = 'true' | 'false' ;
 (* `Null` literal *)
 null = 'null' ;
 
-(* `List` literal *)
-empty_list = '[', TYPE, ']' ;
-full_list = '[', [ expression, { ',', expression } ], ']' ;
-list = empty_list | full_list ;
+(* `Array` literal *)
+empty_array = '[', [ TYPE ], ']' ;
+full_array = '[', [ expression, { ',', expression } ], ']' ;
+array = empty_array | full_array ;
 
 (* Command expression *)
 (* The ordering of command modifiers doesn't matter *)

src/modules/expression/binop/add.rs

@@ -55,7 +55,7 @@ impl TypeCheckModule for Add {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        self.kind = Self::typecheck_allowed_types(meta, "addition", &self.left, &self.right, &[
+        self.kind = Self::typecheck_allowed_types(meta, "addition", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
             Type::Text,

src/modules/expression/binop/and.rs

@@ -52,7 +52,9 @@ impl TypeCheckModule for And {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        Self::typecheck_equality(meta, &self.left, &self.right)?;
+        Self::typecheck_allowed_types(meta, "logical AND", &mut self.left, &mut self.right, &[
+            Type::Bool,
+        ])?;
         Ok(())
     }
 }

src/modules/expression/binop/div.rs

@@ -55,7 +55,7 @@ impl TypeCheckModule for Div {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        self.kind = Self::typecheck_allowed_types(meta, "division", &self.left, &self.right, &[
+        self.kind = Self::typecheck_allowed_types(meta, "division", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
         ])?;

src/modules/expression/binop/eq.rs

@@ -53,7 +53,7 @@ impl TypeCheckModule for Eq {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        Self::typecheck_equality(meta, &self.left, &self.right)?;
+        Self::typecheck_equality(meta, &mut self.left, &mut self.right)?;
         Ok(())
     }
 }

src/modules/expression/binop/ge.rs

@@ -52,7 +52,7 @@ impl TypeCheckModule for Ge {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        Self::typecheck_allowed_types(meta, "comparison", &self.left, &self.right, &[
+        Self::typecheck_allowed_types(meta, "comparison", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
             Type::Text,

src/modules/expression/binop/gt.rs

@@ -52,7 +52,7 @@ impl TypeCheckModule for Gt {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        Self::typecheck_allowed_types(meta, "comparison", &self.left, &self.right, &[
+        Self::typecheck_allowed_types(meta, "comparison", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
             Type::Text,

src/modules/expression/binop/le.rs

@@ -52,7 +52,7 @@ impl TypeCheckModule for Le {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        Self::typecheck_allowed_types(meta, "comparison", &self.left, &self.right, &[
+        Self::typecheck_allowed_types(meta, "comparison", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
             Type::Text,

src/modules/expression/binop/lt.rs

@@ -52,7 +52,7 @@ impl TypeCheckModule for Lt {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        Self::typecheck_allowed_types(meta, "comparison", &self.left, &self.right, &[
+        Self::typecheck_allowed_types(meta, "comparison", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
             Type::Text,

src/modules/expression/binop/mod.rs

@@ -2,7 +2,7 @@ use heraclitus_compiler::prelude::*;
 use crate::modules::types::{Type, Typed};
 use crate::utils::metadata::ParserMetadata;
 use crate::utils::pluralize;
-use super::super::expression::expr::Expr;
+use super::super::expression::expr::{Expr, ExprType};
 use crate::modules::typecheck::TypeCheckModule;
 
 pub mod add;
@@ -28,8 +28,8 @@ pub trait BinOp: SyntaxModule<ParserMetadata> + TypeCheckModule {
     fn typecheck_allowed_types(
         meta: &mut ParserMetadata,
         operator: &str,
-        left: &Expr,
-        right: &Expr,
+        left: &mut Expr,
+        right: &mut Expr,
         allowed_types: &[Type],
     ) -> Result<Type, Failure> {
         let left_type = left.get_type();
@@ -51,13 +51,31 @@ pub trait BinOp: SyntaxModule<ParserMetadata> + TypeCheckModule {
 
     fn typecheck_equality(
         meta: &mut ParserMetadata,
-        left: &Expr,
-        right: &Expr,
+        left: &mut Expr,
+        right: &mut Expr,
     ) -> Result<Type, Failure> {
         match (left.get_type(), right.get_type()) {
             (Type::Int, Type::Num) | (Type::Num, Type::Int) => {
                 Ok(Type::Num)
             }
+            // Array type inference
+            (Type::Array(left_inner), Type::Array(right_inner)) if *left_inner == Type::Generic || *right_inner == Type::Generic => {
+                if *left_inner == Type::Generic && *right_inner != Type::Generic {
+                    if let Some(ExprType::VariableGet(var)) = &left.value {
+                        meta.update_var_type(&var.name, Type::Array(right_inner.clone()));
+                    }
+                    left.kind = Type::Array(right_inner.clone());
+                    Ok(Type::Array(right_inner))
+                } else if *left_inner != Type::Generic && *right_inner == Type::Generic {
+                    if let Some(ExprType::VariableGet(var)) = &right.value {
+                        meta.update_var_type(&var.name, Type::Array(left_inner.clone()));
+                    }
+                    right.kind = Type::Array(left_inner.clone());
+                    Ok(Type::Array(left_inner))
+                } else {
+                    Ok(Type::Array(Box::new(Type::Generic)))
+                }
+            }
             (left_type, right_type) => {
                 if left_type != right_type {
                     let pos = get_binop_position_info(meta, left, right);

src/modules/expression/binop/modulo.rs

@@ -53,7 +53,7 @@ impl TypeCheckModule for Modulo {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        self.kind = Self::typecheck_allowed_types(meta, "modulo", &self.left, &self.right, &[
+        self.kind = Self::typecheck_allowed_types(meta, "modulo", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
         ])?;

src/modules/expression/binop/mul.rs

@@ -54,7 +54,7 @@ impl TypeCheckModule for Mul {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        self.kind = Self::typecheck_allowed_types(meta, "multiplication", &self.left, &self.right, &[
+        self.kind = Self::typecheck_allowed_types(meta, "multiplication", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
         ])?;

src/modules/expression/binop/neq.rs

@@ -53,7 +53,7 @@ impl TypeCheckModule for Neq {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        Self::typecheck_equality(meta, &self.left, &self.right)?;
+        Self::typecheck_equality(meta, &mut self.left, &mut self.right)?;
         Ok(())
     }
 }

src/modules/expression/binop/or.rs

@@ -52,7 +52,9 @@ impl TypeCheckModule for Or {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        Self::typecheck_equality(meta, &self.left, &self.right)?;
+        Self::typecheck_allowed_types(meta, "logical OR", &mut self.left, &mut self.right, &[
+            Type::Bool,
+        ])?;
         Ok(())
     }
 }

src/modules/expression/binop/range.rs

@@ -56,7 +56,7 @@ impl TypeCheckModule for Range {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.from.typecheck(meta)?;
         self.to.typecheck(meta)?;
-        Self::typecheck_allowed_types(meta, "range operator", &self.from, &self.to, &[Type::Int])?;
+        Self::typecheck_allowed_types(meta, "range operator", &mut self.from, &mut self.to, &[Type::Int])?;
         Ok(())
     }
 }

src/modules/expression/binop/sub.rs

@@ -54,7 +54,7 @@ impl TypeCheckModule for Sub {
     fn typecheck(&mut self, meta: &mut ParserMetadata) -> SyntaxResult {
         self.left.typecheck(meta)?;
         self.right.typecheck(meta)?;
-        self.kind = Self::typecheck_allowed_types(meta, "subtraction", &self.left, &self.right, &[
+        self.kind = Self::typecheck_allowed_types(meta, "subtraction", &mut self.left, &mut self.right, &[
             Type::Num,
             Type::Int,
         ])?;

src/modules/expression/literal/array.rs

@@ -29,7 +29,8 @@ impl SyntaxModule<ParserMetadata> for Array {
         token(meta, "[")?;
         let tok = meta.get_current_token();
         if token(meta, "]").is_ok() {
-            return error!(meta, tok, "Expected array type or value before ']'", "Eg. insert 'Num' for empty array or '1, 2, 3' for array with values")
+            self.kind = Type::Array(Box::new(Type::Generic));
+            return Ok(())
         }
         // Try to parse array type
         match try_parse_type(meta) {

src/modules/expression/ternop/ternary.rs

@@ -11,12 +11,13 @@ use super::TernOp;
 pub struct Ternary {
     cond: Box<Expr>,
     true_expr: Box<Expr>,
-    false_expr: Box<Expr>
+    false_expr: Box<Expr>,
+    kind: Type
 }
 
 impl Typed for Ternary {
     fn get_type(&self) -> Type {
-        self.true_expr.get_type()
+        self.kind.clone()
     }
 }
 
@@ -51,7 +52,8 @@ impl SyntaxModule<ParserMetadata> for Ternary {
         Ternary {
             cond: Box::new(Expr::new()),
             true_expr: Box::new(Expr::new()),
-            false_expr: Box::new(Expr::new())
+            false_expr: Box::new(Expr::new()),
+            kind: Type::Null
         }
     }
 
@@ -71,13 +73,34 @@ impl TypeCheckModule for Ternary {
 
         self.true_expr.typecheck(meta)?;
         self.false_expr.typecheck(meta)?;
-        if self.true_expr.get_type() != self.false_expr.get_type() {
+
+        let true_type = self.true_expr.get_type();
+        let false_type = self.false_expr.get_type();
+
+        if true_type == false_type {
+            self.kind = true_type;
+        } else {
+            // Handle Array type inference
+            let mut resolved_type = None;
+            if let (Type::Array(t), Type::Array(f)) = (&true_type, &false_type) {
+                if **t == Type::Generic && **f != Type::Generic {
+                    resolved_type = Some(false_type.clone());
+                } else if **t != Type::Generic && **f == Type::Generic {
+                    resolved_type = Some(true_type.clone());
+                }
+            }
+
+            if let Some(kind) = resolved_type {
+                self.kind = kind;
+                return Ok(());
+            }
+
             let pos = get_binop_position_info(meta, &self.true_expr, &self.false_expr);
             let msg = Message::new_err_at_position(meta, pos)
                 .message("Ternary operation can only evaluate to value of one type.")
                 .comment(format!("Provided branches of type '{}' and '{}'.",
-                    self.true_expr.get_type(),
-                    self.false_expr.get_type()));
+                    true_type,
+                    false_type));
             return Err(Failure::Loud(msg));
         }
         Ok(())

src/modules/function/invocation.rs

@@ -126,6 +126,20 @@ impl TypeCheckModule for FunctionInvocation {
             }
         }
 
+        // Check for type inference on reference arguments
+        for (arg_expr, fun_arg) in izip!(&mut self.args, &function_unit.args) {
+            if fun_arg.is_ref {
+                if let (Type::Array(inner), Type::Array(expected_inner)) = (arg_expr.get_type(), &fun_arg.kind) {
+                    if *inner == Type::Generic && **expected_inner != Type::Generic {
+                        if let Some(ExprType::VariableGet(var)) = &arg_expr.value {
+                            meta.update_var_type(&var.name, Type::array_of(*expected_inner.clone()));
+                            arg_expr.kind = Type::array_of(*expected_inner.clone());
+                        }
+                    }
+                }
+            }
+        }
+
         // Validate arguments and get function variant
         let types = self.args.iter().map(Expr::get_type).collect::<Vec<Type>>();
         let var_refs = self.args.iter().map(is_ref).collect::<Vec<bool>>();

src/modules/shorthand/add.rs

@@ -4,7 +4,7 @@ use crate::modules::expression::expr::Expr;
 use crate::modules::variable::{handle_variable_reference, prevent_constant_mutation, variable_name_extensions};
 use crate::translate::compute::translate_computation_eval;
 use crate::translate::{compute::ArithOp, module::TranslateModule};
-use crate::modules::types::Type;
+use crate::modules::types::{Type, Typed};
 
 use super::shorthand_typecheck_allowed_types;
 
@@ -52,6 +52,15 @@ impl TypeCheckModule for ShorthandAdd {
         self.global_id = variable.global_id;
         self.is_ref = variable.is_ref;
 
+        let right_type = self.expr.get_type();
+        if let (Type::Array(inner_left), Type::Array(inner_right)) = (&self.kind, &right_type) {
+            if *inner_left.as_ref() == Type::Generic && *inner_right.as_ref() != Type::Generic {
+                meta.update_var_type(&self.var, right_type.clone());
+                self.kind = right_type;
+                return Ok(());
+            }
+        }
+
         shorthand_typecheck_allowed_types(meta, "add", &self.kind, &self.expr, &[
             Type::Num,
             Type::Int,

src/modules/types.rs

@@ -35,7 +35,15 @@ impl Type {
     }
 
     pub fn is_allowed_in(&self, other: &Type) -> bool {
-        self == other || self.is_subset_of(other)
+        if self == other || self.is_subset_of(other) {
+            return true;
+        }
+
+        if let (Type::Array(const_type), Type::Array(other_type)) = (self, other) {
+            return **const_type == Type::Generic && **other_type != Type::Generic;
+        }
+
+        false
     }
 
     pub fn is_array(&self) -> bool {
@@ -178,6 +186,7 @@ mod tests {
         let b = Type::Array(Box::new(Type::Generic));
 
         assert!(!b.is_subset_of(&a));
+        assert!(b.is_allowed_in(&a));
     }
 
     #[test]

src/modules/variable/set.rs

@@ -70,15 +70,33 @@ impl TypeCheckModule for VariableSet {
 
         if self.index.is_some() {
             if let Type::Array(kind) = &self.var_type {
+                // Handle type inference
+                if **kind == Type::Generic {
+                    let new_type = Type::array_of(right_type.clone());
+                    meta.update_var_type(&self.name, new_type.clone());
+                    self.var_type = new_type;
+                    return Ok(());
+                }
+
                 if !right_type.is_allowed_in(kind) {
                     let tok = self.expr.get_position();
                     return error_pos!(meta, tok, format!("Cannot assign value of type '{right_type}' to an array of '{kind}'"));
                 }
             }
-        }
-        else if !right_type.is_allowed_in(&self.var_type) {
-            let tok = self.expr.get_position();
-            return error_pos!(meta, tok, format!("Cannot assign value of type '{right_type}' to a variable of type '{}'", self.var_type));
+        } else {
+            // Check for type inference
+            if let (Type::Array(inner_var), Type::Array(inner_right)) = (&self.var_type, &right_type) {
+                if **inner_var == Type::Generic && **inner_right != Type::Generic {
+                    meta.update_var_type(&self.name, right_type.clone());
+                    self.var_type = right_type;
+                    return Ok(());
+                }
+            }
+
+            if !right_type.is_allowed_in(&self.var_type) {
+                let tok = self.expr.get_position();
+                return error_pos!(meta, tok, format!("Cannot assign value of type '{right_type}' to a variable of type '{}'", self.var_type));
+            }
         }
 
         Ok(())

src/tests/erroring/type_inference_array_binary_op_mismatch.ab

@@ -0,0 +1,7 @@
+// Output
+// Expected both operands to be of the same type, but got '[Int]' and '[Text]'.
+
+let a = []
+a = a + [1]
+a = a + ["test"]
+echo a