forspll/main.cpp

#include <algorithm>
#include <cctype>
#include <iostream>
#include <list>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <string_view>
#include <tuple>

#include <llvm/ADT/APInt.h>
#include <llvm/IR/Constants.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Type.h>
#include <llvm/IR/DerivedTypes.h>

static std::unique_ptr<llvm::LLVMContext> llvmContext;
static std::unique_ptr<llvm::Module> llvmModule;
static std::unique_ptr<llvm::IRBuilder<>> llvmBuilder;
static llvm::Constant *llvmSp;
static llvm::Constant *llvmStack;

struct Var {
    llvm::Value *value;
    bool callable;

    Var(llvm::Value *v = nullptr, bool c = false): value(v), callable(c) {}
};
static std::list<std::map<std::string, Var>> llvmVars;

Var llvmVarGet(const std::string& name, int skip = 0) {
    for (auto sc = llvmVars.rbegin(); sc != llvmVars.rend(); ++sc) {
        if (skip > 0) {
            --skip;
            continue;
        }
        if (sc->contains(name))
            return (*sc)[name];
    }

    return {};
}

struct BaseAST
{
    std::string name;

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

    virtual ~BaseAST() = default;

    virtual llvm::Value *codegen() const { return nullptr; }
};

struct NumberAST : public BaseAST
{
    // push number onto stack

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

    llvm::Value *codegen() const override {
        auto inttype = llvm::Type::getInt32Ty(*llvmContext);
        auto stacktype = llvm::VectorType::get(inttype, 12, false);
        auto dspval = llvmBuilder->CreateLoad(inttype, llvmSp);
        auto one = llvm::ConstantInt::get(inttype, 1);
        auto zero = llvm::ConstantInt::get(inttype, 0);
        auto inc = llvmBuilder->CreateAdd(dspval, one);
        llvmBuilder->CreateStore(inc, llvmSp, false);

        auto dsget = llvmBuilder->CreateGEP(stacktype, llvmStack, {zero, dspval});

        auto val = llvm::ConstantInt::get(*llvmContext, llvm::APInt(32, std::stoi(name), true));
        return llvmBuilder->CreateStore(val, dsget);
    }
};

struct PushAST : public BaseAST
{
    // push named value to stack

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

    llvm::Value *codegen() const override {
        if (auto [var, thunk] = llvmVarGet(name); var) {
            auto inttype = llvm::Type::getInt32Ty(*llvmContext);
            auto stacktype = llvm::VectorType::get(inttype, 12, false);
            auto dspval = llvmBuilder->CreateLoad(inttype, llvmSp);
            auto one = llvm::ConstantInt::get(inttype, 1);
            auto zero = llvm::ConstantInt::get(inttype, 0);
            auto inc = llvmBuilder->CreateAdd(dspval, one);
            llvmBuilder->CreateStore(inc, llvmSp, false);

            auto dsget = llvmBuilder->CreateGEP(stacktype, llvmStack, {zero, dspval});

            if (!thunk)
                var = llvmBuilder->CreateLoad(inttype, var);
            return llvmBuilder->CreateStore(var, dsget);
        } else {
            return nullptr;
        }
    }
};

struct PopAST : public BaseAST
{
    // pop value on stack to named var

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

    llvm::Value *codegen() const override {
        auto inttype = llvm::Type::getInt32Ty(*llvmContext);
        auto stacktype = llvm::VectorType::get(inttype, 12, false);
        auto one = llvm::ConstantInt::get(inttype, 1);
        auto zero = llvm::ConstantInt::get(inttype, 0);
        auto dspval = llvmBuilder->CreateLoad(inttype, llvmSp);
        auto dec = llvmBuilder->CreateSub(dspval, one);
        auto gep = llvmBuilder->CreateGEP(stacktype, llvmStack, {zero, dec});
        llvmBuilder->CreateStore(dec, llvmSp, false);

        auto var = new llvm::GlobalVariable(*llvmModule, inttype, false, llvm::GlobalValue::InternalLinkage, zero, name);
        llvmBuilder->CreateStore(llvmBuilder->CreateLoad(inttype, gep), var, false);

        auto [it, _] = llvmVars.back().emplace(name, var);
        return it->second.value;
    }
};

struct CallAST : public BaseAST
{
    // invoke named invocable

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

    llvm::Value *codegen() const override {
        auto ftype = llvm::FunctionType::get(llvm::Type::getVoidTy(*llvmContext), {}, false);

        if (auto [var, call] = llvmVarGet(name); var) {
            if (call) {
                return llvmBuilder->CreateCall(ftype, var);
            } else {
                auto inttype = llvm::Type::getInt32Ty(*llvmContext);
                auto val = llvmBuilder->CreateLoad(inttype, var);
                auto cast = llvmBuilder->CreateIntToPtr(val, inttype->getPointerTo());
                return llvmBuilder->CreateCall(ftype, cast);
            }
        } else {
            auto func = llvm::Function::Create(ftype, llvm::Function::ExternalLinkage, name, llvmModule.get());
            llvmVars.front().emplace(name, Var {func, true});
            return llvmBuilder->CreateCall(ftype, func);
        }
    }
};

struct ThunkAST : public BaseAST
{
    static int tcount;

    std::list<std::unique_ptr<BaseAST>> body;
    llvm::IRBuilderBase::InsertPoint parent;
    llvm::Function *func;

    explicit ThunkAST(): ThunkAST(std::string("__t") + std::to_string(tcount++)) {}

