nif_lib.c

#include <stdint.h>
#include <stdio.h>
#include <stdbool.h>
#include <erl_nif.h>
#include <setjmp.h>
#include <ucontext.h>

__thread static jmp_buf ret_buf;
__thread static ERL_NIF_TERM ret_term;
__thread static bool stack_inited = false;
__thread static void *stack;

// These are generated by the compiler.
// They contain functions that enable you to
// call a generated NIF function with a variable
// number of arguments.
int nifrt_callarg_wrappers_num;
void (**nifrt_callarg_wrappers)(void*, ErlNifEnv*, ERL_NIF_TERM*);
void (*nifrt_callarg_wrappers_arr[10])(void*, ErlNifEnv*, ERL_NIF_TERM*);

ErlNifResourceType *fun_ptr_resource_type;
ERL_NIF_TERM nifrt_native_fun_atom;
ERL_NIF_TERM nifrt_nil_atom;
ERL_NIF_TERM nifrt_false_atom;
ERL_NIF_TERM nifrt_true_atom;
ERL_NIF_TERM nifrt_internal_err_data_atom;
ERL_NIF_TERM nifrt_badarith_atom;
ERL_NIF_TERM nifrt_error_atom;

void free_fun_ptr(ErlNifEnv *env, void *obj) {
}

int on_load(ErlNifEnv *env) {
    fun_ptr_resource_type =
        enif_open_resource_type(env, NULL, "niffy_fun_ptr",
                free_fun_ptr, ERL_NIF_RT_CREATE, NULL);
    nifrt_native_fun_atom = enif_make_atom(env, "native_fun");
    nifrt_nil_atom = enif_make_atom(env, "nil");
    nifrt_false_atom = enif_make_atom(env, "false");
    nifrt_true_atom = enif_make_atom(env, "true");
    nifrt_internal_err_data_atom = enif_make_atom(env, "internal_err_data");
    nifrt_badarith_atom = enif_make_atom(env, "badarith");
    nifrt_error_atom = enif_make_atom(env, "error");
    return 0;
}

typedef struct {
    void (*fun)();
    int arity;
} fun_ptr_t;

ERL_NIF_TERM nifrt_make_fun_ptr(ErlNifEnv *env, void *inner_fun, int arity) {
    void *obj = enif_alloc_resource(fun_ptr_resource_type, sizeof(fun_ptr_t));
    fun_ptr_t *fun = (fun_ptr_t *) obj;

    fun->fun = inner_fun;
    fun->arity = arity;

    ERL_NIF_TERM ret = enif_make_resource(env, obj);
    enif_release_resource(obj);
    return ret;
}



void nifrt_unreachable_fail(int loc_id) {
    printf("unreachable!! (%d)", loc_id);
    exit(-1);
}

static void nifrt_ret_ok_cont(ErlNifEnv *env, ERL_NIF_TERM _closure_env, ERL_NIF_TERM result) {
    ret_term = result;
    longjmp(ret_buf, 1);
}
static void nifrt_ret_err_cont(ErlNifEnv *env, ERL_NIF_TERM _closure_env, ERL_NIF_TERM result) {
    ret_term = enif_raise_exception(env, result);
    longjmp(ret_buf, 2);
}

static ERL_NIF_TERM nifrt_make_fun_empty_env(ErlNifEnv *env, void *fun, int arity) {
    return enif_make_tuple3(
            env,
            nifrt_native_fun_atom,
            nifrt_make_fun_ptr(env, fun, arity),
            nifrt_nil_atom
            );
}

static void nifrt_inner_launchpad(ErlNifEnv *env, ERL_NIF_TERM *argv, int argc, void *nif_fun) {
    // Store stack state
    if (!setjmp(ret_buf)) {
        (nifrt_callarg_wrappers_arr[argc])(nif_fun, env, argv);
    }
}

