#include "ast.hpp"

int ThunkAST::tcount = 0;

NumberAST::NumberAST(const std::string& n): BaseAST(n) {}

llvm::Value *NumberAST::codegen(LLVMState& llvmState) const
{
    auto val = llvmState.createInt(std::stoi(name));
    return llvmState.builder.CreateStore(val, llvmState.createPush());
}

PushAST::PushAST(const std::string& n): BaseAST(n) {}

llvm::Value *PushAST::codegen(LLVMState& llvmState) const
{
    if (auto [var, thunk] = Var::lookup(name); var) {
        auto dsget = llvmState.createPush();
        if (!thunk)
            var = llvmState.builder.CreateLoad(llvmState.inttype, var);

        return llvmState.builder.CreateStore(var, dsget);
    } else {
        return nullptr;
    }
}

PopAST::PopAST(const std::string& n): BaseAST(n) {}

llvm::Value *PopAST::codegen(LLVMState& llvmState) const
{
    auto gep = llvmState.createPop();
    auto var = llvmState.createVariable(name);
    auto load = llvmState.builder.CreateLoad(llvmState.inttype, gep);
    llvmState.builder.CreateStore(load, var, false);

    return Var::addLocal(name, var).value;
}

CallAST::CallAST(const std::string& n): BaseAST(n) {}

llvm::Value *CallAST::codegen(LLVMState& llvmState) const
{
    if (auto [var, call] = Var::lookup(name); var) {
        if (call) {
            return llvmState.builder.CreateCall(llvmState.ftype, var);
        } else {
            auto val = llvmState.builder.CreateLoad(llvmState.inttype, var);
            auto cast = llvmState.builder.CreateIntToPtr(val,
                llvmState.inttype->getPointerTo());
            return llvmState.builder.CreateCall(llvmState.ftype, cast);
        }
    } else {
        auto func = llvmState.createFunction(name);
        Var::addGlobal(name, Var {func, true});
        return llvmState.builder.CreateCall(llvmState.ftype, func);
    }
}

ThunkAST::ThunkAST(LLVMState& llvmState):
    ThunkAST(llvmState, std::string("__t") + std::to_string(tcount++)) {}

ThunkAST::ThunkAST(LLVMState& llvmState, std::string n): BaseAST(n)
{
    parent = llvmState.builder.saveIP();
    func = llvmState.createFunction(name);
    auto BB = llvmState.createEntry(func);
    llvmState.builder.SetInsertPoint(BB);
}

llvm::Value *ThunkAST::codegen(LLVMState& llvmState) const
{
    llvmState.builder.CreateRetVoid();
    llvmState.builder.restoreIP(parent);

    return func;
}