jiffy

// This file is part of Jiffy released under the MIT license.
// See the LICENSE file for more information.

#include <assert.h>
#include <stdlib.h>
#include <string.h>

#include "erl_nif.h"
#include "ffc.h"
#include "jiffy.h"
#include "jiffy_simd.h"
#include "jiffy_utf8.h"

#define STACK_SIZE_INC 64
#define JIFFY_SMALL_ARRAY_SIZE 64

enum {
    st_value=0,
    st_object,
    st_array,
    st_key,
    st_colon,
    st_comma,
    st_done,
    st_invalid
} JsonState;

typedef struct {
    ErlNifEnv*      env;
    jiffy_st*       atoms;

    ERL_NIF_TERM    arg;
    ErlNifBinary    bin;

    size_t          bytes_per_red;
    int             is_partial;
    int             return_maps;
    int             return_trailer;
    int             dedupe_keys;
    int             copy_strings;
    ERL_NIF_TERM    null_term;

    unsigned char*  p;
    size_t          i;
    size_t          len;

    char*           st_data;
    int             st_size;
    int             st_top;
} Decoder;

// Returns an allocated resource or NULL if the term stack
// cannot be allocated
static Decoder*
dec_new(ErlNifEnv* env)
{
    jiffy_st* st = (jiffy_st*) enif_priv_data(env);
    // If enif_alloc_resource cannot allocate it crashes the VM
    Decoder* d = enif_alloc_resource(st->res_dec, sizeof(Decoder));
    assert(d != NULL);
    int i;

    // Zero everything and only set non-0 fields
    memset(d, 0, sizeof(*d));

    d->atoms = st;
    d->bytes_per_red = DEFAULT_BYTES_PER_REDUCTION;
    d->null_term = d->atoms->atom_null;
    d->len = -1;

    d->st_data = (char*) enif_alloc(STACK_SIZE_INC);
    if(d->st_data == NULL) {
        enif_release_resource(d);
        return NULL;
    }

    d->st_size = STACK_SIZE_INC;

    for(i = 0; i < d->st_size; i++) {
        d->st_data[i] = st_invalid;
    }

    d->st_data[0] = st_value;
    d->st_top++;

    return d;
}

static void
dec_init(Decoder* d, ErlNifEnv* env, ERL_NIF_TERM arg, ErlNifBinary* bin)
{
    d->env = env;
    d->arg = arg;

    d->p = bin->data;
    d->len = bin->size;
}

void
dec_destroy(ErlNifEnv* env, void* obj)
{
    Decoder* d = (Decoder*) obj;

    if(d->st_data != NULL) {
        enif_free(d->st_data);
    }
}

static ERL_NIF_TERM
dec_error(Decoder* d, const char* atom)
{
    ERL_NIF_TERM pos = enif_make_int(d->env, d->i+1);
    ERL_NIF_TERM msg = make_atom(d->env, atom);
    ERL_NIF_TERM ret = enif_make_tuple2(d->env, pos, msg);
    return enif_make_tuple2(d->env, d->atoms->atom_error, ret);
}

static inline char
dec_curr(Decoder* d)
{
    assert(d->st_top > 0);
    return d->st_data[d->st_top - 1];
}

static inline int
dec_top(Decoder* d)
{
    return d->st_top;
}

static inline void
dec_push(Decoder* d, char val)
{
    int new_sz;
    int i;

    if(d->st_top == d->st_size) {
        new_sz = d->st_size + STACK_SIZE_INC;
        d->st_data = (char*) enif_realloc(d->st_data, new_sz);
        d->st_size = new_sz;
        for(i = d->st_top; i < d->st_size; i++) {
            d->st_data[i] = st_invalid;
        }
    }

    assert(d->st_top < d->st_size);
    d->st_data[d->st_top++] = val;
}

static inline char
dec_pop(Decoder* d) {
    char current = st_invalid;

    if (d->st_top > 0) {
        current = d->st_data[d->st_top - 1];
        d->st_data[d->st_top - 1] = st_invalid;
        d->st_top--;
    }

    return current;
}

