Parse array assign with multiple levels of indexing

Author: fkettelhoit

src/ast.rs

@@ -273,9 +273,7 @@ pub enum StmtEnum<T> {
     /// Mutable let expression, bind a single variable to an expr.
     LetMut(String, Option<Type>, Expr<T>),
     /// Assignment of a (previously as mutable declared) variable.
-    VarAssign(String, Expr<T>),
-    /// Assignment of an index in a (mutable) array.
-    ArrayAssign(String, Expr<T>, Expr<T>),
+    VarAssign(String, Vec<Accessor<T>>, Expr<T>),
     /// Binds an identifier to each value of an array expr, evaluating the body.
     ForEachLoop(Pattern<T>, Expr<T>, Vec<Stmt<T>>),
     /// Binds an identifier to each joined row of two tables, evaluating the body.
@@ -284,6 +282,19 @@ pub enum StmtEnum<T> {
     Expr(Expr<T>),
 }
 
+/// The different ways a mutable variable can be accessed before assignment.
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum Accessor<T> {
+    /// Accessing an array by index to assing to it, e.g. `array[0] = x;`.
+    ArrayAccess {
+        /// The type of the array being accessed by the index.
+        array_ty: T,
+        /// The expression that indexes the array.
+        index: Expr<T>,
+    },
+}
+
 /// An expression and its location in the source code.
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

src/check.rs

@@ -5,8 +5,8 @@ use std::collections::{HashMap, HashSet};
 
 use crate::{
     ast::{
-        self, ConstDef, ConstExpr, ConstExprEnum, EnumDef, Expr, ExprEnum, Mutability, Op,
-        ParamDef, Pattern, PatternEnum, Stmt, StmtEnum, StructDef, Type, UnaryOp, Variant,
+        self, Accessor, ConstDef, ConstExpr, ConstExprEnum, EnumDef, Expr, ExprEnum, Mutability,
+        Op, ParamDef, Pattern, PatternEnum, Stmt, StmtEnum, StructDef, Type, UnaryOp, Variant,
         VariantExprEnum,
     },
     env::Env,
@@ -682,43 +682,32 @@ impl UntypedStmt {
                 let expr = expr.type_check(top_level_defs, env, fns, defs)?;
                 Ok(Stmt::new(StmtEnum::Expr(expr), meta))
             }
-            ast::StmtEnum::VarAssign(identifier, value) => {
+            ast::StmtEnum::VarAssign(identifier, accessors, value) => {
                 match env.get(identifier) {
-                    Some((Some(ty), Mutability::Mutable)) => {
-                        let mut value = value.type_check(top_level_defs, env, fns, defs)?;
-                        check_type(&mut value, &ty)?;
-                        Ok(Stmt::new(
-                            StmtEnum::VarAssign(identifier.clone(), value),
-                            meta,
-                        ))
-                    }
-                    Some((None, Mutability::Mutable)) => {
-                        // binding does not have a type, must have been caused by a previous error, so
-                        // just ignore the statement here
-                        Err(vec![None])
-                    }
-                    Some((_, Mutability::Immutable)) => Err(vec![Some(TypeError(
-                        TypeErrorEnum::IdentifierNotDeclaredAsMutable(identifier.clone()),
-                        meta,
-                    ))]),
-                    None => Err(vec![Some(TypeError(
-                        TypeErrorEnum::UnknownIdentifier(identifier.clone()),
-                        meta,
-                    ))]),
-                }
-            }
-            ast::StmtEnum::ArrayAssign(identifier, index, value) => {
-                match env.get(identifier) {
-                    Some((Some(array_ty), Mutability::Mutable)) => {
-                        let elem_ty = expect_array_type(&array_ty, meta)?;
-
-                        let mut index = index.type_check(top_level_defs, env, fns, defs)?;
-                        check_or_constrain_unsigned(&mut index, UnsignedNumType::Usize)?;
+                    Some((Some(mut elem_ty), Mutability::Mutable)) => {
+                        let mut typed_accessors = vec![];
+                        for access in accessors {
+                            let typed = match access {
+                                Accessor::ArrayAccess { index, .. } => {
+                                    let array_ty = elem_ty.clone();
+                                    elem_ty = expect_array_type(&elem_ty, meta)?;
 
+                                    let mut index =
+                                        index.type_check(top_level_defs, env, fns, defs)?;
+                                    check_or_constrain_unsigned(
+                                        &mut index,
+                                        UnsignedNumType::Usize,
+                                    )?;
+                                    println!("Accessor for {array_ty} with {index:?}");
+                                    Accessor::ArrayAccess { array_ty, index }
+                                }
+                            };
+                            typed_accessors.push(typed);
+                        }
                         let mut value = value.type_check(top_level_defs, env, fns, defs)?;
                         check_type(&mut value, &elem_ty)?;
                         Ok(Stmt::new(
-                            StmtEnum::ArrayAssign(identifier.clone(), index, value),
+                            StmtEnum::VarAssign(identifier.clone(), typed_accessors, value),
                             meta,
                         ))
                     }

src/parse.rs

@@ -4,8 +4,9 @@ use std::{collections::HashMap, iter::Peekable, vec::IntoIter};
 
 use crate::{
     ast::{
-        ConstDef, ConstExpr, ConstExprEnum, EnumDef, Expr, ExprEnum, FnDef, Op, ParamDef, Pattern,
-        PatternEnum, Program, Stmt, StmtEnum, StructDef, Type, UnaryOp, Variant, VariantExprEnum,
+        Accessor, ConstDef, ConstExpr, ConstExprEnum, EnumDef, Expr, ExprEnum, FnDef, Op, ParamDef,
+        Pattern, PatternEnum, Program, Stmt, StmtEnum, StructDef, Type, UnaryOp, Variant,
+        VariantExprEnum,
     },
     scan::Tokens,
     token::{MetaInfo, SignedNumType, Token, TokenEnum, UnsignedNumType},
@@ -455,7 +456,7 @@ impl Parser {
                     if !self.peek(&TokenEnum::RightBrace) && !self.peek(&TokenEnum::Comma) {
                         self.expect(&TokenEnum::Semicolon)?;
                     }
-                    Stmt::new(StmtEnum::VarAssign(identifier, value), meta)
+                    Stmt::new(StmtEnum::VarAssign(identifier, vec![], value), meta)
                 } else if let Some(Token(next, _)) = self.tokens.peek() {
                     let op = match next {
                         TokenEnum::AddAssign => Some(Op::Add),
@@ -489,7 +490,7 @@ impl Parser {
                             ),
                             meta,
                         );
-                        Stmt::new(StmtEnum::VarAssign(identifier, binary_op), meta)
+                        Stmt::new(StmtEnum::VarAssign(identifier, vec![], binary_op), meta)
                     } else {
                         let expr = Expr::untyped(ExprEnum::Identifier(identifier), meta);
                         if !self.peek(&TokenEnum::RightBrace) && !self.peek(&TokenEnum::Comma) {
@@ -514,7 +515,14 @@ impl Parser {
                         self.expect(&TokenEnum::Semicolon)?;
                     }
                     Stmt::new(
-                        StmtEnum::ArrayAssign(identifier.clone(), *index, value),
+                        StmtEnum::VarAssign(
+                            identifier.clone(),
+                            vec![Accessor::ArrayAccess {
+                                array_ty: (),
+                                index: *index,
+                            }],
+                            value,
+                        ),
                         meta,
                     )
                 } else {