Make nested arrays work

reid-llvm-lib/src/builder.rs

@@ -3,8 +3,6 @@
 
 use std::{cell::RefCell, rc::Rc};
 
-use llvm_sys::core::LLVMBuildArrayAlloca;
-
 use crate::{
     BlockData, ConstValue, FunctionData, Instr, InstructionData, ModuleData, TerminatorKind, Type,
     util::match_types,
@@ -270,6 +268,7 @@ impl Builder {
                 }
                 Store(ptr, _) => {
                     if let Ok(ty) = ptr.get_type(&self) {
+                        dbg!(&ty);
                         if let Type::Ptr(_) = ty {
                             Ok(())
                         } else {
@@ -279,24 +278,11 @@ impl Builder {
                         Err(())
                     }
                 }
-                Extract(arr, idx) => {
-                    let arr_ty = arr.get_type(&self)?;
-                    if let Type::Array(_, len) = arr_ty {
-                        if len > idx { Ok(()) } else { Err(()) }
-                    } else {
-                        Err(())
-                    }
-                }
                 ArrayAlloca(_, _) => Ok(()),
-                Insert(arr, idx, val) => {
+                GetElemPtr(arr, _) => {
                     let arr_ty = arr.get_type(&self)?;
-                    let val_ty = val.get_type(&self)?;
-                    if let Type::Array(elem_ty, len) = arr_ty {
-                        if val_ty == *elem_ty && len > idx {
-                            Ok(())
-                        } else {
-                            Err(())
-                        }
+                    if let Type::Ptr(_) = arr_ty {
+                        Ok(())
                     } else {
                         Err(())
                     }
@@ -362,15 +348,8 @@ impl InstructionValue {
                 Alloca(_, ty) => Ok(Type::Ptr(Box::new(ty.clone()))),
                 Load(_, ty) => Ok(ty.clone()),
                 Store(_, value) => value.get_type(builder),
-                Extract(arr, _) => match arr.get_type(builder) {
-                    Ok(Type::Array(elem_t, _)) => Ok(*elem_t),
-                    Ok(_) => Err(()),
-                    Err(_) => Err(()),
-                },
-                ArrayAlloca(ty, len_value) => {
-                    Ok(Type::Array(Box::new(ty.clone()), len_value.clone()))
-                }
-                Insert(_, _, val) => val.get_type(builder),
+                ArrayAlloca(ty, _) => Ok(Type::Ptr(Box::new(ty.clone()))),
+                GetElemPtr(ptr, _) => ptr.get_type(builder),
             }
         }
     }
@@ -391,15 +370,6 @@ impl ConstValue {
             ConstValue::U64(_) => U64,
             ConstValue::U128(_) => U128,
             ConstValue::Bool(_) => Bool,
-            // ConstValue::Array(arr) => Array(
-            //     Box::new(
-            //         arr.iter()
-            //             .map(|a| a.get_type(builder).unwrap())
-            //             .next()
-            //             .unwrap_or(Void),
-            //     ),
-            //     arr.len() as u32,
-            // ),
         }
     }
 }
@@ -420,7 +390,6 @@ impl Type {
             Type::Bool => true,
             Type::Void => false,
             Type::Ptr(_) => false,
-            Type::Array(_, _) => false,
         }
     }
 
@@ -439,7 +408,6 @@ impl Type {
             Type::Bool => false,
             Type::Void => false,
             Type::Ptr(_) => false,
-            Type::Array(_, _) => false,
         }
     }
 }

reid-llvm-lib/src/compile.rs

@@ -349,33 +349,9 @@ impl InstructionHolder {
                     module.values.get(&val).unwrap().value_ref,
                     module.values.get(&ptr).unwrap().value_ref,
                 ),