static void
dec_pop_assert(Decoder* d, char val)
{
    char current = dec_pop(d);
    assert(current == val && "popped invalid state.");
    (void)current;
}

static int
dec_string(Decoder* d, ERL_NIF_TERM* value)
{
    int has_escape = 0;
    int num_escapes = 0;
    int st;
    size_t ulen;
    size_t ui;
    int hi;
    int lo;
    char* chrbuf;
    int chrpos;

    if(d->p[d->i] != '\"') {
        return 0;
    }
    d->i++;

    st = d->i;

    while(d->i < d->len) {
        if(d->p[d->i] < 0x20) {
            return 0;
        } else if(d->p[d->i] == '\"') {
            d->i++;
            goto parse;
        } else if(d->p[d->i] == '\\') {
            if(d->i+1 >= d->len) {
                return 0;
            }
            has_escape = 1;
            num_escapes += 1;
            d->i++;
            switch(d->p[d->i]) {
                case '\"':
                case '\\':
                case '/':
                case 'b':
                case 'f':
                case 'n':
                case 'r':
                case 't':
                    d->i++;
                    break;
                case 'u':
                    hi = 0;
                    lo = 0;
                    d->i++;
                    if(d->i + 4 >= d->len) {
                        return 0;
                    }
                    hi = int_from_hex(&(d->p[d->i]));
                    if(hi < 0) {
                        return 0;
                    }
                    d->i += 4;
                    if(hi >= 0xD800 && hi < 0xDC00) {
                        if(d->i + 6 >= d->len) {
                            return 0;
                        }
                        if(d->p[d->i++] != '\\') {
                            return 0;
                        } else if(d->p[d->i++] != 'u') {
                            return 0;
                        }
                        lo = int_from_hex(&(d->p[d->i]));
                        if(lo < 0) {
                            return 0;
                        }
                        hi = unicode_from_pair(hi, lo);
                        if(hi < 0) {
                            return 0;
                        }
                    }
                    hi = utf8_len(hi);
                    if(hi < 0) {
                        return 0;
                    }
                    if(lo == 0) {
                        num_escapes += 5 - hi;
                    } else {
                        num_escapes += 11 - hi;
                    }
                    break;
                default:
                    return 0;
            }
        } else if(d->p[d->i] < 0x80) {
            // Scan ahead plain ASCII as an optimization. The first
            // byte has already been checked, so start at i+1.
            d->i = jiffy_scan_ascii_string_body(d->p, d->len, d->i + 1);
        } else {
            ulen = utf8_validate(&(d->p[d->i]), d->len - d->i);
            if(ulen == 0) {
                return 0;
            }
            d->i += ulen;
        }
    }

    // The goto above ensures that we only
    // hit this when a string is not terminated
    // correctly.
    return 0;

parse:
    if(!has_escape && !d->copy_strings) {
        *value = enif_make_sub_binary(d->env, d->arg, st, (d->i - st - 1));
        return 1;
    } else if(!has_escape) {
        ulen = d->i - 1 - st;
        chrbuf = (char*) enif_make_new_binary(d->env, ulen, value),
        memcpy(chrbuf, &(d->p[st]), ulen);
        return 1;
    }

    hi = 0;
    lo = 0;

    ulen = (d->i - 1) - st - num_escapes;
    chrbuf = (char*) enif_make_new_binary(d->env, ulen, value);
    chrpos = 0;
    ui = st;
    while(ui < d->i - 1) {
        if(d->p[ui] != '\\') {
            chrbuf[chrpos++] = d->p[ui++];
            continue;
        }
        ui++;
        switch(d->p[ui]) {
            case '\"':
            case '\\':
            case '/':
                chrbuf[chrpos++] = d->p[ui];
                ui++;
                break;
            case 'b':
                chrbuf[chrpos++] = '\b';
                ui++;
                break;
            case 'f':
                chrbuf[chrpos++] = '\f';
                ui++;
                break;
            case 'n':
                chrbuf[chrpos++] = '\n';
                ui++;
                break;
            case 'r':
                chrbuf[chrpos++] = '\r';
                ui++;
                break;
            case 't':
                chrbuf[chrpos++] = '\t';
                ui++;
                break;
            case 'u':
                ui++;
                hi = int_from_hex(&(d->p[ui]));
                if(hi < 0) {
                    return 0;
                }
                if(hi >= 0xD800 && hi < 0xDC00) {
                    lo = int_from_hex(&(d->p[ui+6]));
                    if(lo < 0) {
                        return 0;
                    }
                    hi = unicode_from_pair(hi, lo);
                    ui += 10;
                } else {
                    ui += 4;
                }
                hi = unicode_to_utf8(hi, (unsigned char*) chrbuf+chrpos);
                if(hi < 0) {
                    return 0;
                }
                chrpos += hi;
                break;
            default:
                return 0;
        }
    }

    return 1;
}

