From 4f8153643be3f38118ba7bdaa40c211794e29ac4 Mon Sep 17 00:00:00 2001
From: sofia <sofia@teascade.net>
Date: Sun, 13 Jul 2025 19:26:20 +0300
Subject: [PATCH] Perform type inference for array and indexing expressions

---
 reid/src/mir/typecheck.rs     |  2 ++
 reid/src/mir/typeinference.rs | 47 ++++++++++++++++++++++++++++++++---
 reid/src/mir/typerefs.rs      |  5 ++++
 3 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/reid/src/mir/typecheck.rs b/reid/src/mir/typecheck.rs
index 5e86726..3c48be6 100644
--- a/reid/src/mir/typecheck.rs
+++ b/reid/src/mir/typecheck.rs
@@ -40,6 +40,8 @@ pub enum ErrorKind {
     TypeNotInferrable(TypeKind),
     #[error("Expected branch type to be {0}, found {1} instead")]
     BranchTypesDiffer(TypeKind, TypeKind),
+    #[error("Attempted to index a non-array type of {0}")]
+    TriedIndexingNonArray(TypeKind),
 }
 
 /// Struct used to implement a type-checking pass that can be performed on the
diff --git a/reid/src/mir/typeinference.rs b/reid/src/mir/typeinference.rs
index bae1414..42ea667 100644
--- a/reid/src/mir/typeinference.rs
+++ b/reid/src/mir/typeinference.rs
@@ -4,7 +4,9 @@
 //! must then be passed through TypeCheck with the same [`TypeRefs`] in order to
 //! place the correct types from the IDs and check that there are no issues.
 
-use std::iter;
+use std::{convert::Infallible, iter};
+
+use crate::util::try_all;
 
 use super::{
     pass::{Pass, PassState, ScopeVariable},
@@ -261,8 +263,47 @@ impl Expression {
                     ReturnKind::Soft => Ok(block_ref.1),
                 }
             }
-            ExprKind::Index(expression, _) => todo!("type inference for index expression"),
-            ExprKind::Array(expressions) => todo!("type inference for array expression"),
+            ExprKind::Index(expression, _) => {
+                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()),
+                    _ => Err(ErrorKind::TriedIndexingNonArray(kind)),
+                }
+            }
+            ExprKind::Array(expressions) => {
+                let mut expr_types_result = try_all(
+                    expressions
+                        .iter_mut()
+                        .map(|e| (*e).infer_types(state, type_refs))
+                        .collect(),
+                );
+                match &mut expr_types_result {
+                    Ok(expr_types) => {
+                        let mut iter = expr_types.iter_mut();
+                        if let Some(first) = iter.next() {
+                            while let Some(other) = iter.next() {
+                                first.narrow(other);
+                            }
+
+                            Ok(type_refs
+                                .from_type(&Array(
+                                    Box::new(first.as_type()),
+                                    expressions.len() as u64,
+                                ))
+                                .unwrap())
+                        } else {
+                            todo!();
+                        }
+                    }
+                    Err(errors) => {
+                        for error in errors {
+                            state.ok::<_, Infallible>(Err(error.clone()), self.1);
+                        }
+                        todo!();
+                    }
+                }
+            }
         }
     }
 }
diff --git a/reid/src/mir/typerefs.rs b/reid/src/mir/typerefs.rs
index 423b21a..cbcd043 100644
--- a/reid/src/mir/typerefs.rs
+++ b/reid/src/mir/typerefs.rs
@@ -140,6 +140,11 @@ impl<'outer> ScopeTypeRefs<'outer> {
                 self.types.type_refs.borrow().get(inner_idx).cloned()?
             }
             TypeKind::Vague(_) => self.types.new(ty),
+            TypeKind::Array(elem_ty, length) => {
+                let elem_ty = self.from_type(elem_ty)?;
+                self.types
+                    .new(&TypeKind::Array(Box::new(elem_ty.as_type()), *length))
+            }
             _ => {
                 if let Some(ty_ref) = self.types.find(ty) {
                     ty_ref