Plugin.cpp

#include <clang/Frontend/CompilerInstance.h>
#include <clang/AST/ASTConsumer.h>
#include <clang/AST/RecursiveASTVisitor.h>
#include <clang/Rewrite/Core/Rewriter.h>
#include <clang/Frontend/FrontendPluginRegistry.h>
#include <clang/Frontend/FrontendAction.h>
#include <clang/Lex/Pragma.h>

#include <sstream>
#include <memory>


struct Module {
    std::string name;
    // char just_a_lot_of_data[64];
};

class ModulePragmaHandler : public clang::PragmaHandler {
public:
    ModulePragmaHandler(Module &module) : PragmaHandler("module"), _module(module)
    {}

    [[clang::no_sanitize("address")]]
    virtual void HandlePragma(clang::Preprocessor &preproc, clang::PragmaIntroducer intro, clang::Token &ft) override
    {
        auto loc = ft.getLocation();
        auto &sm = preproc.getSourceManager();

        llvm::outs() << "Found #pragma lua module at " << loc.printToString(preproc.getSourceManager()) << "\n";

        clang::Token tk = {};
        std::ostringstream directive = {};

        while(tk.isNot(clang::tok::eod)) {
            preproc.Lex(tk);
            if (tk.isNot(clang::tok::eod))
                directive << preproc.getSpelling(tk);
        }

        std::string s = directive.str();
        if (s.empty()) {
            llvm::errs() << "#pragma lua module must be followed with a module name";
            return;
        }

        _module.name = s.substr(1).substr(0, s.length()-2);
        llvm::outs() << "Module name: " << _module.name << "\n";
    }

private:
    Module &_module;
};

class PragmaASTConsumer : public clang::ASTConsumer {
public:
    PragmaASTConsumer(clang::CompilerInstance &comp, Module &module) : _handler(module)
    {
        comp.getPreprocessor().AddPragmaHandler("lua", &_handler);
    }

private:
    ModulePragmaHandler _handler;
};

class ASTConsumer : public clang::ASTConsumer {
public:
    ASTConsumer(clang::CompilerInstance &comp, Module &module)
        : _compiler(comp), _module(module), _pragma_consumer(comp, module)
    {}

private:
    clang::CompilerInstance &_compiler;
    Module &_module;
    PragmaASTConsumer _pragma_consumer;
};

class LPlugin : public clang::PluginASTAction {
public:
    LPlugin() : _module() {}

protected:
        std::unique_ptr<clang::ASTConsumer> CreateASTConsumer(clang::CompilerInstance &comp, llvm::StringRef) override
        {
            return std::make_unique<ASTConsumer>(comp, _module);
        }

        bool ParseArgs(const clang::CompilerInstance &, const std::vector<std::string> &) override
        { return true; }

private:
    Module _module;
};

static clang::FrontendPluginRegistry::Add<LPlugin> X("lua", "Plugin for automating the creation of lua modules");