// Compile me with: (NOTE: LLVM must be compiled with LLVM_USE_PERF=On)
//   gcc -O2 -Wall sample.c $(llvm-config --cflags --ldflags --libs --system-libs) -lstdc++
// Run me with:
//   perf record -k 1 ./a.out
// View perf result with:
//   perf inject -j -i perf.data -o perf.data.jitted && perf report -i perf.data.jitted


#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <llvm-c/Core.h>
#include <llvm-c/Target.h>
#include <llvm-c/Orc.h>
#include <llvm-c/OrcEE.h>
#include <llvm-c/LLJIT.h>
#include <llvm-c/Analysis.h>
#include <llvm-c/Transforms/PassManagerBuilder.h>

static LLVMOrcLLJITRef jit;
static LLVMOrcThreadSafeContextRef tsctx;
static LLVMContextRef ctx;
static LLVMBuilderRef bldr;
static LLVMOrcJITDylibRef jd;
static LLVMPassManagerBuilderRef pmb;
static LLVMPassManagerRef mpm;
static LLVMTypeRef int_ty;
static LLVMValueRef int_0, int_1;

typedef int (*MyFunc)(int n);

static inline void check_error(LLVMErrorRef e)
{
    if (e) {
        char *m = LLVMGetErrorMessage(e);
        puts(m);
        LLVMDisposeErrorMessage(m);
        exit(1);
    }
}
static inline void dump_module(LLVMModuleRef mdl)
{
    char *str = LLVMPrintModuleToString(mdl);
    puts(str);
    LLVMDisposeMessage(str);
}

static MyFunc create_function(const char *name)
{
    LLVMModuleRef mod;
    LLVMValueRef fn, next;
    LLVMTypeRef fn_type;
    LLVMBasicBlockRef entry_bb, recurse_bb, return_bb;
    LLVMOrcJITTargetAddress addr;
    LLVMOrcThreadSafeModuleRef tsm;

    mod = LLVMModuleCreateWithNameInContext(name, ctx);
    fn_type = LLVMFunctionType(int_ty, &int_ty, 1, 0);

    fn = LLVMAddFunction(mod, name, fn_type);
    entry_bb = LLVMAppendBasicBlockInContext(ctx, fn, "entrybb");
    recurse_bb = LLVMAppendBasicBlockInContext(ctx, fn, "recursebb");
    return_bb = LLVMAppendBasicBlockInContext(ctx, fn, "returnbb");

    LLVMPositionBuilderAtEnd(bldr, entry_bb);
    LLVMBuildCondBr(bldr,
        LLVMBuildICmp(bldr,
            LLVMIntEQ,
            LLVMGetParam(fn, 0),
            int_0, ""),
        return_bb,
        recurse_bb);
    
    LLVMPositionBuilderAtEnd(bldr, return_bb);
    LLVMBuildRet(bldr, int_1);

    LLVMPositionBuilderAtEnd(bldr, recurse_bb);
    next = LLVMBuildSub(bldr,
        LLVMGetParam(fn, 0),
        int_1, "");
    LLVMBuildRet(bldr,
        LLVMBuildMul(bldr,
            LLVMBuildCall(bldr, fn, &next, 1, ""),
            LLVMGetParam(fn, 0), ""));

    LLVMVerifyFunction(fn, LLVMAbortProcessAction);

    LLVMRunPassManager(mpm, mod);
    //dump_module(mod);

    tsm = LLVMOrcCreateNewThreadSafeModule(mod, tsctx);
    check_error(LLVMOrcLLJITAddLLVMIRModule(jit, jd, tsm));
    LLVMOrcLLJITLookup(jit, &addr, name);
    
    printf("COMPILE DONE! %s = %p\n", name, (void *)addr);
    return (void *)addr;
}

static LLVMOrcObjectLayerRef create_oll(
    void *Ctx, LLVMOrcExecutionSessionRef ES, const char *Triple)
{
    LLVMOrcObjectLayerRef ol;
    LLVMJITEventListenerRef perf;
    ol = LLVMOrcCreateRTDyldObjectLinkingLayerWithSectionMemoryManager(ES);
    perf = LLVMCreatePerfJITEventListener();
    assert(perf); // LLVM must be compiled with LLVM_USE_PERF=On
    LLVMOrcRTDyldObjectLinkingLayerRegisterJITEventListener(ol, perf);
    return ol;
}

int main()
{
    int i;
    MyFunc factorial, test1, test2;
    LLVMOrcLLJITBuilderRef jb;

    LLVMInitializeNativeTarget();
    LLVMInitializeNativeAsmPrinter();
    LLVMInitializeNativeAsmParser();

    jb = LLVMOrcCreateLLJITBuilder();
    LLVMOrcLLJITBuilderSetObjectLinkingLayerCreator(jb, create_oll, NULL);

    check_error(LLVMOrcCreateLLJIT(&jit, jb));
    tsctx = LLVMOrcCreateNewThreadSafeContext();
    ctx = LLVMOrcThreadSafeContextGetContext(tsctx);
    bldr = LLVMCreateBuilderInContext(ctx);
    jd = LLVMOrcLLJITGetMainJITDylib(jit);
    
    pmb = LLVMPassManagerBuilderCreate();
    mpm = LLVMCreatePassManager();
    LLVMPassManagerBuilderSetOptLevel(pmb, 2);
    LLVMPassManagerBuilderPopulateModulePassManager(pmb, mpm);

    int_ty = LLVMInt32TypeInContext(ctx);
    int_0 = LLVMConstInt(int_ty, 0, 0);
    int_1 = LLVMConstInt(int_ty, 1, 0);

    factorial = create_function("factorial");
    printf("10! = %d\n", factorial(10)); // should be 3628800

    puts("TEST BEGIN!");
    test1 = create_function("test1");
    for (i = 0; i < 10000; i++) {
        test1(1000000);
    }
    test2 = create_function("test2");
    for (i = 0; i < 10000; i++) {
        test1(1000000);
        test2(1000000);
    }
    puts("TEST END!");
    return 0;
}