Refactor array handling to use Rc<RefCell<Vec<Value>>> to gain O(1) read and write (#7)

* Refactor array handling to use Rc<RefCell<Vec<Value>>> for nested arrays

* Enhance nested array support by updating index type to Vec<Expression> and adjusting parser and compiler logic for multi-dimensional indexing

Author: shubhiscoding

src/ast/mod.rs

@@ -1,3 +1,5 @@
+use std::{cell::RefCell, rc::Rc};
+
 #[derive(Debug)]
 pub enum Statement {
     Expression(Expression),
@@ -39,7 +41,7 @@ pub enum Statement {
     },
     AssignmentIndex {
         array: String,
-        index: Expression,
+        index: Vec<Expression>,
         value: Expression,
     },
 }
@@ -87,5 +89,5 @@ pub enum Expression {
 pub enum Value {
     Number(i32),
     String(String),
-    Array(Vec<Value>),
+    Array(Rc<RefCell<Vec<Value>>>),
 }

src/compiler/mod.rs

@@ -262,13 +262,14 @@ pub fn compile_statements(
                 instructions.push(Instruction::LoadVar(array.clone()));
 
                 // index is expression
-                compile_expr(instructions, &index);
-
+                for idx in &index {
+                    compile_expr(instructions, idx);
+                    instructions.push(Instruction::LoadIndex);
+                }
+                instructions.pop();
                 // value is expression
                 compile_expr(instructions, &value);
-
                 instructions.push(Instruction::StoreIndex);
-                instructions.push(Instruction::AssignVar(array));
             }
         }
     }

src/parser/mod.rs

