Refactor trapping and trap more often during decode

c_src/decoder.c

@@ -715,7 +715,9 @@ decode_iter(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     ERL_NIF_TERM val = argv[2];
     ERL_NIF_TERM trailer;
     ERL_NIF_TERM ret;
-    size_t bytes_read = 0;
+
+    size_t start;
+    size_t bytes_processed = 0;
 
     if(argc != 5) {
         return enif_make_badarg(env);
@@ -733,20 +735,29 @@ decode_iter(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     objs = argv[3];
     curr = argv[4];
 
+    start = d->i;
+
     while(d->i < bin.size) {
-        if(should_yield(env, &bytes_read, d->bytes_per_red)) {
-            return enif_make_tuple5(
-                    env,
-                    st->atom_iter,
-                    argv[1],
-                    val,
-                    objs,
-                    curr
-                );
+        bytes_processed = d->i - start;
+
+        if(should_yield(env, bytes_processed, d->bytes_per_red)) {
+            ERL_NIF_TERM tmp_argv[5];
+
+            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);
         }
 
-        bytes_read += 1;
-
         switch(dec_curr(d)) {
             case st_value:
                 switch(d->p[d->i]) {
@@ -1040,5 +1051,7 @@ decode_done:
     }
 
 done:
+    bump_used_reds(env, bytes_processed, d->bytes_per_red);
+
     return ret;
 }

c_src/encoder.c

@@ -823,7 +823,7 @@ encode_iter(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     double dval;
 
     size_t start;
-    size_t bytes_written = 0;
+    size_t bytes_processed = 0;
 
     if(argc != 3) {
         return enif_make_badarg(env);
@@ -846,20 +846,24 @@ encode_iter(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     start = e->iosize + e->i;
 
     while(!termstack_is_empty(&stack)) {
-        bytes_written += (e->iosize + e->i) - start;
+        size_t bytes_processed = (e->iosize + e->i) - start;
 
-        if(should_yield(env, &bytes_written, e->bytes_per_red)) {
-            ERL_NIF_TERM saved_stack = termstack_save(env, &stack);
+        if(should_yield(env, bytes_processed, e->bytes_per_red)) {
+            ERL_NIF_TERM tmp_argv[3];
+
+            tmp_argv[0] = argv[0];
+            tmp_argv[1] = termstack_save(env, &stack);
+            tmp_argv[2] = e->iolist;
 
             termstack_destroy(&stack);
 
-            return enif_make_tuple4(
-                    env,
-                    st->atom_iter,
-                    argv[0],
-                    saved_stack,
-                    e->iolist
-                );
+            bump_used_reds(env, bytes_processed, e->bytes_per_red);
+            return enif_schedule_nif(env,
+                                     "nif_encode_iter",
+                                     0,
+                                     encode_iter,
+                                     3,
+                                     tmp_argv);
         }
 
         curr = termstack_pop(&stack);
@@ -1042,6 +1046,7 @@ encode_iter(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     }
 
 done:
+    bump_used_reds(env, bytes_processed, e->bytes_per_red);
     termstack_destroy(&stack);
 
     return ret;

c_src/jiffy.c

@@ -83,10 +83,8 @@ unload(ErlNifEnv* env, void* priv)
 
 static ErlNifFunc funcs[] =
 {
-    {"nif_decode_init", 2, decode_init},
-    {"nif_decode_iter", 5, decode_iter},
-    {"nif_encode_init", 2, encode_init},
-    {"nif_encode_iter", 3, encode_iter}
+    {"nif_decode", 2, decode_init},
+    {"nif_encode", 2, encode_init}
 };
 
 ERL_NIF_INIT(jiffy, funcs, &load, &reload, &upgrade, &unload);

c_src/jiffy.h

@@ -53,7 +53,8 @@ ERL_NIF_TERM make_obj_error(jiffy_st* st, ErlNifEnv* env, const char* error,
 int get_bytes_per_iter(ErlNifEnv* env, ERL_NIF_TERM val, size_t* bpi);
 int get_bytes_per_red(ErlNifEnv* env, ERL_NIF_TERM val, size_t* bpr);
 int get_null_term(ErlNifEnv* env, ERL_NIF_TERM val, ERL_NIF_TERM *null_term);
-int should_yield(ErlNifEnv* env, size_t* used, size_t bytes_per_red);
+int should_yield(ErlNifEnv* env, size_t used, size_t bytes_per_red);
+void bump_used_reds(ErlNifEnv* env, size_t used, size_t bytes_per_red);
 
 ERL_NIF_TERM decode_init(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 ERL_NIF_TERM decode_iter(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);

c_src/util.c

@@ -122,20 +122,34 @@ get_null_term(ErlNifEnv* env, ERL_NIF_TERM val, ERL_NIF_TERM *null_term)
 }
 
 int
-should_yield(ErlNifEnv* env, size_t* used, size_t bytes_per_red)
+should_yield(ErlNifEnv* env, size_t used, size_t bytes_per_red)
 {
-#if(ERL_NIF_MAJOR_VERSION >= 2 && ERL_NIF_MINOR_VERSION >= 4)
+    (void)env;
+    (void)used;
+    (void)bytes_per_red;
 
-    if(((*used) / bytes_per_red) >= 20) {
-        *used = 0;
-        return enif_consume_timeslice(env, 1);
-    }
+    return (used / bytes_per_red) >= DEFAULT_ERLANG_REDUCTION_COUNT;
+}
 
-    return 0;
+void
+bump_used_reds(ErlNifEnv* env, size_t used, size_t bytes_per_red)
+{
+#if(ERL_NIF_MAJOR_VERSION >= 2 && ERL_NIF_MINOR_VERSION >= 4)
+    size_t reds_used, pct_used;
 
-#else
+    reds_used = used / bytes_per_red;
+    pct_used = 100 * reds_used / DEFAULT_ERLANG_REDUCTION_COUNT;
 
-    return ((*used) / bytes_per_red) >= DEFAULT_ERLANG_REDUCTION_COUNT;
+    if(pct_used > 0) {
+        if(pct_used > 100) {
+            pct_used = 100;
+        }
 
+        enif_consume_timeslice(env, pct_used);
+    }
 #endif
-}
+
+    (void)env;
+    (void)used;
+    (void)bytes_per_red;
+}

src/jiffy.erl

@@ -66,13 +66,11 @@ decode(Data) ->
 
 -spec decode(iolist() | binary(), decode_options()) -> jiffy_decode_result().
 decode(Data, Opts) when is_binary(Data), is_list(Opts) ->
-    case nif_decode_init(Data, Opts) of
+    case nif_decode(Data, Opts) of
         {error, _} = Error ->
             throw(Error);
         {partial, EJson} ->
             finish_decode(EJson);
-        {iter, Decoder, Val, Objs, Curr} ->
-            decode_loop(Data, Decoder, Val, Objs, Curr);
         EJson ->
             EJson
     end;
@@ -88,7 +86,7 @@ encode(Data) ->
 -spec encode(json_value(), encode_options()) -> iodata().
 encode(Data, Options) ->
     ForceUTF8 = lists:member(force_utf8, Options),
-    case nif_encode_init(Data, Options) of
+    case nif_encode(Data, Options) of
         {error, {invalid_string, _}} when ForceUTF8 == true ->
             FixedData = jiffy_utf8:fix(Data),
             encode(FixedData, Options -- [force_utf8]);
@@ -99,8 +97,6 @@ encode(Data, Options) ->
             throw(Error);
         {partial, IOData} ->
             finish_encode(IOData, []);
-        {iter, Encoder, Stack, IOBuf} ->
-            encode_loop(Data, Options, Encoder, Stack, IOBuf);
         IOData ->
             IOData
     end.
@@ -177,51 +173,11 @@ init() ->
     end,
     erlang:load_nif(filename:join(PrivDir, "jiffy"), 0).
 
-
-decode_loop(Data, Decoder, Val, Objs, Curr) ->
-    case nif_decode_iter(Data, Decoder, Val, Objs, Curr) of
-        {error, _} = Error ->
-            throw(Error);
-        {partial, EJson} ->
-            finish_decode(EJson);
-        {iter, NewDecoder, NewVal, NewObjs, NewCurr} ->
-            decode_loop(Data, NewDecoder, NewVal, NewObjs, NewCurr);
-        EJson ->
-            EJson
-    end.
-
-
-encode_loop(Data, Options, Encoder, Stack, IOBuf) ->
-    ForceUTF8 = lists:member(force_utf8, Options),
-    case nif_encode_iter(Encoder, Stack, IOBuf) of
-        {error, {invalid_string, _}} when ForceUTF8 == true ->
-            FixedData = jiffy_utf8:fix(Data),
-            encode(FixedData, Options -- [force_utf8]);
-        {error, {invalid_object_member_key, _}} when ForceUTF8 == true ->
-            FixedData = jiffy_utf8:fix(Data),
-            encode(FixedData, Options -- [force_utf8]);
-        {error, _} = Error ->
-            throw(Error);
-        {partial, IOData} ->
-            finish_encode(IOData, []);
-        {iter, NewEncoder, NewStack, NewIOBuf} ->
-            encode_loop(Data, Options, NewEncoder, NewStack, NewIOBuf);
-        IOData ->
-            IOData
-    end.
-
-
 not_loaded(Line) ->
     erlang:nif_error({not_loaded, [{module, ?MODULE}, {line, Line}]}).
 
-nif_decode_init(_Data, _Opts) ->
-    ?NOT_LOADED.
-
-nif_decode_iter(_Data, _Decoder, _, _, _) ->
-    ?NOT_LOADED.
-
-nif_encode_init(_Data, _Options) ->
+nif_decode(_Data, _Opts) ->
     ?NOT_LOADED.
 
-nif_encode_iter(_Encoder, _Stack, _IoList) ->
+nif_encode(_Data, _Options) ->
     ?NOT_LOADED.