static int
dec_number(Decoder* d, ERL_NIF_TERM* value)
{
    // ffc validates, parses, and picks int-vs-double in a single call
    const unsigned char* JIFFY_RESTRICT p = d->p;
    const size_t start = d->i;
    const char* nstart = (const char*)&p[start];
    const char* nend_max = (const char*)&p[d->len];

    ffc_json_number jn;
    ffc_result r = ffc_parse_json_number(nstart, nend_max, &jn);

    // After parsing r.ptr point to where parsing stops:
    //   OK - first byte past the number
    //   OUT_OF_RANGE - first byte past the number (same span, just doesn't fit)
    //   INVALID - the offending byte
    d->i = start + (size_t)(r.ptr - nstart);

    if(r.outcome == FFC_OUTCOME_INVALID_INPUT) {
        return 0;
    }

    if(r.outcome == FFC_OUTCOME_OK) {
        if(jn.kind == FFC_JSON_NUM_KIND_INT64) {
            *value = enif_make_int64(d->env, jn.value.i64);
        } else {
            *value = enif_make_double(d->env, jn.value.f64);
        }
        return 1;
    }

    ERL_NIF_TERM num_type = (jn.kind == FFC_JSON_NUM_KIND_INT64)
        ? d->atoms->atom_bignum
        : d->atoms->atom_bigdbl;

    d->is_partial = 1;
    const size_t num_len = (size_t)(r.ptr - nstart);
    *value = enif_make_sub_binary(d->env, d->arg, start, num_len);
    *value = enif_make_tuple2(d->env, num_type, *value);
    return 1;
}

static ERL_NIF_TERM
make_empty_object(ErlNifEnv* env, int ret_map)
{
    if(ret_map) {
        return enif_make_new_map(env);
    }

    return enif_make_tuple1(env, enif_make_list(env, 0));
}

static inline ERL_NIF_TERM
make_array(ErlNifEnv* env, ERL_NIF_TERM list)
{
    ERL_NIF_TERM item;

    unsigned int count = 0;
    enif_get_list_length(env, list, &count);

    if(count == 0) {
        return enif_make_list(env, 0);
    }

    ERL_NIF_TERM small_buf[JIFFY_SMALL_ARRAY_SIZE];
    ERL_NIF_TERM* arr = (count <= JIFFY_SMALL_ARRAY_SIZE)
        ? small_buf
        : enif_alloc(count * sizeof(ERL_NIF_TERM));

    // Fill array in reverse since list was reversed after parsing
    unsigned int i = count;
    while(enif_get_list_cell(env, list, &item, &list)) {
        arr[--i] = item;
    }

    ERL_NIF_TERM ret = enif_make_list_from_array(env, arr, count);

    if(arr != small_buf) {
        enif_free(arr);
    }

    return ret;
}

// Scan ahead and look for whitespace. This is used when we saw at least one
// whitespace character and then expect to find more.
static inline size_t
skip_whitespace(const unsigned char* JIFFY_RESTRICT p, size_t len, size_t i)
{
    while(i < len && (
            p[i] == ' ' ||
            p[i] == '\n' ||
            p[i] == '\r' ||
            p[i] == '\t')) {
        i++;
    }
    return i;
}