@@ -369,28 +369,28 @@ impl Parser {
             }
             Statement::Expression(call_expr)
         } else if let Some(Token::SquareLeft) = self.peek() {
-            let index = self.parse_index(name);
-            match self.advance() {
-                Some(Token::Equals) => {
-                    // It's an array assignment: arr[index] = value;
-                    let value = self.parse_first();
-                    match self.advance() {
-                        Some(Token::Semicolon) => {}
-                        _ => panic!("Expected ';' after array assignment"),
-                    }
-                    match index {
-                        Expression::Index { array, index } => Statement::AssignmentIndex {
-                            array: match *array {
-                                Expression::Identifier(name) => name,
-                                _ => panic!("Expected identifier for array name in assignment"),
-                            },
-                            index: *index,
-                            value,
-                        },
-                        _ => panic!("Expected array indexing expression for array assignment"),
-                    }
+            let mut indices = Vec::new();
+            while Some(&Token::SquareLeft) == self.peek() {
+                self.advance();
+                indices.push(self.parse_first());
+                match self.advance() {
+                    Some(Token::SquareRight) => {}
+                    _ => panic!("Expected ']' after array index"),
                 }
-                _ => Statement::Expression(index), // Just an array access expression
+            }
+            match self.advance() {
+                Some(Token::Equals) => {}
+                _ => panic!("Expected '=' after array index"),
+            }
+            let value = self.parse_first();
+            match self.advance() {
+                Some(Token::Semicolon) => {}
+                _ => panic!("Expected ';' after array assignment"),
+            }
+            Statement::AssignmentIndex {
+                array: name,
+                index: indices,
+                value,
             }
         } else {
             self.parse_assignment(name)

src/vm/mod.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, fmt};
+use std::{cell::RefCell, collections::HashMap, fmt, rc::Rc};
 
 use crate::{
     ast::{BinaryOperation, Value},
@@ -11,7 +11,8 @@ impl fmt::Display for Value {
             Value::Number(n) => write!(f, "{}", n),
             Value::String(s) => write!(f, "{}", s),
             Value::Array(arr) => {
-                let elements = arr
+                let vec = arr.borrow();
+                let elements = vec
                     .iter()
                     .map(|v| format!("{}", v))
                     .collect::<Vec<String>>()
@@ -308,25 +309,26 @@ pub fn execute(program: Program, runtime: &mut Runtime) {
                 continue;
             }
             Instruction::CreateArray(len) => {
-                let mut elements = Vec::new();
+                let elements = Rc::new(RefCell::new(Vec::new()));
                 for _ in 0..len {
                     if let Some(value) = runtime.stack.pop() {
-                        elements.push(value);
+                        elements.borrow_mut().push(value);
                     } else {
                         panic!("Not enough values on stack to create array");
                     }
                 }
-                elements.reverse();
+                elements.borrow_mut().reverse();
                 runtime.stack.push(Value::Array(elements));
             }
             Instruction::LoadIndex => {
                 let index = runtime.pop_or_panic_stack();
                 let array = runtime.pop_or_panic_stack();
                 if let (Value::Array(arr), Value::Number(idx)) = (array, index) {
-                    if idx < 0 || (idx as usize) >= arr.len() {
+                    let vec = arr.borrow();
+                    if idx < 0 || (idx as usize) >= vec.len() {
                         panic!("Array index out of bounds");
-                    }
-                    runtime.stack.push(arr[idx as usize].clone());
+                    } // immutable borrow
+                    runtime.stack.push(vec[idx as usize].clone());
                 } else {
                     panic!("LoadIndex requires an array and a number index");
                 }
@@ -336,12 +338,12 @@ pub fn execute(program: Program, runtime: &mut Runtime) {
                 let index = runtime.pop_or_panic_stack();
                 let array = runtime.pop_or_panic_stack();
 
-                if let (Value::Array(mut arr), Value::Number(idx)) = (array, index) {
-                    if idx < 0 || (idx as usize) >= arr.len() {
+                if let (Value::Array(arr), Value::Number(idx)) = (&array, index) {
+                    let mut vec = arr.borrow_mut();
+                    if idx < 0 || (idx as usize) >= vec.len() {
                         panic!("Array index out of bounds");
                     }
-                    arr[idx as usize] = value;
-                    runtime.stack.push(Value::Array(arr));
+                    vec[idx as usize] = value;
                 } else {
                     panic!("StoreIndex requires an array and a number index");
                 }

tests/integration_tests.rs

@@ -1019,3 +1019,93 @@ fn test_nested_array_print_whole() {
     let out = run_rts("let arr = [[1, 2], [3, 4]];\nprint arr;");
     assert_eq!(out, "[[1, 2], [3, 4]]");
 }
+
+// ========== Nested Array Assignment ==========
+
+#[test]
+fn test_nested_array_write_depth_2() {
+    let code = r#"
+let x = [[1, 2], [3, 4]];
+x[0][1] = 99;
+print x;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "[[1, 99], [3, 4]]");
+}
+
+#[test]
+fn test_nested_array_write_depth_3() {
+    let code = r#"
+let x = [10, [3, [6, 11]], 5];
+x[1][1][0] = 99;
+print x;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "[10, [3, [99, 11]], 5]");
+}
+
+#[test]
+fn test_nested_array_write_preserves_siblings() {
+    let code = r#"
+let x = [[1, 2], [3, 4]];
+x[1][0] = 99;
+print x[0];
+print x[1];
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "[1, 2]\n[99, 4]");
+}
+
+#[test]
+fn test_nested_array_write_then_read() {
+    let code = r#"
+let x = [[10, 20], [30, 40]];
+x[0][0] = 5;
+print x[0][0];
+print x[0][1];
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "5\n20");
+}
+
+#[test]
+fn test_nested_array_write_with_expression_index() {
+    let code = r#"
+let x = [[1, 2], [3, 4]];
+let i = 1;
+x[i][i - 1] = 50;
+print x;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "[[1, 2], [50, 4]]");
+}
+
+#[test]
+fn test_nested_array_write_in_loop() {
+    let code = r#"
+let x = [[0, 0], [0, 0]];
+for (let i = 0; i < 2; i++) {
+    for (let j = 0; j < 2; j++) {
+        x[i][j] = i * 2 + j;
+    }
+}
+print x;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "[[0, 1], [2, 3]]");
+}
+
+#[test]
+fn test_nested_array_write_via_function() {
+    let code = r#"
+function setCell(grid, r, c, val) {
+    grid[r][c] = val;
+    return 0;
+}
+let g = [[1, 2], [3, 4]];
+setCell(g, 0, 1, 99);
+print g;
+"#;
+    let out = run_rts(code);
+    assert_eq!(out, "[[1, 99], [3, 4]]");
+}