-                Extract(arr, idx) => {
-                    let t = arr.get_type(module.builder).unwrap();
-                    let Type::Array(elem_t, _) = t else { panic!() };
-
-                    let indices = &mut [ConstValue::I32(*idx as i32).as_llvm(module.context_ref)];
-                    let ptr = LLVMBuildGEP2(
-                        module.builder_ref,
-                        elem_t.as_llvm(module.context_ref),
-                        module.values.get(arr).unwrap().value_ref,
-                        indices.as_mut_ptr(),
-                        1,
-                        c"insert_gep".as_ptr(),
-                    );
-                    LLVMBuildLoad2(
-                        module.builder_ref,
-                        elem_t.as_llvm(module.context_ref),
-                        ptr,
-                        c"load".as_ptr(),
-                    )
-                }
                 ArrayAlloca(ty, len) => {
                     let array_len = ConstValue::U16(*len as u16).as_llvm(module.context_ref);
-                    let array_ty = Type::Array(Box::new(ty.clone()), *len);
-                    dbg!(
-                        &ty.as_llvm(module.context_ref),
-                        &array_ty.as_llvm(module.context_ref)
-                    );
+                    let array_ty = Type::Ptr(Box::new(ty.clone()));
                     LLVMBuildArrayAlloca(
                         module.builder_ref,
                         ty.as_llvm(module.context_ref),
@@ -383,31 +359,23 @@ impl InstructionHolder {
                         c"array_alloca".as_ptr(),
                     )
                 }
-                Insert(arr, idx, val) => {
-                    let indices = &mut [ConstValue::I32(*idx as i32).as_llvm(module.context_ref)];
-                    let ptr = LLVMBuildGEP2(
+                GetElemPtr(arr, indices) => {
+                    let t = arr.get_type(module.builder).unwrap();
+                    let Type::Ptr(elem_t) = t else { panic!() };
+
+                    let mut indices: Vec<_> = indices
+                        .iter()
+                        .map(|idx| ConstValue::U32(*idx).as_llvm(module.context_ref))
+                        .collect();
+
+                    LLVMBuildGEP2(
                         module.builder_ref,
-                        val.get_type(module.builder)
-                            .unwrap()
-                            .as_llvm(module.context_ref),
+                        elem_t.as_llvm(module.context_ref),
                         module.values.get(arr).unwrap().value_ref,
                         indices.as_mut_ptr(),
-                        1,
-                        c"insert_gep".as_ptr(),
-                    );
-                    LLVMBuildStore(
-                        module.builder_ref,
-                        module.values.get(val).unwrap().value_ref,
-                        ptr,
+                        indices.len() as u32,
+                        c"array_gep".as_ptr(),
                     )
-
-                    // LLVMBuildInsertValue(
-                    //     module.builder_ref,
-                    //     ptr,
-                    //     module.values.get(val).unwrap().value_ref,
-                    //     *idx,
-                    //     c"insert".as_ptr(),
-                    // )
                 }
             }
         };
@@ -505,7 +473,6 @@ impl Type {
                 Bool => LLVMInt1TypeInContext(context),
                 Void => LLVMVoidType(),
                 Ptr(ty) => LLVMPointerType(ty.as_llvm(context), 0),
-                Array(elem_t, _) => LLVMPointerType(elem_t.as_llvm(context), 0),
             }
         }
     }

reid-llvm-lib/src/debug.rs

@@ -100,11 +100,18 @@ impl Debug for Instr {
             Instr::Alloca(name, ty) => write!(f, "alloca<{:?}>({})", ty, name),
             Instr::Load(val, ty) => write!(f, "load<{:?}>({:?})", ty, val),
             Instr::Store(ptr, val) => write!(f, "store({:?} = {:?})", ptr, val),
-            Instr::Extract(instruction_value, idx) => fmt_index(f, instruction_value, idx),
             Instr::ArrayAlloca(ty, instruction_value) => {
                 write!(f, "array_alloca<{:?}>({:?})", ty, instruction_value)
             }
-            Instr::Insert(arr, idx, val) => write!(f, "{:?}[{}] = {:?}", arr, idx, val),
+            Instr::GetElemPtr(instruction_value, items) => fmt_index(
+                f,
+                instruction_value,
+                &items
+                    .iter()
+                    .map(|i| i.to_string())
+                    .collect::<Vec<_>>()
+                    .join(", "),
+            ),
         }
     }
 }

