parser.hpp

#ifndef __PARSER_H
#define __PARSER_H

#include <vector>
#include <iostream>
#include "./lexer.hpp"

enum NODE_TYPE{
  NODE_ROOT,
  NODE_VARIABLE,
  NODE_RETURN,
  NODE_PRINT,
  NODE_STRING,
  NODE_INT,
  NODE_BINARY_OP,
  NODE_IF,
  NODE_WHILE,
  NODE_BLOCK,
  NODE_FUNCTION,
  NODE_FUNCTION_CALL,
  NODE_IDENTIFIER
};

struct AST_NODE{
    enum NODE_TYPE TYPE;
    std::string *VALUE;
    AST_NODE *CHILD;
    AST_NODE *LEFT;
    AST_NODE *RIGHT;
    std::vector<AST_NODE *> SUB_STATEMENTS;
    std::vector<AST_NODE *> PARAMETERS;
    std::vector<AST_NODE *> ARGUMENTS;
    
    AST_NODE() : VALUE(nullptr), CHILD(nullptr), LEFT(nullptr), RIGHT(nullptr) {}
};

class Parser{
    public:
        Parser(std::vector<Token *> tokens){
            parserTokens=tokens;
            index=0;
            limit=tokens.size();
            current = (limit > 0) ? parserTokens.at(index) : nullptr;
        }
        
        AST_NODE* parse(){
            AST_NODE *ROOT = new AST_NODE();
            ROOT->TYPE = NODE_ROOT;
            
            while(current && current->TYPE!=TOKEN_EOF){
                switch(current->TYPE){
                    case TOKEN_ID:{
                        ROOT->SUB_STATEMENTS.push_back(parseId());
                        break;
                    }    
                    case TOKEN_IF:{
                        ROOT->SUB_STATEMENTS.push_back(parseIf());
                        break;
                    }
                    case TOKEN_WHILE:{
                        ROOT->SUB_STATEMENTS.push_back(parseWhile());
                        break;
                    }
                    case TOKEN_FUNCTION:{
                        ROOT->SUB_STATEMENTS.push_back(parseFunction());
                        break;
                    }
                    case TOKEN_RETURN:{
                        ROOT->SUB_STATEMENTS.push_back(parseReturn());
                        break;
                    }
                    case TOKEN_PRINT:{
                        ROOT->SUB_STATEMENTS.push_back(parsePrint());
                        break;
                    }
                    default:{
                        std::cout<<"SYNTAX ERROR! Unexpected token: "<<current->VALUE<<" at "<<current->line<<":"<<current->column<<std::endl;
                        exit(1);
                    }
                }
            }
            
            return ROOT;
        }
    private:
        int limit;
        int index;
        Token* current;
        std::vector<Token *> parserTokens;
        
        Token* proceed(enum type TYPE){
            if(!current || TYPE != current->TYPE){
                std::cout<<"SYNTAX ERROR! Expected: "<<TYPE<<" got: "<<(current ? current->TYPE : -1)<<" at "<<(current ? current->line : 0)<<":"<<(current ? current->column : 0)<<std::endl;
                exit(1);
            }
            index++;
            if (index >= limit) {
                current = nullptr;
            } else {
                current = parserTokens.at(index);
            }
            return current;
        }
        
        AST_NODE* parseInt(){
            if(!current || current->TYPE!= TOKEN_INT){
                std::cout<<"SYNTAX ERROR! Expected integer at "<<(current ? current->line : 0)<<":"<<(current ? current->column : 0)<<std::endl;
                exit(1);
            }
            
            AST_NODE* newNode = new AST_NODE();
            newNode->TYPE = NODE_INT;
            newNode->VALUE = &current->VALUE;
            proceed(TOKEN_INT);
            
            return newNode;
        }
        
        AST_NODE* parseExpression(){
            AST_NODE* left = parseTerm();
            
            while(current && (current->TYPE == TOKEN_PLUS || current->TYPE == TOKEN_MINUS)){
                enum type op = current->TYPE;
                std::string* opValue = &current->VALUE;
                proceed(op);
                
                AST_NODE* right = parseTerm();
                
                AST_NODE* binaryOp = new AST_NODE();
                binaryOp->TYPE = NODE_BINARY_OP;
                binaryOp->LEFT = left;
                binaryOp->RIGHT = right;
                binaryOp->VALUE = opValue;
                
                left = binaryOp;
            }
            
            return left;
        }
        
