c-compiler/src/ast.h

#ifndef AST_H
#define AST_H

#include <string>
#include <vector>
#include <memory>

#include "codegen.h"
#include "types.h"
#include "binops.h"
#include "tokens.h"
#include "typechecker.h"

namespace AST {
    class Node {
    public:
        token::Metadata m_meta;

        Node(token::Metadata meta) : m_meta{ meta } {}
        virtual std::string formatted() = 0;
        virtual ~Node() = default;
    };

    class Expression : public Node {
    public:
        Expression(token::Metadata meta) : Node{ meta } {}
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) = 0;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) = 0;
    };

    class Statement : public Node {
    public:
        Statement(token::Metadata meta) : Node{ meta } {}
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) = 0;
    };

    /// @brief Any integer literal
    class IntLiteralExpression : public Expression {
    private:
        int m_value;
        std::shared_ptr<types::Type> m_ty;
    public:
        IntLiteralExpression(token::Metadata meta, int value)
            : Expression{ meta }
            , m_value{ value }
            , m_ty{ { std::shared_ptr<types::Type>{
                new types::FundamentalType{types::FundamentalTypeKind::Int}
            } } } {
        }
        virtual ~IntLiteralExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    /// @brief i.e. "contents"
    class StringLiteralExpression : public Expression {
    private:
        std::string m_value;
    public:
        StringLiteralExpression(token::Metadata meta, std::string value) : Expression{ meta }, m_value{ value } {}
        virtual ~StringLiteralExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    /// @brief Anything that references by value
    class ValueReferenceExpression : public Expression {
    private:
        std::string m_name;
    public:
        ValueReferenceExpression(token::Metadata meta, std::string name) : Expression{ meta }, m_name{ name } {}
        virtual ~ValueReferenceExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    /// @brief Any binary operation (i.e. lhs + rhs)
    class BinaryOperationExpression : public Expression {
    private:
        std::unique_ptr<Expression> m_lhs;
        types::BinOp m_binop;
        std::unique_ptr<Expression> m_rhs;
    public:
        BinaryOperationExpression(
            token::Metadata meta,
            std::unique_ptr<Expression> lhs,
            types::BinOp op,
            std::unique_ptr<Expression> rhs)
            : Expression{ meta }
            , m_lhs{ std::move(lhs) }
            , m_binop{ op }
            , m_rhs{ std::move(rhs) } {
        }
        virtual ~BinaryOperationExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    /// @brief Same as fn()/// @brief 
    class FunctionCallExpression : public Expression {
    private:
        std::unique_ptr<Expression> m_fn_expr;
        std::vector<std::unique_ptr<Expression>> m_args;
    public:
        FunctionCallExpression(
            token::Metadata meta,
            std::unique_ptr<Expression> fn_expr,
            std::vector<std::unique_ptr<Expression>> args)
            : Expression{ meta }
            , m_fn_expr{ std::move(fn_expr) }
            , m_args{ std::move(args) } {
        }
        virtual ~FunctionCallExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    /// @brief Same as (type) value
    class CastExpression : public Expression {
    private:
        std::shared_ptr<types::Type> m_ty;
        std::unique_ptr<Expression> m_expr;
    public:
        CastExpression(
            token::Metadata meta,
            std::shared_ptr<types::Type> type,
            std::unique_ptr<Expression> expr)
            : Expression{ meta }
            , m_ty{ type }
            , m_expr{ std::move(expr) } {
        }
        virtual ~CastExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    /// @brief Same as &value
    class RefExpression : public Expression {
    private:
        std::unique_ptr<Expression> m_expr;
    public:
        RefExpression(
            token::Metadata meta,
            std::unique_ptr<Expression> expr)
            : Expression{ meta }
            , m_expr{ std::move(expr) } {
        }
        virtual ~RefExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    /// @brief Same as *value
    class DerefExpression : public Expression {
    private:
        std::unique_ptr<Expression> m_expr;
    public:
        DerefExpression(
            token::Metadata meta,
            std::unique_ptr<Expression> expr)
            : Expression{ meta }
            , m_expr{ std::move(expr) } {
        }
        virtual ~DerefExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    /// @brief Same as value[num]
    class IndexAccessExpression : public Expression {
    private:
        std::unique_ptr<Expression> m_expr;
        uint32_t m_num;
    public:
        IndexAccessExpression(
            token::Metadata meta,
            std::unique_ptr<Expression> expr,
            uint32_t num)
            : Expression{ meta }
            , m_expr{ std::move(expr) }
            , m_num{ num } {
        }
        virtual ~IndexAccessExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };

    class ListInitializerExpression : public Expression {
    private:
        std::vector<std::unique_ptr<Expression>> m_expressions;
        std::shared_ptr<types::Type> m_ty;
    public:
        ListInitializerExpression(
            token::Metadata meta,
            std::vector<std::unique_ptr<Expression>> expressions,
            std::shared_ptr<types::Type> type)
            : Expression{ meta }
            , m_expressions{ std::move(expressions) }
            , m_ty{ type } {
        }
        virtual ~ListInitializerExpression() override = default;
        virtual std::string formatted() override;
        virtual codegen::StackValue codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> get_codegen_type(codegen::Scope& scope) override;
        virtual std::shared_ptr<types::Type> typecheck(
            typecheck::State& state,
            typecheck::Scope& scope,
            std::optional<std::shared_ptr<types::Type>> expected_ty
        ) override;
    };


    class ReturnStatement : public Statement {
    private:
        std::unique_ptr<Expression> m_expr;
    public:
        ReturnStatement(token::Metadata meta, std::unique_ptr<Expression> expr)
            : Statement{ meta }, m_expr{ std::move(expr) } {
        }
        virtual ~ReturnStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };

    class InitializationStatement : public Statement {
    private:
        std::shared_ptr<types::Type> m_type;
        std::string m_name;
        std::optional<std::unique_ptr<Expression>> m_expr;
    public:
        InitializationStatement(
            token::Metadata meta,
            std::shared_ptr<types::Type> ty,
            std::string name,
            std::optional<std::unique_ptr<Expression>> expr)
            : Statement{ meta }
            , m_type{ std::move(ty) }
            , m_name{ name }
            , m_expr{ std::move(expr) } {
        }
        virtual ~InitializationStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };

    class ExpressionStatement : public Statement {
    private:
        std::unique_ptr<Expression> m_expr;
    public:
        ExpressionStatement(token::Metadata meta, std::unique_ptr<Expression> expr)
            : Statement{ meta }, m_expr{ std::move(expr) } {
        }
        virtual ~ExpressionStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };

    class IfStatement : public Statement {
    private:
        std::unique_ptr<Expression> m_condition;
        std::unique_ptr<Statement> m_then;
        std::optional<std::unique_ptr<Statement>> m_else;
    public:
        IfStatement(token::Metadata meta,
            std::unique_ptr<Expression> condition,
            std::unique_ptr<Statement> then_statement,
            std::optional<std::unique_ptr<Statement>> else_statement)
            : Statement{ meta }
            , m_condition{ std::move(condition) }
            , m_then{ std::move(then_statement) }
            , m_else{ std::move(else_statement) } {
        }
        virtual ~IfStatement() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };

    class TopLevelStatement : public Node {
    public:
        TopLevelStatement(token::Metadata meta) : Node{ meta } {}
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) = 0;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) = 0;
    };

    class Function : public TopLevelStatement {
    private:
        std::shared_ptr<types::Type> m_return_ty;
        std::vector<std::pair<std::optional<std::string>, std::shared_ptr<types::Type>>> m_params;
        bool m_is_vararg;
        std::string m_name;
        std::optional<std::vector<std::unique_ptr<Statement>>> m_statements;
    public:
        Function(
            token::Metadata meta,
            std::shared_ptr<types::Type> return_ty,
            std::vector<std::pair<std::optional<std::string>, std::shared_ptr<types::Type>>> params,
            bool is_vararg,
            std::string name,
            std::optional<std::vector<std::unique_ptr<Statement>>> statements)
            : TopLevelStatement{ meta }
            , m_return_ty{ return_ty }
            , m_params{ params }
            , m_is_vararg{ is_vararg }
            , m_name{ name }
            , m_statements{ std::move(statements) } {
        }
        virtual ~Function() override = default;
        virtual std::string formatted() override;
        virtual void codegen(codegen::Builder& builder, codegen::Scope& scope) override;
        virtual void typecheck(typecheck::State& state, typecheck::Scope& scope) override;
    };
}

#endif