ERL_NIF_TERM
decode_init(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
    Decoder* d;
    jiffy_st* st = (jiffy_st*) enif_priv_data(env);
    ERL_NIF_TERM tmp_argv[5];
    ERL_NIF_TERM opts;
    ERL_NIF_TERM val;

    if(argc != 2) {
        return enif_make_badarg(env);
    }

    d = dec_new(env);
    if(d == NULL) {
        return enif_make_tuple2(env, st->atom_error,
                make_atom(env, "internal_error"));
    }

    tmp_argv[0] = argv[0];
    tmp_argv[1] = enif_make_resource(env, d);
    tmp_argv[2] = st->atom_error;
    tmp_argv[3] = enif_make_list(env, 0);
    tmp_argv[4] = enif_make_list(env, 0);

    enif_release_resource(d);

    opts = argv[1];
    if(!enif_is_list(env, opts)) {
        return enif_make_badarg(env);
    }

    while(enif_get_list_cell(env, opts, &val, &opts)) {
        if(get_bytes_per_iter(env, val, &(d->bytes_per_red))) {
            continue;
        } else if(get_bytes_per_red(env, val, &(d->bytes_per_red))) {
            continue;
        } else if(enif_is_identical(val, d->atoms->atom_return_maps)) {
            d->return_maps = 1;
        } else if(enif_is_identical(val, d->atoms->atom_return_trailer)) {
            d->return_trailer = 1;
        } else if(enif_is_identical(val, d->atoms->atom_dedupe_keys)) {
            d->dedupe_keys = 1;
        } else if(enif_is_identical(val, d->atoms->atom_copy_strings)) {
            d->copy_strings = 1;
        } else if(enif_is_identical(val, d->atoms->atom_use_nil)) {
            d->null_term = d->atoms->atom_nil;
        } else if(get_null_term(env, val, &(d->null_term))) {
            continue;
        } else {
            return enif_make_badarg(env);
        }
    }

    return decode_iter(env, 5, tmp_argv);
}