ERL_NIF_TERM nifrt_launchpad(ErlNifEnv *env, ERL_NIF_TERM *argv, int argc, void *nif_fun) {

    // If the stack hasn't been allocated for this thread, do that.
    if (!stack_inited) {
        stack = malloc(1048576);
        stack_inited = true;
    }

    // Reset ret term
    ret_term = 0;

    // Construct args to native NIF function
    // (env, ok_cont, err_cont, args..)
    ERL_NIF_TERM argsi[argc+3];
    argsi[0] = nifrt_nil_atom;
    argsi[1] = nifrt_make_fun_empty_env(env, nifrt_ret_ok_cont, 2);
    argsi[2] = nifrt_make_fun_empty_env(env, nifrt_ret_err_cont, 2);
    for (int i = 0; i < argc; i++) {
        argsi[i+3] = argv[i];
    }

    //for (int i = 0; i < argc+3; i++) {
    //    enif_fprintf(stdout, "arg: %T\n", argsi[i]);
    //}

    // Set up contexts
    ucontext_t ret_ctx;
    ucontext_t exec_ctx;

    getcontext(&exec_ctx);
    exec_ctx.uc_stack.ss_sp = stack;
    exec_ctx.uc_stack.ss_size = 1048576;
    exec_ctx.uc_link = &ret_ctx;
    makecontext(&exec_ctx, (void(*)(void)) &nifrt_inner_launchpad, 4, env, argsi, argc+3, nif_fun);

    swapcontext(&ret_ctx, &exec_ctx);

    return ret_term;
}

void nifrt_return_cont(ErlNifEnv *env, void *closure_env, ERL_NIF_TERM ret) {
    ret_term = ret;
    longjmp(ret_buf, 1);
}

void nifrt_cont_call(ErlNifEnv *env, ERL_NIF_TERM cont, ERL_NIF_TERM ret) {
    int arity;
    const ERL_NIF_TERM *terms;
    int res;
    void *fun_void;

    if (enif_get_tuple(env, cont, &arity, &terms)) {
        if (arity == 3 && terms[0] == nifrt_native_fun_atom) {
            res = enif_get_resource(env, terms[1], 
                    fun_ptr_resource_type, &fun_void);
            if (res) {
                enif_keep_resource(fun_void);
                fun_ptr_t *fun_ptr = (fun_ptr_t *) fun_void;
                if (fun_ptr->arity == 2) {
                    fun_ptr->fun(env, terms[2], ret);
                }
            }
        }
    }

    printf("Cont call fail!\n");
    exit(-1);
}

void (*nifrt_unpack_fun_ptr(ErlNifEnv *env, int i_arity, ERL_NIF_TERM fun))() {
    int arity;
    const ERL_NIF_TERM *terms;
    int res;
    void *fun_void;

    enif_fprintf(stdout, "%T\n", fun);

    if (enif_get_tuple(env, fun, &arity, &terms)) {
        if (arity == 3 && terms[0] == nifrt_native_fun_atom) {
            res = enif_get_resource(env, terms[1], 
                    fun_ptr_resource_type, &fun_void);
            if (res) {
                enif_keep_resource(fun_void);
                fun_ptr_t *fun_ptr = (fun_ptr_t *) fun_void;
                enif_fprintf(stdout, "%d %d\n", fun_ptr->arity, i_arity);
                if (fun_ptr->arity == i_arity) {
                    return fun_ptr->fun;
                }
            }
        }
    }

    printf("Unpack fun ptr fail!\n");
    exit(-1);
}

// ==================================
// ==== Function implementations ====
// ==================================

// Function name mangling scheme is as follows:
// GNIF{module_len}_{module}{name_len}_{name}{arity}_{env_num}_{lambda_num}
//
// Mangling scheme of identifiers is as follows:
// "+" replaced by "_p"
// "_" replaced by "__"

