From 72790ca4f5eb687395378ea9c225b48743c36885 Mon Sep 17 00:00:00 2001
From: Sergey Urbanovich <sergey.urbanovich@gmail.com>
Date: Sun, 11 Dec 2011 18:11:56 +0400
Subject: [PATCH] Add pretty print

---
 c_src/encoder.c      | 56 +++++++++++++++++++++++++++++++++++++++-----
 c_src/jiffy.c        |  1 +
 c_src/jiffy.h        |  1 +
 test/jiffy_tests.erl |  4 ++--
 4 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/c_src/encoder.c b/c_src/encoder.c
index a751fd7..e3a7e3a 100644
--- a/c_src/encoder.c
+++ b/c_src/encoder.c
@@ -15,7 +15,9 @@ typedef struct {
     ErlNifEnv*      env;
     jiffy_st*       atoms;
     int             uescape;
+    int             pretty;
 
+    int             shiftcnt;
     int             count;
 
     int             iolen;
@@ -36,6 +38,8 @@ enc_init(Encoder* e, ErlNifEnv* env, ERL_NIF_TERM opts, ErlNifBinary* bin)
     e->env = env;
     e->atoms = enif_priv_data(env);
     e->uescape = 0;
+    e->pretty = 0;
+    e->shiftcnt = 0;
     e->count = 0;
 
     if(!enif_is_list(env, opts)) {
@@ -45,6 +49,8 @@ enc_init(Encoder* e, ErlNifEnv* env, ERL_NIF_TERM opts, ErlNifBinary* bin)
     while(enif_get_list_cell(env, opts, &val, &opts)) {
         if(enif_compare(val, e->atoms->atom_uescape) == 0) {
             e->uescape = 1;
+        } else if(enif_compare(val, e->atoms->atom_pretty) == 0) {
+            e->pretty = 1;
         } else {
             return 0;
         }
@@ -376,42 +382,80 @@ enc_char(Encoder* e, char c)
     return 1;
 }
 
+static inline int
+enc_shift(Encoder* e) {
+    int i = 0;
+    if (!e->pretty)
+        return 1;
+
+    switch (e->shiftcnt) {
+        case 0 :
+            return enc_char(e, '\n');
+        case 1 :
+            return enc_literal(e, "\n  ", 3);
+        case 2 :
+            return enc_literal(e, "\n    ", 5);
+        case 3 :
+            return enc_literal(e, "\n      ", 7);
+        case 4 :
+            return enc_literal(e, "\n        ", 9);
+        case 5 :
+            return enc_literal(e, "\n          ", 11);
+        default :
+            if (!enc_char(e, '\n'))
+                return 1;
+            for (;i < e->shiftcnt; i++) {
+                if (!enc_literal(e, "  ", 2))
+                    return 0;
+            }
+            return 1;
+    }
+}
+
 static inline int
 enc_start_object(Encoder* e)
 {
     e->count++;
-    return enc_char(e, '{');
+    e->shiftcnt++;
+    return enc_char(e, '{') && enc_shift(e);
 }
 
 static inline int
 enc_end_object(Encoder* e)
 {
-    return enc_char(e, '}');
+    e->shiftcnt--;
+    return enc_shift(e) && enc_char(e, '}');
 }
 
 static inline int
 enc_start_array(Encoder* e)
 {
     e->count++;
-    return enc_char(e, '[');
+    e->shiftcnt++;
+    return enc_char(e, '[') && enc_shift(e);
 }
 
 static inline int
 enc_end_array(Encoder* e)
 {
-    return enc_char(e, ']');
+    e->shiftcnt--;
+    return enc_shift(e) && enc_char(e, ']');
 }
 
 static inline int
 enc_colon(Encoder* e)
 {
-    return enc_char(e, ':');
+    if (!e->pretty)
+        return enc_char(e, ':');
+    return enc_literal(e, " : ", 3);
 }
 
 static inline int
 enc_comma(Encoder* e)
 {
-    return enc_char(e, ',');
+    if (!e->pretty)
+        return enc_char(e, ',');
+    return enc_literal(e, ", ", 2) && enc_shift(e);
 }
 
 ERL_NIF_TERM
diff --git a/c_src/jiffy.c b/c_src/jiffy.c
index c1dd8b1..8fdde2f 100644
--- a/c_src/jiffy.c
+++ b/c_src/jiffy.c
@@ -21,6 +21,7 @@ load(ErlNifEnv* env, void** priv, ERL_NIF_TERM info)
     st->atom_bigdbl = make_atom(env, "bigdbl");
     st->atom_partial = make_atom(env, "partial");
     st->atom_uescape = make_atom(env, "uescape");
+    st->atom_pretty = make_atom(env, "pretty");
 
     // 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 9b1f700..327657a 100644
--- a/c_src/jiffy.h
+++ b/c_src/jiffy.h
@@ -17,6 +17,7 @@ typedef struct {
     ERL_NIF_TERM    atom_bigdbl;
     ERL_NIF_TERM    atom_partial;
     ERL_NIF_TERM    atom_uescape;
+    ERL_NIF_TERM    atom_pretty;
 
     ERL_NIF_TERM    ref_object;
     ERL_NIF_TERM    ref_array;
diff --git a/test/jiffy_tests.erl b/test/jiffy_tests.erl
index f7075c0..4c83e36 100644
--- a/test/jiffy_tests.erl
+++ b/test/jiffy_tests.erl
@@ -17,10 +17,10 @@ proper_test_() ->
 
 
 prop_encode_decode() ->
-    ?FORALL(Data, json(),
+    ?FORALL({Data, Opts}, {json(), union([[], [pretty]])},
         begin
             %io:format(standard_error, "Data: ~p~n", [Data]),
-            Data == jiffy:decode(jiffy:encode(Data))
+            Data == jiffy:decode(jiffy:encode(Data, Opts))
         end
     ).