#include "stack_allocator.h"
#include "ast.h"

namespace codegen {
    void StackAllocator::push_alloca(std::shared_ptr<types::Type> ty, std::string name) {
        auto type = ty->codegen(this->m_builder, this->m_structs);
        auto value = this->m_builder.builder->CreateAlloca(type, nullptr, name);
        this->m_stack_ptrs.push_back(std::pair<llvm::Value*, std::shared_ptr<types::Type>>{ value, ty });
    }

    llvm::Value* StackAllocator::pop_alloca(std::shared_ptr<types::Type> ty) {
        assert(this->m_stack_ptrs.size() > 0);
        auto first = this->m_stack_ptrs[0];
        assert(types::types_equal(first.second, ty));
        this->m_stack_ptrs.erase(this->m_stack_ptrs.begin());
        return first.first;
    }
}

namespace AST {
    void IntLiteralExpression::codegen_alloca(codegen::StackAllocator&) {
    }

    void StringLiteralExpression::codegen_alloca(codegen::StackAllocator&) {
    }

    void ValueReferenceExpression::codegen_alloca(codegen::StackAllocator&) {
    }

    void BinaryOperationExpression::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_lhs->codegen_alloca(allocator);
        this->m_rhs->codegen_alloca(allocator);
    }

    void FunctionCallExpression::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_fn_expr->codegen_alloca(allocator);
        for (auto& arg : this->m_args) {
            arg->codegen_alloca(allocator);
        }
    }

    void CastExpression::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_expr->codegen_alloca(allocator);
    }

    void RefExpression::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_expr->codegen_alloca(allocator);
    }

    void DerefExpression::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_expr->codegen_alloca(allocator);
    }

    void IndexAccessExpression::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_expr->codegen_alloca(allocator);
    }

    void FieldAccessExpression::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_expr->codegen_alloca(allocator);
    }

    void ListInitializerExpression::codegen_alloca(codegen::StackAllocator& allocator) {
        allocator.push_alloca(this->m_ty, "list");
        for (auto& expr : this->m_expressions) {
            expr->codegen_alloca(allocator);
        }
    }

    void ReturnStatement::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_expr->codegen_alloca(allocator);
    }

    void InitializationStatement::codegen_alloca(codegen::StackAllocator& allocator) {
        allocator.push_alloca(this->m_type, this->m_name);
        if (this->m_expr) {
            (*this->m_expr)->codegen_alloca(allocator);
        }
    }

    void ExpressionStatement::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_expr->codegen_alloca(allocator);
    }

    void IfStatement::codegen_alloca(codegen::StackAllocator& allocator) {
        this->m_condition->codegen_alloca(allocator);
        this->m_then->codegen_alloca(allocator);
        if (this->m_else) {
            (*this->m_else)->codegen_alloca(allocator);
        }
    }
}