// erlang:'+'/2
void GNIF6_erlang2__p2_n_n(ErlNifEnv* env, ERL_NIF_TERM c_env, 
        ERL_NIF_TERM ok_cont, ERL_NIF_TERM err_cont,
        ERL_NIF_TERM term1, ERL_NIF_TERM term2) {
    int i1, i2;
    ERL_NIF_TERM ret;
    if (!enif_get_int(env, term1, &i1)) {
        ret = enif_make_tuple3(env, 
                nifrt_error_atom, 
                nifrt_badarith_atom, 
                nifrt_internal_err_data_atom
                );
        nifrt_cont_call(env, err_cont, ret);
    }
    if (!enif_get_int(env, term2, &i2)) {
        ret = enif_make_tuple3(env, 
                nifrt_error_atom, 
                nifrt_badarith_atom, 
                nifrt_internal_err_data_atom
                );
        nifrt_cont_call(env, err_cont, ret);
    }
    ret = enif_make_int(env, i1 + i2);
    nifrt_cont_call(env, ok_cont, ret);
}

// erlang:'-'/2
void GNIF6_erlang2__m2_n_n(ErlNifEnv* env, ERL_NIF_TERM c_env, 
        ERL_NIF_TERM ok_cont, ERL_NIF_TERM err_cont,
        ERL_NIF_TERM term1, ERL_NIF_TERM term2) {
    int i1, i2;
    ERL_NIF_TERM ret;
    if (!enif_get_int(env, term1, &i1)) {
        ret = enif_make_tuple3(env, 
                nifrt_error_atom, 
                nifrt_badarith_atom, 
                nifrt_internal_err_data_atom
                );
        nifrt_cont_call(env, err_cont, ret);
    }
    if (!enif_get_int(env, term2, &i2)) {
        ret = enif_make_tuple3(env, 
                nifrt_error_atom, 
                nifrt_badarith_atom, 
                nifrt_internal_err_data_atom
                );
        nifrt_cont_call(env, err_cont, ret);
    }
    ret = enif_make_int(env, i1 - i2);
    nifrt_cont_call(env, ok_cont, ret);
}

// erlang:nif_error/1
void GNIF6_erlang10_nif__error1_n_n(ErlNifEnv *env, ERL_NIF_TERM c_env,
        ERL_NIF_TERM ok_cont, ERL_NIF_TERM err_cont,
        ERL_NIF_TERM arg) {
    ERL_NIF_TERM ret;
        ret = enif_make_tuple3(env, 
                nifrt_error_atom, 
                arg, 
                nifrt_internal_err_data_atom
                );
    nifrt_cont_call(env, err_cont, ret);
}

// erlang:is_map/1
void GNIF6_erlang7_is__map1_n_n(ErlNifEnv *env, ERL_NIF_TERM c_env,
        ERL_NIF_TERM ok_cont, ERL_NIF_TERM err_cont,
        ERL_NIF_TERM arg) {
    //enif_fprintf(stdout, "%T\n", arg);
    if (enif_is_map(env, arg)) {
        nifrt_cont_call(env, ok_cont, nifrt_true_atom);
    } else {
        nifrt_cont_call(env, ok_cont, nifrt_false_atom);
    }
}

// erlang:'>'/2
void GNIF6_erlang2__g2_n_n(ErlNifEnv* env, ERL_NIF_TERM c_env, 
        ERL_NIF_TERM ok_cont, ERL_NIF_TERM err_cont,
        ERL_NIF_TERM term1, ERL_NIF_TERM term2) {
    if (enif_compare(term1, term2) > 0) {
        nifrt_cont_call(env, ok_cont, nifrt_true_atom);
    } else {
        nifrt_cont_call(env, ok_cont, nifrt_false_atom);
    }
}

// erlang:element/2
// erlang:is_list/1
// erlang:=</2
// lists:reverse/1
// maps:fold/3
// erlang:is_atom/1
// Elixir.Enumerable:reduce/3