reid-llvm-lib/src/lib.rs

@@ -178,8 +178,7 @@ pub enum Instr {
     Load(InstructionValue, Type),
     Store(InstructionValue, InstructionValue),
     ArrayAlloca(Type, u32),
-    Insert(InstructionValue, u32, InstructionValue),
-    Extract(InstructionValue, u32),
+    GetElemPtr(InstructionValue, Vec<u32>),
 
     /// Integer Comparison
     ICmp(CmpPredicate, InstructionValue, InstructionValue),
@@ -202,7 +201,6 @@ pub enum Type {
     Bool,
     Void,
     Ptr(Box<Type>),
-    Array(Box<Type>, u32),
 }
 
 #[derive(Debug, Clone, Hash)]

reid/examples/reid/array.reid

@@ -1,7 +1,7 @@
 // Arithmetic, function calls and imports!
 
-fn array() -> [u16; 4] {
-    return [10, 15, 7, 9];
+fn array() -> [[u16; 4]; 1] {
+    return [[10, 15, 7, 9]];
 }
 
 fn main() -> u16 {
@@ -9,7 +9,7 @@ fn main() -> u16 {
 
     let mut list = array();
 
-    // list[1] = 5;
+    list[0][3] = 5;
 
-    return list[3];
+    return list[0][3];
 }

reid/src/ast/process.rs

@@ -171,9 +171,11 @@ impl ast::Expression {
             ast::ExpressionKind::Array(expressions) => {
                 mir::ExprKind::Array(expressions.iter().map(|e| e.process()).collect())
             }
-            ast::ExpressionKind::Index(expression, idx) => {
-                mir::ExprKind::Index(Box::new(expression.process()), *idx)
-            }
+            ast::ExpressionKind::Index(expression, idx) => mir::ExprKind::Index(
+                Box::new(expression.process()),
+                mir::TypeKind::Vague(mir::VagueType::Unknown),
+                *idx,
+            ),
         };
 
         mir::Expression(kind, self.1.into())

reid/src/codegen.rs

@@ -5,7 +5,7 @@ use reid_lib::{
     TerminatorKind as Term, Type,
 };
 
-use crate::mir::{self, types::ReturnType, NamedVariableRef, TypeKind};
+use crate::mir::{self, types::ReturnType, IndexedVariableReference, NamedVariableRef, TypeKind};
 
 /// Context that contains all of the given modules as complete codegenerated
 /// LLIR that can then be finally compiled into LLVM IR.
@@ -153,6 +153,31 @@ impl<'ctx, 'a> Scope<'ctx, 'a> {
     }
 }
 
+impl IndexedVariableReference {
+    fn get_stack_value(&self, scope: &mut Scope) -> Option<(StackValue, Vec<u32>)> {
+        use StackValueKind as Kind;
+
+        match &self.kind {
+            mir::IndexedVariableReferenceKind::Named(NamedVariableRef(_, name, _)) => scope
+                .stack_values
+                .get(name)
+                .cloned()
+                .map(|v| (v, Vec::new())),
+            mir::IndexedVariableReferenceKind::Index(inner, idx) => {
+                let (inner_val, mut indices) = inner.get_stack_value(scope)?;
+
+                match &inner_val.1 {
+                    Type::Ptr(_) => {
+                        indices.push(*idx as u32);
+                        Some((inner_val, indices))
+                    }
+                    _ => panic!("Tried to codegen indexing a non-indexable value!"),
+                }
+            }
+        }
+    }
+}
+
 impl mir::Statement {
     fn codegen<'ctx, 'a>(&self, scope: &mut Scope<'ctx, 'a>) -> Option<InstructionValue> {
         match &self.0 {
@@ -163,14 +188,17 @@ impl mir::Statement {
                     StackValue(
                         match mutable {
                             false => StackValueKind::Immutable(value),
-                            true => StackValueKind::Mutable({
-                                let alloca = scope
-                                    .block
-                                    .build(Instr::Alloca(name.clone(), ty.get_type()))
-                                    .unwrap();
-                                scope.block.build(Instr::Store(alloca, value)).unwrap();
-                                alloca
-                            }),
+                            true => match ty {
+                                TypeKind::Array(_, _) => StackValueKind::Mutable(value),
+                                _ => StackValueKind::Mutable({
+                                    let alloca = scope
+                                        .block
+                                        .build(Instr::Alloca(name.clone(), ty.get_type()))
+                                        .unwrap();
+                                    scope.block.build(Instr::Store(alloca, value)).unwrap();
+                                    alloca
+                                }),
+                            },
                         },
                         ty.get_type(),
                     ),
@@ -178,20 +206,31 @@ impl mir::Statement {
                 None
             }
             mir::StmtKind::Set(var, val) => {
-                todo!("Re-think how set needs to work with arrays");
-                // if let Some(StackValue(kind, _)) = scope.stack_values.get(&var.1).cloned() {
-                //     match kind {
-                //         StackValueKind::Immutable(_) => {
-                //             panic!("Tried to mutate an immutable variable")
-                //         }
-                //         StackValueKind::Mutable(ptr) => {
-                //             let expression = val.codegen(scope).unwrap();
-                //             Some(scope.block.build(Instr::Store(ptr, expression)).unwrap())
-                //         }
-                //     }
-                // } else {
-                //     panic!("")
-                // }
+                if let Some((StackValue(kind, mut ty), indices)) = var.get_stack_value(scope) {
+                    match kind {
+                        StackValueKind::Immutable(_) => {
+                            panic!("Tried to mutate an immutable variable")
+                        }
+                        StackValueKind::Mutable(mut ptr) => {
+                            for (i, idx) in indices.iter().enumerate() {
+                                let Type::Ptr(inner) = ty else { panic!() };
+                                ty = *inner;
+
+                                ptr = scope
+                                    .block
+                                    .build(Instr::GetElemPtr(ptr, vec![*idx]))
+                                    .unwrap();
+                                if i < (indices.len() - 1) {
+                                    ptr = scope.block.build(Instr::Load(ptr, ty.clone())).unwrap()
+                                }
+                            }
+                            let expression = val.codegen(scope).unwrap();
+                            Some(scope.block.build(Instr::Store(ptr, expression)).unwrap())
+                        }
+                    }
+                } else {
+                    panic!("")
+                }
             }
             // mir::StmtKind::If(if_expression) => if_expression.codegen(scope),
             mir::StmtKind::Import(_) => todo!(),
@@ -267,9 +306,11 @@ impl mir::Expression {
                     .expect("Variable reference not found?!");
                 Some(match v.0 {
                     StackValueKind::Immutable(val) => val.clone(),
-                    StackValueKind::Mutable(val) => {
-                        scope.block.build(Instr::Load(val, v.1.clone())).unwrap()
-                    }
+                    StackValueKind::Mutable(val) => match v.1 {
+                        // TODO probably wrong ..?
+                        Type::Ptr(_) => val,
+                        _ => scope.block.build(Instr::Load(val, v.1.clone())).unwrap(),
+                    },
                 })
             }
             mir::ExprKind::Literal(lit) => Some(lit.as_const(&mut scope.block)),
@@ -334,12 +375,17 @@ impl mir::Expression {
                     None
                 }
             }
-            mir::ExprKind::Index(expression, idx) => {
-                let expr = expression.codegen(scope)?;
+            mir::ExprKind::Index(expression, val_t, idx) => {
+                let array = expression.codegen(scope)?;
+                let ptr = scope
+                    .block
+                    .build(Instr::GetElemPtr(array, vec![*idx as u32]))
+                    .unwrap();
+
                 Some(
                     scope
                         .block
-                        .build(Instr::Extract(expr, *idx as u32))
+                        .build(Instr::Load(ptr, val_t.get_type()))
                         .unwrap(),
                 )
             }
@@ -354,7 +400,6 @@ impl mir::Expression {
                     .next()
                     .unwrap_or(TypeKind::Void);
 
-                dbg!(&instr_t);
                 let array = scope
                     .block
                     .build(Instr::ArrayAlloca(
@@ -364,10 +409,11 @@ impl mir::Expression {
                     .unwrap();
 
                 for (i, instr) in instr_list.iter().enumerate() {
-                    scope
+                    let ptr = scope
                         .block
-                        .build(Instr::Insert(array, i as u32, *instr))
+                        .build(Instr::GetElemPtr(array, vec![i as u32]))
                         .unwrap();
+                    scope.block.build(Instr::Store(ptr, *instr)).unwrap();
                 }
 
                 Some(array)
@@ -450,7 +496,7 @@ impl TypeKind {
             TypeKind::U64 => Type::U64,
             TypeKind::U128 => Type::U128,
             TypeKind::Bool => Type::Bool,
-            TypeKind::Array(elem_t, len) => Type::Array(Box::new(elem_t.get_type()), *len as u32),
+            TypeKind::Array(elem_t, _) => Type::Ptr(Box::new(elem_t.get_type())),
             TypeKind::Void => panic!("Void not a supported type"),
             TypeKind::Vague(_) => panic!("Tried to compile a vague type!"),
         }

reid/src/mir/display.rs

@@ -128,8 +128,9 @@ impl Display for ExprKind {
             ExprKind::FunctionCall(fc) => Display::fmt(fc, f),
             ExprKind::If(if_exp) => Display::fmt(&if_exp, f),
             ExprKind::Block(block) => Display::fmt(block, f),
-            ExprKind::Index(expression, idx) => {
+            ExprKind::Index(expression, elem_ty, idx) => {
                 Display::fmt(&expression, f)?;
+                write!(f, "<{}>", elem_ty)?;
                 write_index(f, *idx)
             }
             ExprKind::Array(expressions) => {

reid/src/mir/mod.rs

@@ -208,7 +208,7 @@ pub struct Import(pub String, pub Metadata);
 #[derive(Debug)]
 pub enum ExprKind {
     Variable(NamedVariableRef),
-    Index(Box<Expression>, u64),
+    Index(Box<Expression>, TypeKind, u64),
     Array(Vec<Expression>),
     Literal(Literal),
     BinOp(BinaryOperator, Box<Expression>, Box<Expression>),

reid/src/mir/typecheck.rs

@@ -7,7 +7,7 @@ use TypeKind::*;
 use VagueType::*;
 
 use super::{
-    pass::{Pass, PassState, ScopeFunction, ScopeVariable},
+    pass::{Pass, PassState, ScopeFunction, ScopeVariable, Storage},
     typerefs::TypeRefs,
     types::ReturnType,
 };
@@ -169,41 +169,46 @@ impl Block {
                     None
                 }
                 StmtKind::Set(variable_reference, expression) => {
-                    todo!("Re-think how set needs to work with arrays")
-                    // if let Some(var) = state.scope.variables.get(&variable_reference.1).cloned() {
-                    //     // Typecheck expression and coerce to variable type
-                    //     let res = expression.typecheck(&mut state, &hints, Some(&var.ty));
+                    if let Some(var) = state
+                        .ok(
+                            variable_reference.get_variable(&state.scope.variables),
+                            variable_reference.meta,
+                        )
+                        .flatten()
+                    {
+                        // Typecheck expression and coerce to variable type
+                        let res = expression.typecheck(&mut state, &hints, Some(&var.ty));
 
-                    //     // If expression resolution itself was erronous, resolve as
-                    //     // Unknown.
-                    //     let expr_ty = state.or_else(res, Vague(Unknown), expression.1);
+                        // If expression resolution itself was erronous, resolve as
+                        // Unknown.
+                        let expr_ty = state.or_else(res, Vague(Unknown), expression.1);
 
-                    //     // Make sure the expression and variable type to really
-                    //     // be the same
-                    //     let res_t = state.or_else(
-                    //         expr_ty.collapse_into(&variable_reference.0.resolve_hinted(&hints)),
-                    //         Vague(Unknown),
-                    //         variable_reference.2 + expression.1,
-                    //     );
+                        // Make sure the expression and variable type to really
+                        // be the same
+                        let res_t = state.or_else(
+                            expr_ty.collapse_into(&var.ty.resolve_hinted(&hints)),
+                            Vague(Unknown),
+                            variable_reference.meta + expression.1,
+                        );
 
-                    //     // Update typing to be more accurate
-                    //     variable_reference.0 = res_t;
+                        // Update typing to be more accurate
+                        variable_reference.update_type(&res_t);
 
-                    //     if !var.mutable {
-                    //         state.ok::<_, Infallible>(
-                    //             Err(ErrorKind::VariableNotMutable(variable_reference.1.clone())),
-                    //             variable_reference.2,
-                    //         );
-                    //     }
+                        if !var.mutable {
+                            state.ok::<_, Infallible>(
+                                Err(ErrorKind::VariableNotMutable(variable_reference.get_name())),
+                                variable_reference.meta,
+                            );
+                        }
 
-                    //     None
-                    // } else {
-                    //     state.ok::<_, Infallible>(
-                    //         Err(ErrorKind::VariableNotDefined(variable_reference.1.clone())),
-                    //         variable_reference.2,
-                    //     );
-                    //     None
-                    // }
+                        None
+                    } else {
+                        state.ok::<_, Infallible>(
+                            Err(ErrorKind::VariableNotDefined(variable_reference.get_name())),
+                            variable_reference.meta,
+                        );
+                        None
+                    }
                 }
                 StmtKind::Import(_) => todo!(), // TODO
                 StmtKind::Expression(expression) => {
@@ -380,7 +385,7 @@ impl Expression {
                 Ok(collapsed)
             }
             ExprKind::Block(block) => block.typecheck(state, &hints, hint_t),
-            ExprKind::Index(expression, idx) => {
+            ExprKind::Index(expression, elem_ty, idx) => {
                 // Try to unwrap hint type from array if possible
                 let hint_t = hint_t.map(|t| match t {
                     Array(type_kind, _) => &type_kind,
@@ -388,11 +393,17 @@ impl Expression {
                 });
 
                 let expr_t = expression.typecheck(state, hints, hint_t)?;
-                if let TypeKind::Array(elem_t, len) = expr_t {
+                if let TypeKind::Array(inferred_ty, len) = expr_t {
                     if len < *idx {
                         return Err(ErrorKind::IndexOutOfBounds(*idx, len));
                     }
-                    Ok(*elem_t)
+                    let ty = state.or_else(
+                        elem_ty.resolve_hinted(hints).collapse_into(&inferred_ty),
+                        TypeKind::Vague(Unknown),
+                        self.1,
+                    );
+                    *elem_ty = ty.clone();
+                    Ok(ty)
                 } else {
                     Err(ErrorKind::TriedIndexingNonArray(expr_t))
                 }
@@ -432,6 +443,32 @@ impl Expression {
     }
 }
 
+impl IndexedVariableReference {
+    fn get_variable(
+        &self,
+        storage: &Storage<ScopeVariable>,
+    ) -> Result<Option<ScopeVariable>, ErrorKind> {
+        match &self.kind {
+            IndexedVariableReferenceKind::Named(NamedVariableRef(_, name, _)) => {
+                Ok(storage.get(&name).cloned())
+            }
+            IndexedVariableReferenceKind::Index(inner_ref, _) => {
+                if let Some(var) = inner_ref.get_variable(storage)? {
+                    match &var.ty {
+                        TypeKind::Array(inner_ty, _) => Ok(Some(ScopeVariable {
+                            ty: *inner_ty.clone(),
+                            mutable: var.mutable,
+                        })),
+                        _ => Err(ErrorKind::TriedIndexingNonArray(var.ty.clone())),
+                    }
+                } else {
+                    Ok(None)
+                }
+            }
+        }
+    }
+}
+
 impl Literal {
     /// Try to coerce this literal, ie. convert it to a more specific type in
     /// regards to the given hint if any.

reid/src/mir/typeinference.rs

@@ -13,8 +13,8 @@ use super::{
     typecheck::ErrorKind,
     typerefs::{ScopeTypeRefs, TypeRef, TypeRefs},
     types::{pick_return, ReturnType},
-    Block, ExprKind, Expression, FunctionDefinition, FunctionDefinitionKind, IfExpression, Module,
-    ReturnKind, StmtKind,
+    Block, ExprKind, Expression, FunctionDefinition, FunctionDefinitionKind, IfExpression,
+    IndexedVariableReference, Module, NamedVariableRef, ReturnKind, StmtKind,
     TypeKind::*,
     VagueType::*,
 };
@@ -113,24 +113,23 @@ impl Block {
                     }
                 }
                 StmtKind::Set(var, expr) => {
-                    todo!("Re-think how set needs to work with arrays")
-                    // // Get the TypeRef for this variable declaration
-                    // let var_ref = inner_hints.find_hint(&var.1);
+                    // Get the TypeRef for this variable declaration
+                    let var_ref = var.find_hint(&inner_hints)?;
 
-                    // // If ok, update the MIR type to this TypeRef
-                    // if let Some((_, var_ref)) = &var_ref {
-                    //     var.0 = var_ref.as_type()
-                    // }
+                    // If ok, update the MIR type to this TypeRef
+                    if let Some((_, var_ref)) = &var_ref {
+                        var.update_type(&var_ref.as_type());
+                    }
 
-                    // // Infer hints for the expression itself
-                    // let inferred = expr.infer_types(&mut state, &inner_hints);
-                    // let expr_ty_ref = state.ok(inferred, expr.1);
+                    // Infer hints for the expression itself
+                    let inferred = expr.infer_types(&mut state, &inner_hints);
+                    let expr_ty_ref = state.ok(inferred, expr.1);
 
-                    // // Try to narrow the variable type declaration with the
-                    // // expression
-                    // if let (Some((_, mut var_ref)), Some(expr_ty_ref)) = (var_ref, expr_ty_ref) {
-                    //     var_ref.narrow(&expr_ty_ref);
-                    // }
+                    // Try to narrow the variable type declaration with the
+                    // expression
+                    if let (Some((_, mut var_ref)), Some(expr_ty_ref)) = (var_ref, expr_ty_ref) {
+                        var_ref.narrow(&expr_ty_ref);
+                    }
                 }
                 StmtKind::Import(_) => todo!(),
                 StmtKind::Expression(expr) => {
@@ -161,6 +160,33 @@ impl Block {
     }
 }
 
+impl IndexedVariableReference {
+    fn find_hint<'s>(
+        &self,
+        hints: &'s ScopeTypeRefs,
+    ) -> Result<Option<(bool, TypeRef<'s>)>, ErrorKind> {
+        match &self.kind {
+            super::IndexedVariableReferenceKind::Named(NamedVariableRef(_, name, _)) => {
+                Ok(hints.find_hint(&name))
+            }
+            super::IndexedVariableReferenceKind::Index(inner, _) => {
+                if let Some((mutable, inner_ref)) = inner.find_hint(hints)? {
+                    let inner_ty = inner_ref.as_type();
+                    match inner_ty {
+                        Array(type_kind, _) => Ok(hints
+                            .from_type(&type_kind)
+                            .clone()
+                            .map(|t_ref| (mutable, t_ref))),
+                        _ => Err(ErrorKind::TriedIndexingNonArray(inner_ty.clone())),
+                    }
+                } else {
+                    Ok(None)
+                }
+            }
+        }
+    }
+}
+
 impl Expression {
     fn infer_types<'s>(
         &mut self,
@@ -263,11 +289,16 @@ impl Expression {
                     ReturnKind::Soft => Ok(block_ref.1),
                 }
             }
-            ExprKind::Index(expression, _) => {
+            ExprKind::Index(expression, index_ty, _) => {
                 let expr_ty = expression.infer_types(state, type_refs)?;
+
                 let kind = unsafe { expr_ty.resolve_type() };
                 match kind {
-                    Array(type_kind, _) => Ok(type_refs.from_type(&type_kind).unwrap()),
+                    Array(type_kind, _) => {
+                        let elem_ty = type_refs.from_type(&type_kind).unwrap();
+                        *index_ty = elem_ty.as_type().clone();
+                        Ok(elem_ty)
+                    }
                     _ => Err(ErrorKind::TriedIndexingNonArray(kind)),
                 }
             }

reid/src/mir/types.rs

@@ -1,6 +1,8 @@
-use crate::ast::VariableReference;
-
-use super::*;
+use super::{
+    pass::{ScopeVariable, Storage},
+    typecheck::ErrorKind,
+    *,
+};
 
 #[derive(Debug, Clone)]
 pub enum ReturnTypeOther {
@@ -89,7 +91,7 @@ impl ReturnType for Expression {
             Block(block) => block.return_type(),
             FunctionCall(fcall) => fcall.return_type(),
             If(expr) => expr.return_type(),
-            Index(expression, _) => {
+            Index(expression, _, _) => {
                 let expr_type = expression.return_type()?;
                 if let (_, TypeKind::Array(elem_ty, _)) = expr_type {
                     Ok((ReturnKind::Soft, *elem_ty))
@@ -162,3 +164,23 @@ pub fn pick_return<T>(lhs: (ReturnKind, T), rhs: (ReturnKind, T)) -> (ReturnKind
         (_, _) => (Soft, lhs.1),
     }
 }
+
+impl IndexedVariableReference {
+    pub fn get_name(&self) -> String {
+        match &self.kind {
+            IndexedVariableReferenceKind::Named(NamedVariableRef(_, name, _)) => name.clone(),
+            IndexedVariableReferenceKind::Index(inner, idx) => {
+                format!("{}[{}]", inner.get_name(), idx)
+            }
+        }
+    }
+
+    pub fn update_type(&mut self, new_ty: &TypeKind) {
+        match &mut self.kind {
+            IndexedVariableReferenceKind::Named(NamedVariableRef(ty, _, _)) => {
+                *ty = new_ty.clone();
+            }
+            IndexedVariableReferenceKind::Index(inner, _) => inner.update_type(new_ty),
+        }
+    }
+}

test.ll

@@ -0,0 +1,38 @@
+; ModuleID = 'test'
+source_filename = "test"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define ptr @array() {
+array:
+  %array_alloca = alloca i16, i16 4, align 2
+  %array_gep = getelementptr i16, ptr %array_alloca, i32 0
+  store i16 10, ptr %array_gep, align 2
+  %array_gep1 = getelementptr i16, ptr %array_alloca, i32 1
+  store i16 15, ptr %array_gep1, align 2
+  %array_gep2 = getelementptr i16, ptr %array_alloca, i32 2
+  store i16 7, ptr %array_gep2, align 2
+  %array_gep3 = getelementptr i16, ptr %array_alloca, i32 3
+  store i16 9, ptr %array_gep3, align 2
+
+  %array_alloca4 = alloca ptr, i16 1, align 8
+  %array_gep5 = getelementptr ptr, ptr %array_alloca4, i32 0
+  store ptr %array_alloca, ptr %array_gep5, align 8
+  ret ptr %array_alloca4
+}
+
+define i16 @main() {
+main:
+  %call = call ptr @array()
+  %array_gep = getelementptr ptr, ptr %call, i32 0
+  %load5 = load ptr, ptr %array_gep2, align 8
+  %array_gep1 = getelementptr i16, ptr %load5, i32 3
+  store i16 5, ptr %array_gep1, align 2
+
+  %array_gep2 = getelementptr ptr, ptr %call, i32 0
+  %load = load ptr, ptr %array_gep2, align 8
+  %array_gep3 = getelementptr i16, ptr %load, i32 3
+  %load4 = load i16, ptr %array_gep3, align 2
+  
+  ret i16 %load4
+}