From 8d50e28ac8d2e59dc1192b109cdc234597524cd4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20Grodzi=C5=84ski?= <rafal.grodzinski@gmail.com>
Date: Thu, 10 Jul 2025 14:28:13 +0900
Subject: [PATCH] Array assignment

---
 src/Compiler/ModuleBuilder.cpp               | 15 ++++++++++-
 src/Parser/Parser.cpp                        | 26 +++++++++++++++++---
 src/Parser/Statement/StatementAssignment.cpp |  8 ++++--
 src/Parser/Statement/StatementAssignment.h   |  4 ++-
 4 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/src/Compiler/ModuleBuilder.cpp b/src/Compiler/ModuleBuilder.cpp
index 0215dda..4e1aaa1 100644
--- a/src/Compiler/ModuleBuilder.cpp
+++ b/src/Compiler/ModuleBuilder.cpp
@@ -159,7 +159,20 @@ void ModuleBuilder::buildAssignment(shared_ptr<StatementAssignment> statement) {
         return;
 
     llvm::Value *value = valueForExpression(statement->getExpression());
-    builder->CreateStore(value, alloca);
+
+    if (statement->getIndexExpression()) {
+        llvm::Value *indexValue = valueForExpression(statement->getIndexExpression());
+        llvm::Value *index[] = {
+            builder->getInt32(0),
+            indexValue
+        };
+        llvm::ArrayType *type = (llvm::ArrayType *)alloca->getAllocatedType();
+        llvm::Value *elementPtr = builder->CreateGEP(type, alloca, index, format("{}[]", statement->getName()));
+
+        builder->CreateStore(value, elementPtr);
+    } else {
+        builder->CreateStore(value, alloca);
+    }
 }
 
 void ModuleBuilder::buildBlock(shared_ptr<StatementBlock> statement) {
diff --git a/src/Parser/Parser.cpp b/src/Parser/Parser.cpp
index d7aecec..c8922ed 100644
--- a/src/Parser/Parser.cpp
+++ b/src/Parser/Parser.cpp
@@ -250,17 +250,35 @@ shared_ptr<Statement> Parser::matchStatementBlock(vector<TokenKind> terminalToke
 }
 
 shared_ptr<Statement> Parser::matchStatementAssignment() {
-    if (!tryMatchingTokenKinds({TokenKind::IDENTIFIER, TokenKind::LEFT_ARROW}, true, false))
+    int startIndex = currentIndex;
+
+    if (!tryMatchingTokenKinds({TokenKind::IDENTIFIER}, true, false))
         return nullptr;
-     
     shared_ptr<Token> identifierToken = tokens.at(currentIndex++);
-    currentIndex++; // arrow
+    shared_ptr<Expression> indexExpression;
+
+    if (tryMatchingTokenKinds({TokenKind::LEFT_SQUARE_BRACKET}, true, true)) {
+        indexExpression = nextExpression();
+        if (indexExpression == nullptr)
+            return nullptr;
+
+        if (!tryMatchingTokenKinds({TokenKind::RIGHT_SQUARE_BRACKET}, true, true)) {
+            markError(TokenKind::RIGHT_SQUARE_BRACKET, {});
+            return nullptr;
+        }
+    }
+
+    // assignment requires left arrow, otherwise abort
+    if (!tryMatchingTokenKinds({TokenKind::LEFT_ARROW}, true, true)) {
+        currentIndex = startIndex;
+        return nullptr;
+    }
 
     shared_ptr<Expression> expression = nextExpression();
     if (expression == nullptr)
         return nullptr;
 
-    return make_shared<StatementAssignment>(identifierToken->getLexme(), expression);
+    return make_shared<StatementAssignment>(identifierToken->getLexme(), indexExpression, expression);
 }
 
 shared_ptr<Statement> Parser::matchStatementReturn() {
diff --git a/src/Parser/Statement/StatementAssignment.cpp b/src/Parser/Statement/StatementAssignment.cpp
index 24fcff6..ae0d24f 100644
--- a/src/Parser/Statement/StatementAssignment.cpp
+++ b/src/Parser/Statement/StatementAssignment.cpp
@@ -1,12 +1,16 @@
 #include "StatementAssignment.h"
 
-StatementAssignment::StatementAssignment(string name, shared_ptr<Expression> expression):
-Statement(StatementKind::ASSIGNMENT), name(name), expression(expression) { }
+StatementAssignment::StatementAssignment(string name, shared_ptr<Expression> indexExpression, shared_ptr<Expression> expression):
+Statement(StatementKind::ASSIGNMENT), name(name), indexExpression(indexExpression), expression(expression) { }
 
 string StatementAssignment::getName() {
     return name;
 }
 
+shared_ptr<Expression> StatementAssignment::getIndexExpression() {
+    return indexExpression;
+}
+
 shared_ptr<Expression> StatementAssignment::getExpression() {
     return expression;
 }
\ No newline at end of file
diff --git a/src/Parser/Statement/StatementAssignment.h b/src/Parser/Statement/StatementAssignment.h
index 0641e3f..0bb7f63 100644
--- a/src/Parser/Statement/StatementAssignment.h
+++ b/src/Parser/Statement/StatementAssignment.h
@@ -5,10 +5,12 @@ class Expression;
 class StatementAssignment: public Statement {
 private:
     string name;
+    shared_ptr<Expression> indexExpression;
     shared_ptr<Expression> expression;
 
 public:
-    StatementAssignment(string name, shared_ptr<Expression> expression);
+    StatementAssignment(string name, shared_ptr<Expression> indexExpressio, shared_ptr<Expression> expression);
     string getName();
+    shared_ptr<Expression> getIndexExpression();
     shared_ptr<Expression> getExpression();
 };
\ No newline at end of file