    explicit ThunkAST(std::string n): BaseAST(n) {
        parent = llvmBuilder->saveIP();
        auto ftype = llvm::FunctionType::get(llvm::Type::getVoidTy(*llvmContext), {}, false);
        func = llvm::Function::Create(ftype, llvm::Function::ExternalLinkage, name.c_str(), llvmModule.get());
        auto BB = llvm::BasicBlock::Create(*llvmContext, "entry", func);
        llvmBuilder->SetInsertPoint(BB);
    }

    llvm::Value *codegen() const override {
        llvmBuilder->CreateRetVoid();
        llvmBuilder->restoreIP(parent);

        return func;
    }
};
int ThunkAST::tcount = 0;

enum class Token {
    none,
    ThunkOpen,
    ThunkClose,
    Quote,
    PopVar,
    PushVar,
    Var,
    Number
};

static std::string name;
static std::list<ThunkAST> scope;

bool isname(char ch) {
    return !isspace(ch) && ch != ')';
}

std::string_view extractName(std::string_view sv)
{
    name.clear();

    while (!sv.empty()) {
        const auto ch = sv.front();

        if (isname(ch)) {
            name += ch;
            sv.remove_prefix(1);
        } else {
            break;
        }
    }

    return sv;
}

std::pair<std::string_view, Token> nextToken(std::string_view sv)
{
    if (sv.empty())
        return {sv, Token::none};

    while (std::isspace(sv.front()))
        sv.remove_prefix(1);

    if (sv.empty())
        return {sv, Token::none};

    const auto ch = sv.front();
    if (ch == ';') {
        return {{}, Token::none};
    } else if (ch == '(') {
        return {sv.substr(1), Token::ThunkOpen};
    } else if (ch == ')') {
        return {sv.substr(1), Token::ThunkClose};
    } else if (ch == '\'') {
        return {extractName(sv.substr(1)), Token::Quote};
    } else if (ch == '$') {
        return {extractName(sv.substr(1)), Token::PopVar};
    } else if (ch == '^') {
        return {extractName(sv.substr(1)), Token::PushVar};
    } else if (isdigit(ch) || (ch == '-' && sv.size() > 1 && isdigit(sv[1]))) {
        return {extractName(sv), Token::Number};
    } else if (isname(ch)) {
        return {extractName(sv), Token::Var};
    }

    return {sv, Token::none};
}

void printToken(Token tok)
{
    switch (tok) {
    case Token::ThunkOpen:
        std::cout << "ThunkOpen ";
        break;
    case Token::ThunkClose:
        std::cout << "ThunkClose ";
        break;
    case Token::Quote:
        std::cout << "Quote ";
        break;
    case Token::PopVar:
        std::cout << "PopVar ";
        break;
    case Token::PushVar:
        std::cout << "PushVar ";
        break;
    case Token::Var:
        std::cout << "Var ";
        break;
    case Token::Number:
        std::cout << "Number ";
        break;
    case Token::none:
        //std::cout << "none ";
        break;
    }
}

bool parseString(std::string_view sv)
{
    do {
        const auto [nsv, tok] = nextToken(sv);

        //printToken(tok);

        if (tok == Token::none && !nsv.empty()) {
            std::cerr << "unknown " << nsv << std::endl;
            break;
        } else {
            std::unique_ptr<BaseAST> expr;

            switch (tok) {
            case Token::ThunkOpen:
                if (scope.empty())
                    scope.emplace_back("main");
                else
                    scope.emplace_back();
                llvmVars.emplace_back();
                break;
            case Token::ThunkClose:
                {
                auto& thunk = scope.back();
                auto gen = thunk.codegen();
                llvmVars.pop_back();
                llvmVars.back().emplace(thunk.name, Var {gen, true});
                expr.reset(new PushAST {thunk.name});
                scope.pop_back();
                }
                break;
            case Token::Quote:
                break;
            case Token::PopVar:
                expr.reset(new PopAST {name});
                break;
            case Token::PushVar:
                expr.reset(new PushAST {name});
                break;
            case Token::Var:
                expr.reset(new CallAST {name});
                break;
            case Token::Number:
                expr.reset(new NumberAST {name});
                break;
            case Token::none:
                break;
            }

            if (expr && !scope.empty()) {
                expr->codegen();
                //scope.back().body.emplace_back().swap(expr);
            }
        }

        sv = nsv;
    } while (!sv.empty());

    return true;
}

int main()
{
    llvmContext = std::make_unique<llvm::LLVMContext>();
    llvmModule = std::make_unique<llvm::Module>("forsp", *llvmContext);
    llvmBuilder = std::make_unique<llvm::IRBuilder<>>(*llvmContext);
    auto inttype = llvm::Type::getInt32Ty(*llvmContext);
    auto stacktype = llvm::VectorType::get(inttype, 12, false);
    auto zero = llvm::ConstantInt::get(inttype, 0);
    auto zerovec = llvm::ConstantVector::get(llvm::ArrayRef(dynamic_cast<llvm::Constant *>(zero)));
    llvmStack = new llvm::GlobalVariable(*llvmModule, stacktype, false, llvm::GlobalValue::ExternalLinkage, zerovec, "stack");
    llvmSp = new llvm::GlobalVariable(*llvmModule, inttype, false, llvm::GlobalValue::ExternalLinkage, zero, "sp");

    llvmVars.emplace_back();

    std::string line;
    while (std::cin.good()) {
        std::getline(std::cin, line);
        parseString(line);
        //std::cout << std::endl;
    }

    //std::cout << "LLVM:" << std::endl;
    llvmModule->print(llvm::errs(), nullptr);

    std::cout << std::endl;
}