        AST_NODE* parseTerm(){
            AST_NODE* left = parseFactor();
            
            while(current && (current->TYPE == TOKEN_MULTIPLY || current->TYPE == TOKEN_DIVIDE)){
                enum type op = current->TYPE;
                std::string* opValue = &current->VALUE;
                proceed(op);
                
                AST_NODE* right = parseFactor();
                
                AST_NODE* binaryOp = new AST_NODE();
                binaryOp->TYPE = NODE_BINARY_OP;
                binaryOp->LEFT = left;
                binaryOp->RIGHT = right;
                binaryOp->VALUE = opValue;
                
                left = binaryOp;
            }
            
            return left;
        }
        
        AST_NODE* parseFactor(){
            if(!current) {
                std::cout<<"SYNTAX ERROR! Unexpected end of input"<<std::endl;
                exit(1);
            }
            
            switch(current->TYPE){
                case TOKEN_INT:{
                    return parseInt();
                }
                case TOKEN_ID:{
                    std::string* id = &current->VALUE;
                    proceed(TOKEN_ID);
                    
                    // Check if it's a function call
                    if(current && current->TYPE == TOKEN_LEFT_PAREN){
                        return parseFunctionCall(id);
                    }
                    
                    AST_NODE* newNode = new AST_NODE();
                    newNode->TYPE = NODE_IDENTIFIER;
                    newNode->VALUE = id;
                    return newNode;
                }
                case TOKEN_LEFT_PAREN:{
                    proceed(TOKEN_LEFT_PAREN);
                    AST_NODE* expr = parseExpression();
                    proceed(TOKEN_RIGHT_PAREN);
                    return expr;
                }
                default:{
                    std::cout<<"SYNTAX ERROR! Unexpected token in expression: "<<current->VALUE<<" at "<<current->line<<":"<<current->column<<std::endl;
                    exit(1);
                }
            }
        }
        
        AST_NODE* parseFunctionCall(std::string* functionName){
            proceed(TOKEN_LEFT_PAREN);
            
            AST_NODE* callNode = new AST_NODE();
            callNode->TYPE = NODE_FUNCTION_CALL;
            callNode->VALUE = functionName;
            
            // Parse arguments
            while(current && current->TYPE != TOKEN_RIGHT_PAREN){
                callNode->ARGUMENTS.push_back(parseExpression());
                if(current && current->TYPE == TOKEN_COMMA){
                    proceed(TOKEN_COMMA);
                }
            }
            
            proceed(TOKEN_RIGHT_PAREN);
            return callNode;
        }
        
        AST_NODE* parseId(){
            std::string *buffer = &current->VALUE;
            proceed(TOKEN_ID);
            proceed(TOKEN_EQUALS);
            
            AST_NODE *newNode = new AST_NODE();
            newNode->TYPE = NODE_VARIABLE;
            newNode->CHILD = parseExpression();
            newNode->VALUE = buffer;
            
            // Consume semicolon if present
            if(current && current->TYPE == TOKEN_SEMICOLON) {
                proceed(TOKEN_SEMICOLON);
            }
            
            return newNode;
        }
        
        AST_NODE* parseReturn(){
            proceed(TOKEN_RETURN);
            AST_NODE * newNode = new AST_NODE();
            newNode->TYPE = NODE_RETURN;
            newNode->CHILD = parseExpression();
            
            // Consume semicolon if present
            if(current && current->TYPE == TOKEN_SEMICOLON) {
                proceed(TOKEN_SEMICOLON);
            }
            
            return newNode;
        }
        
        AST_NODE* parseString(){
            if(!current || current->TYPE != TOKEN_STRING){
                std::cout<<"Parser error! Expected string at "<<(current ? current->line : 0)<<":"<<(current ? current->column : 0)<<std::endl;
                exit(1);
            }
            
            AST_NODE* newNode = new AST_NODE();
            newNode->TYPE=NODE_STRING;
            newNode->VALUE=&current->VALUE;
            proceed(TOKEN_STRING);
            return newNode;
        }
        