ERL_NIF_TERM
decode_iter(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
    Decoder* d;
    jiffy_st* st = (jiffy_st*) enif_priv_data(env);

    ErlNifBinary bin;

    ERL_NIF_TERM objs;
    ERL_NIF_TERM curr;
    ERL_NIF_TERM val = argv[2];
    ERL_NIF_TERM trailer;
    ERL_NIF_TERM ret;
    ERL_NIF_TERM tmp_argv[5];

    void* res;

    size_t start;
    size_t bytes_processed = 0;

    if(!enif_inspect_binary(env, argv[0], &bin)) {
        return enif_make_badarg(env);
    } else if(!enif_get_resource(env, argv[1], st->res_dec, &res)) {
        return enif_make_badarg(env);
    }

    d = (Decoder*) res;

    dec_init(d, env, argv[0], &bin);
    objs = argv[3];
    curr = argv[4];

    start = d->i;

    const size_t yt = yield_threshold(d->bytes_per_red);
    const unsigned char* JIFFY_RESTRICT p = d->p;
    const size_t len = d->len;

    while(d->i < bin.size) {
        bytes_processed = d->i - start;

        if(bytes_processed >= yt) {
            assert(enif_is_list(env, objs));
            assert(enif_is_list(env, curr));

            tmp_argv[0] = argv[0];
            tmp_argv[1] = argv[1];
            tmp_argv[2] = val;
            tmp_argv[3] = objs;
            tmp_argv[4] = curr;

            bump_used_reds(env, bytes_processed, d->bytes_per_red);

            return enif_schedule_nif(
                    env,
                    "nif_decode_iter",
                    0,
                    decode_iter,
                    5,
                    tmp_argv
                );
        }

        switch(dec_curr(d)) {
            case st_value:
                switch(d->p[d->i]) {
                    case ' ':
                    case '\n':
                    case '\r':
                    case '\t':
                        d->i = skip_whitespace(p, len, d->i + 1);
                        break;
                    case 'n':
                        if(d->i + 3 >= d->len) {
                            ret = dec_error(d, "invalid_literal");
                            goto done;
                        }
                        if(memcmp(&(d->p[d->i]), "null", 4) != 0) {
                            ret = dec_error(d, "invalid_literal");
                            goto done;
                        }
                        val = d->null_term;
                        dec_pop_assert(d, st_value);
                        d->i += 4;
                        break;
                    case 't':
                        if(d->i + 3 >= d->len) {
                            ret = dec_error(d, "invalid_literal");
                            goto done;
                        }
                        if(memcmp(&(d->p[d->i]), "true", 4) != 0) {
                            ret = dec_error(d, "invalid_literal");
                            goto done;
                        }
                        val = d->atoms->atom_true;
                        dec_pop_assert(d, st_value);
                        d->i += 4;
                        break;
                    case 'f':
                        if(d->i + 4 >= bin.size) {
                            ret = dec_error(d, "invalid_literal");
                            goto done;
                        }
                        if(memcmp(&(d->p[d->i]), "false", 5) != 0) {
                            ret = dec_error(d, "invalid_literal");
                            goto done;
                        }
                        val = d->atoms->atom_false;
                        dec_pop_assert(d, st_value);
                        d->i += 5;
                        break;
                    case '\"':
                        if(!dec_string(d, &val)) {
                            ret = dec_error(d, "invalid_string");
                            goto done;
                        }
                        dec_pop_assert(d, st_value);
                        break;
                    case '-':
                    case '0':
                    case '1':
                    case '2':
                    case '3':
                    case '4':
                    case '5':
                    case '6':
                    case '7':
                    case '8':
                    case '9':
                        if(!dec_number(d, &val)) {
                            ret = dec_error(d, "invalid_number");
                            goto done;
                        }
                        dec_pop_assert(d, st_value);
                        break;
                    case '{':
                        dec_push(d, st_object);
                        dec_push(d, st_key);
                        objs = enif_make_list_cell(env, curr, objs);
                        curr = enif_make_list(env, 0);
                        d->i++;
                        break;
                    case '[':
                        dec_push(d, st_array);
                        dec_push(d, st_value);
                        objs = enif_make_list_cell(env, curr, objs);
                        curr = enif_make_list(env, 0);
                        d->i++;
                        break;
                    case ']':
                        if(!enif_is_empty_list(env, curr)) {
                            ret = dec_error(d, "invalid_json");
                            goto done;
                        }
                        dec_pop_assert(d, st_value);
                        if(dec_pop(d) != st_array) {
                            ret = dec_error(d, "invalid_json");
                            goto done;
                        }
                        dec_pop_assert(d, st_value);
                        val = curr; // curr is []
                        if(!enif_get_list_cell(env, objs, &curr, &objs)) {
                            ret = dec_error(d, "internal_error");
                            goto done;
                        }
                        d->i++;
                        break;
                    default:
                        ret = dec_error(d, "invalid_json");
                        goto done;
                }
                if(dec_top(d) == 0) {
                    dec_push(d, st_done);
                } else if(dec_curr(d) != st_value && dec_curr(d) != st_key) {
                    dec_push(d, st_comma);
                    curr = enif_make_list_cell(env, val, curr);
                }
                break;

            case st_key:
                switch(d->p[d->i]) {
                    case ' ':
                    case '\n':
                    case '\r':
                    case '\t':
                        d->i = skip_whitespace(p, len, d->i + 1);
                        break;
                    case '\"':
                        if(!dec_string(d, &val)) {
                            ret = dec_error(d, "invalid_string");
                            goto done;
                        }
                        dec_pop_assert(d, st_key);
                        dec_push(d, st_colon);
                        curr = enif_make_list_cell(env, val, curr);
                        break;
                    case '}':
                        if(!enif_is_empty_list(env, curr)) {
                            ret = dec_error(d, "invalid_json");
                            goto done;
                        }
                        dec_pop_assert(d, st_key);
                        dec_pop_assert(d, st_object);
                        dec_pop_assert(d, st_value);
                        val = make_empty_object(env, d->return_maps);
                        if(!enif_get_list_cell(env, objs, &curr, &objs)) {
                            ret = dec_error(d, "internal_error");
                            goto done;
                        }
                        if(dec_top(d) == 0) {
                            dec_push(d, st_done);
                        } else {
                            dec_push(d, st_comma);
                            curr = enif_make_list_cell(env, val, curr);
                        }
                        d->i++;
                        break;
                    default:
                        ret = dec_error(d, "invalid_json");
                        goto done;
                }
                break;

            case st_colon:
                switch(d->p[d->i]) {
                    case ' ':
                    case '\n':
                    case '\r':
                    case '\t':
                        d->i = skip_whitespace(p, len, d->i + 1);
                        break;
                    case ':':
                        dec_pop_assert(d, st_colon);
                        dec_push(d, st_value);
                        d->i++;
                        break;
                    default:
                        ret = dec_error(d, "invalid_json");
                        goto done;
                }
                break;

            case st_comma:
                switch(d->p[d->i]) {
                    case ' ':
                    case '\n':
                    case '\r':
                    case '\t':
                        d->i = skip_whitespace(p, len, d->i + 1);
                        break;
                    case ',':
                        dec_pop_assert(d, st_comma);
                        switch(dec_curr(d)) {
                            case st_object:
                                dec_push(d, st_key);
                                break;
                            case st_array:
                                dec_push(d, st_value);
                                break;
                            default:
                                ret = dec_error(d, "internal_error");
                                goto done;
                        }
                        d->i++;
                        break;
                    case '}':
                        dec_pop_assert(d, st_comma);
                        if(dec_pop(d) != st_object) {
                            ret = dec_error(d, "invalid_json");
                            goto done;
                        }
                        dec_pop_assert(d, st_value);
                        if(!make_object(env, curr, &val,
                                d->return_maps, d->dedupe_keys)) {
                            ret = dec_error(d, "internal_object_error");
                            goto done;
                        }
                        if(!enif_get_list_cell(env, objs, &curr, &objs)) {
                            ret = dec_error(d, "internal_error");
                            goto done;
                        }
                        if(dec_top(d) > 0) {
                            dec_push(d, st_comma);
                            curr = enif_make_list_cell(env, val, curr);
                        } else {
                            dec_push(d, st_done);
                        }
                        d->i++;
                        break;
                    case ']':
                        dec_pop_assert(d, st_comma);
                        if(dec_pop(d) != st_array) {
                            ret = dec_error(d, "invalid_json");
                            goto done;
                        }
                        dec_pop_assert(d, st_value);
                        val = make_array(env, curr);
                        if(!enif_get_list_cell(env, objs, &curr, &objs)) {
                            ret = dec_error(d, "internal_error");
                            goto done;
                        }
                        if(dec_top(d) > 0) {
                            dec_push(d, st_comma);
                            curr = enif_make_list_cell(env, val, curr);
                        } else {
                            dec_push(d, st_done);
                        }
                        d->i++;
                        break;
                    default:
                        ret = dec_error(d, "invalid_json");
                        goto done;
                }
                break;

            case st_done:
                switch(d->p[d->i]) {
                    case ' ':
                    case '\n':
                    case '\r':
                    case '\t':
                        d->i = skip_whitespace(p, len, d->i + 1);
                        break;
                    default:
                        goto decode_done;
                }
                break;

            default:
                ret = dec_error(d, "invalid_internal_state");
                goto done;
        }
    }

decode_done:

    if(d->i < bin.size && d->return_trailer) {
        trailer = enif_make_sub_binary(env, argv[0], d->i, bin.size - d->i);
        val = enif_make_tuple3(env, d->atoms->atom_has_trailer, val, trailer);
    } else if(d->i < bin.size) {
        ret = dec_error(d, "invalid_trailing_data");
        goto done;
    }

    if(dec_pop(d) != st_done) {
        ret = dec_error(d, "truncated_json");
    } else if(d->is_partial) {
        ret = enif_make_tuple2(env, d->atoms->atom_partial, val);
    } else {
        ret = val;
    }

done:
    bump_used_reds(env, bytes_processed, d->bytes_per_red);

    return ret;
}