From cd1b2632083c3724b246c8e8107845d9903bc578 Mon Sep 17 00:00:00 2001
From: Jeremie Lasalle Ratelle <jeremie.ratelle@gmail.com>
Date: Fri, 24 Apr 2015 12:35:32 -0400
Subject: [PATCH] Add an option to control null value decoding atom

This is pretty much a generalisation of the use_nil option to support an
arbitrary atom.
---
 c_src/decoder.c                 | 10 ++++++----
 c_src/jiffy.c                   |  1 +
 c_src/jiffy.h                   |  2 ++
 c_src/util.c                    | 28 ++++++++++++++++++++++++++++
 test/jiffy_02_literal_tests.erl | 11 +++++++++++
 5 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/c_src/decoder.c b/c_src/decoder.c
index ad3e8d3..f39575a 100644
--- a/c_src/decoder.c
+++ b/c_src/decoder.c
@@ -52,7 +52,7 @@ typedef struct {
     size_t          bytes_per_iter;
     int             is_partial;
     int             return_maps;
-    int             use_nil;
+    ERL_NIF_TERM    null_term;
 
     char*           p;
     unsigned char*  u;
@@ -80,7 +80,7 @@ dec_new(ErlNifEnv* env)
     d->bytes_per_iter = DEFAULT_BYTES_PER_ITER;
     d->is_partial = 0;
     d->return_maps = 0;
-    d->use_nil = 0;
+    d->null_term = d->atoms->atom_null;
 
     d->p = NULL;
     d->u = NULL;
@@ -711,7 +711,9 @@ decode_init(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
             return enif_make_badarg(env);
 #endif
         } else if(enif_compare(val, d->atoms->atom_use_nil) == 0) {
-            d->use_nil = 1;
+            d->null_term = d->atoms->atom_nil;
+        } else if(get_null_term(env, val, &(d->null_term))) {
+            continue;
         } else {
             return enif_make_badarg(env);
         }
@@ -783,7 +785,7 @@ decode_iter(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
                             ret = dec_error(d, "invalid_literal");
                             goto done;
                         }
-                        val = d->use_nil ? d->atoms->atom_nil : d->atoms->atom_null;
+                        val = d->null_term;
                         dec_pop(d, st_value);
                         d->i += 4;
                         break;
diff --git a/c_src/jiffy.c b/c_src/jiffy.c
index 31aa854..c18a111 100644
--- a/c_src/jiffy.c
+++ b/c_src/jiffy.c
@@ -28,6 +28,7 @@ load(ErlNifEnv* env, void** priv, ERL_NIF_TERM info)
     st->atom_return_maps = make_atom(env, "return_maps");
     st->atom_nil = make_atom(env, "nil");
     st->atom_use_nil = make_atom(env, "use_nil");
+    st->atom_null_term = make_atom(env, "null_term");
 
     // Markers used in encoding
     st->ref_object = make_atom(env, "$object_ref$");
diff --git a/c_src/jiffy.h b/c_src/jiffy.h
index 6d869a5..ecad9da 100644
--- a/c_src/jiffy.h
+++ b/c_src/jiffy.h
@@ -30,6 +30,7 @@ typedef struct {
     ERL_NIF_TERM    atom_return_maps;
     ERL_NIF_TERM    atom_nil;
     ERL_NIF_TERM    atom_use_nil;
+    ERL_NIF_TERM    atom_null_term;
 
     ERL_NIF_TERM    ref_object;
     ERL_NIF_TERM    ref_array;
@@ -44,6 +45,7 @@ ERL_NIF_TERM make_error(jiffy_st* st, ErlNifEnv* env, const char* error);
 ERL_NIF_TERM make_obj_error(jiffy_st* st, ErlNifEnv* env, const char* error,
         ERL_NIF_TERM obj);
 int get_bytes_per_iter(ErlNifEnv* env, ERL_NIF_TERM val, size_t* bpi);
+int get_null_term(ErlNifEnv* env, ERL_NIF_TERM val, ERL_NIF_TERM *null_term);
 int should_yield(size_t used, size_t limit);
 int consume_timeslice(ErlNifEnv* env, size_t used, size_t limit);
 
diff --git a/c_src/util.c b/c_src/util.c
index c30a6d7..f05e1e8 100644
--- a/c_src/util.c
+++ b/c_src/util.c
@@ -62,6 +62,34 @@ get_bytes_per_iter(ErlNifEnv* env, ERL_NIF_TERM val, size_t* bpi)
     return 1;
 }
 
+int
+get_null_term(ErlNifEnv* env, ERL_NIF_TERM val, ERL_NIF_TERM *null_term)
+{
+    jiffy_st* st = (jiffy_st*) enif_priv_data(env);
+    const ERL_NIF_TERM* tuple;
+    int arity;
+
+    if(!enif_get_tuple(env, val, &arity, &tuple)) {
+      return 0;
+    }
+
+    if(arity != 2) {
+      return 0;
+    }
+
+    if(enif_compare(tuple[0], st->atom_null_term) != 0) {
+      return 0;
+    }
+
+    if(!enif_is_atom(env, tuple[1])) {
+      return 0;
+    }
+
+    *null_term = tuple[1];
+
+    return 1;
+}
+
 int
 should_yield(size_t used, size_t limit)
 {
diff --git a/test/jiffy_02_literal_tests.erl b/test/jiffy_02_literal_tests.erl
index 98b163c..2699199 100644
--- a/test/jiffy_02_literal_tests.erl
+++ b/test/jiffy_02_literal_tests.erl
@@ -31,3 +31,14 @@ nil_test_() ->
         {"Decode", ?_assertEqual(nil, jiffy:decode(<<"null">>, [use_nil]))},
         {"Encode", ?_assertEqual(<<"null">>, jiffy:encode(nil, [use_nil]))}
     ]}.
+
+null_term_test_() ->
+    T = [
+        {undefined, [{null_term, undefined}]},
+        {whatever, [{null_term, whatever}]},
+        {undefined, [use_nil, {null_term, undefined}]},
+        {nil, [{null_term, undefined}, use_nil]},
+        {whatever, [{null_term, undefined}, {null_term, whatever}]}
+    ],
+    {"null_term",
+        [?_assertEqual(R, jiffy:decode(<<"null">>, O)) || {R, O} <- T]}.