        AST_NODE* parsePrint(){
            proceed(TOKEN_PRINT);
            proceed(TOKEN_LEFT_PAREN);
            proceed(TOKEN_QUOTES);
            
            AST_NODE* newNode = new AST_NODE();
            newNode->TYPE=NODE_PRINT;
            newNode->CHILD=parseString();
            
            proceed(TOKEN_QUOTES);
            proceed(TOKEN_RIGHT_PAREN);
            
            // Consume semicolon if present
            if(current && current->TYPE == TOKEN_SEMICOLON) {
                proceed(TOKEN_SEMICOLON);
            }
            
            return newNode;
        }
        
        AST_NODE* parseIf(){
            proceed(TOKEN_IF);
            proceed(TOKEN_LEFT_PAREN);
            
            AST_NODE* ifNode = new AST_NODE();
            ifNode->TYPE = NODE_IF;
            ifNode->LEFT = parseExpression(); // condition
            
            proceed(TOKEN_RIGHT_PAREN);
            proceed(TOKEN_LEFT_BRACE);
            
            // Parse if block
            while(current && current->TYPE != TOKEN_RIGHT_BRACE){
                ifNode->SUB_STATEMENTS.push_back(parseStatement());
            }
            
            proceed(TOKEN_RIGHT_BRACE);
            
            // Check for else
            if(current && current->TYPE == TOKEN_ELSE){
                proceed(TOKEN_ELSE);
                proceed(TOKEN_LEFT_BRACE);
                
                while(current && current->TYPE != TOKEN_RIGHT_BRACE){
                    ifNode->SUB_STATEMENTS.push_back(parseStatement());
                }
                
                proceed(TOKEN_RIGHT_BRACE);
            }
            
            return ifNode;
        }
        
        AST_NODE* parseWhile(){
            proceed(TOKEN_WHILE);
            proceed(TOKEN_LEFT_PAREN);
            
            AST_NODE* whileNode = new AST_NODE();
            whileNode->TYPE = NODE_WHILE;
            whileNode->LEFT = parseExpression(); // condition
            
            proceed(TOKEN_RIGHT_PAREN);
            proceed(TOKEN_LEFT_BRACE);
            
            // Parse while block
            while(current && current->TYPE != TOKEN_RIGHT_BRACE){
                whileNode->SUB_STATEMENTS.push_back(parseStatement());
            }
            
            proceed(TOKEN_RIGHT_BRACE);
            
            return whileNode;
        }
        
        AST_NODE* parseFunction(){
            proceed(TOKEN_FUNCTION);
            
            AST_NODE* funcNode = new AST_NODE();
            funcNode->TYPE = NODE_FUNCTION;
            funcNode->VALUE = &current->VALUE; // function name
            proceed(TOKEN_ID);
            
            proceed(TOKEN_LEFT_PAREN);
            
            // Parse parameters
            while(current && current->TYPE != TOKEN_RIGHT_PAREN){
                AST_NODE* param = new AST_NODE();
                param->TYPE = NODE_IDENTIFIER;
                param->VALUE = &current->VALUE;
                funcNode->PARAMETERS.push_back(param);
                proceed(TOKEN_ID);
                
                if(current && current->TYPE == TOKEN_COMMA){
                    proceed(TOKEN_COMMA);
                }
            }
            
            proceed(TOKEN_RIGHT_PAREN);
            proceed(TOKEN_LEFT_BRACE);
            
            // Parse function body
            while(current && current->TYPE != TOKEN_RIGHT_BRACE){
                funcNode->SUB_STATEMENTS.push_back(parseStatement());
            }
            
            proceed(TOKEN_RIGHT_BRACE);
            
            return funcNode;
        }
        
        AST_NODE* parseStatement(){
            switch(current->TYPE){
                case TOKEN_ID:{
                    if(peek() && peek()->TYPE == TOKEN_EQUALS){
                        return parseId();
                    } else {
                        return parseFunctionCall(&current->VALUE);
                    }
                }
                case TOKEN_RETURN:{
                    return parseReturn();
                }
                case TOKEN_PRINT:{
                    return parsePrint();
                }
                case TOKEN_IF:{
                    return parseIf();
                }
                case TOKEN_WHILE:{
                    return parseWhile();
                }
                default:{
                    std::cout<<"SYNTAX ERROR! Unexpected statement: "<<current->VALUE<<" at "<<current->line<<":"<<current->column<<std::endl;
                    exit(1);
                }
            }
        }
        
        Token* peek(){
            if(index + 1 < limit){
                return parserTokens.at(index + 1);
            }
            return nullptr;
        }
};

#endif