From aafbf34beaa870a5c691465e7311cb761fb7fcc0 Mon Sep 17 00:00:00 2001
From: MaksimShagov <43129418+MaksimShagov@users.noreply.github.com>
Date: Mon, 29 Apr 2024 14:01:20 +0300
Subject: [PATCH] list accessor in expressions

---
 compiler/include/compiler/ast/node_type.hpp |  1 +
 compiler/lib/ast/node.cpp                   |  3 ++
 compiler/lib/frontend/parser/parser.cpp     | 28 +++++++++++++---
 compiler/tests/frontend/parser.cpp          | 37 +++++++++++++++++++++
 4 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/compiler/include/compiler/ast/node_type.hpp b/compiler/include/compiler/ast/node_type.hpp
index d43c68a7..ad199cdb 100644
--- a/compiler/include/compiler/ast/node_type.hpp
+++ b/compiler/include/compiler/ast/node_type.hpp
@@ -63,6 +63,7 @@ enum class NodeType {
     VariableName,
     WhileStatement,
     ListStatement,
+    ListAccessor,
 };
 
 } // namespace ast
diff --git a/compiler/lib/ast/node.cpp b/compiler/lib/ast/node.cpp
index 26074d7a..87b06535 100644
--- a/compiler/lib/ast/node.cpp
+++ b/compiler/lib/ast/node.cpp
@@ -194,6 +194,9 @@ void Node::dump(std::ostream &stream, int depth) const {
     case NodeType::ListStatement:
         stream << "ListStatement\n";
         break;
+    case NodeType::ListAccessor:
+        stream << "ListAccessor\n";
+        break;
     default:
         stream << "Unknown\n";
     }
diff --git a/compiler/lib/frontend/parser/parser.cpp b/compiler/lib/frontend/parser/parser.cpp
index 84114c0f..73595767 100644
--- a/compiler/lib/frontend/parser/parser.cpp
+++ b/compiler/lib/frontend/parser/parser.cpp
@@ -195,6 +195,10 @@ bool isFunctionCall(const TokenIterator &tokenIter) {
     return tokenIter->type == TokenType::Identifier && std::next(tokenIter)->is(Operator::LeftBrace);
 }
 
+bool isListAccessor(const TokenIterator &tokenIter) {
+    return tokenIter->type == TokenType::Identifier && std::next(tokenIter)->is(Operator::RectLeftBrace);
+}
+
 void buildExpressionSubtree(std::stack<SubExpression> postfixForm, ast::Node::Ptr root, ErrorBuffer &errors) {
     ast::Node::Ptr currNode = root;
     while (!postfixForm.empty()) {
@@ -240,10 +244,11 @@ void buildExpressionSubtree(std::stack<SubExpression> postfixForm, ast::Node::Pt
                 }
             }
         } else {
-            ast::Node::Ptr funcCallNode = std::get<ast::Node::Ptr>(subexpr);
-            assert(funcCallNode->type == ast::NodeType::FunctionCall);
-            funcCallNode->parent = currNode;
-            currNode->children.push_front(funcCallNode);
+            // can be FunctionCall node and list ListAccessor
+            ast::Node::Ptr callNode = std::get<ast::Node::Ptr>(subexpr);
+            assert(callNode->type == ast::NodeType::FunctionCall or callNode->type == ast::NodeType::ListAccessor);
+            callNode->parent = currNode;
+            currNode->children.push_front(callNode);
         }
         while (currNode->children.size() >= getOperandCount(getOperationType(*currNode)))
             currNode = currNode->parent;
@@ -289,6 +294,21 @@ std::stack<SubExpression> generatePostfixForm(TokenIterator tokenIterBegin, Toke
             tokenIter = argsEnd;
             continue;
         }
+        if (isListAccessor(tokenIter)) {
+            ast::Node::Ptr listAccessorNode = std::make_shared<ast::Node>(ast::NodeType::ListAccessor);
+            auto node = ParserContext::pushChildNode(listAccessorNode, ast::NodeType::VariableName, token.ref);
+            node->value = token.id();
+            auto exprBegin = std::next(tokenIter, 2);
+            auto it = exprBegin;
+            while (!it->is(Operator::RectRightBrace))
+                it++;
+            std::stack<SubExpression> argPostfixForm = generatePostfixForm(exprBegin, it, errors);
+            auto exprNode = ParserContext::pushChildNode(listAccessorNode, ast::NodeType::Expression, token.ref);
+            buildExpressionSubtree(argPostfixForm, exprNode, errors);
+            postfixForm.push(listAccessorNode);
+            tokenIter = it++;
+            continue;
+        }
         OperationType opType = getOperationType(token);
         ExpressionTokenType expType = getExpressionTokenType(token);
         if (expType == ExpressionTokenType::Operand) {
diff --git a/compiler/tests/frontend/parser.cpp b/compiler/tests/frontend/parser.cpp
index bbcac643..f18b3299 100644
--- a/compiler/tests/frontend/parser.cpp
+++ b/compiler/tests/frontend/parser.cpp
@@ -545,3 +545,40 @@ TEST(Parser, can_parse_list_defenition) {
                            "                IntegerLiteralValue: 3\n";
     ASSERT_EQ(expected, tree.dump());
 }
+
+TEST(Parser, can_parse_list_access) {
+    StringVec source = {
+        "def main() -> None:",
+        "    mylist : list[int] = [1, 2, 3]",
+        "    x : int = mylist[0]",
+    };
+    TokenList tokens = Lexer::process(source);
+    SyntaxTree tree = Parser::process(tokens);
+    std::string expected = "ProgramRoot\n"
+                           "  FunctionDefinition\n"
+                           "    FunctionName: main\n"
+                           "    FunctionArguments\n"
+                           "    FunctionReturnType: NoneType\n"
+                           "    BranchRoot\n"
+                           "      VariableDeclaration\n"
+                           "        TypeName: ListType\n"
+                           "          TypeName: IntType\n"
+                           "        VariableName: mylist\n"
+                           "        Expression\n"
+                           "          ListStatement\n"
+                           "            Expression\n"
+                           "              IntegerLiteralValue: 1\n"
+                           "            Expression\n"
+                           "              IntegerLiteralValue: 2\n"
+                           "            Expression\n"
+                           "              IntegerLiteralValue: 3\n"
+                           "      VariableDeclaration\n"
+                           "        TypeName: IntType\n"
+                           "        VariableName: x\n"
+                           "        Expression\n"
+                           "          ListAccessor\n"
+                           "            VariableName: mylist\n"
+                           "            Expression\n"
+                           "              IntegerLiteralValue: 0\n";
+    ASSERT_EQ(expected, tree.dump());
+}