feature: added new cjson.array_mt metatable to allow enforcing JSON array encoding.

Signed-off-by: Yichun Zhang (agentzh) <agentzh@gmail.com>

3 files changed, 155 insertions, 21 deletions

diff --git a/README.md b/README.md
index 7282a4d..1cdb53b 100644
--- a/README.md
+++ b/README.md
@@ -11,6 +11,7 @@ Table of Contents
 * [Additions to mpx/lua](#additions)
     * [encode_empty_table_as_object](#encode_empty_table_as_object)
     * [empty_array](#empty_array)
+    * [array_mt](#array_mt)
     * [empty_array_mt](#empty_array_mt)
     * [encode_number_precision](#encode_number_precision)
 
@@ -76,6 +77,39 @@ This will generate:
 
 [Back to TOC](#table-of-contents)
 
+array_mt
+--------
+**syntax:** `setmetatable({}, cjson.array_mt)`
+
+When lua-cjson encodes a table with this metatable, it will systematically
+encode it as a JSON Array. The resulting, encoded Array will contain the array
+part of the table, and will be of the same length as the `#` operator on that
+table. Holes in the table will be encoded with the `null` JSON value.
+
+Example:
+
+```lua
+local t = { "hello", "world" }
+setmetatable(t, cjson.array_mt)
+cjson.encode(t) -- ["hello","world"]
+```
+
+Or:
+
+```lua
+local t = {}
+t[1] = "one"
+t[2] = "two"
+t[4] = "three"
+t.foo = "bar"
+setmetatable(t, cjson.array_mt)
+cjson.encode(t) -- ["one","two",null,"three"]
+```
+
+This value was introduced in the `2.1.0.5` release of this module.
+
+[Back to TOC](#table-of-contents)
+
 empty_array_mt
 --------------
 **syntax:** `setmetatable({}, cjson.empty_array_mt)`
diff --git a/lua_cjson.c b/lua_cjson.c
index d04f151..8d6e313 100644
--- a/lua_cjson.c
+++ b/lua_cjson.c
@@ -92,6 +92,7 @@
 #endif
 
 static const char * const *json_empty_array;
+static const char * const *json_array;
 
 typedef enum {
     T_OBJ_BEGIN,
@@ -713,21 +714,37 @@ static void json_append_data(lua_State *l, json_config_t *cfg,
     case LUA_TTABLE:
         current_depth++;
         json_check_encode_depth(l, cfg, current_depth, json);
-        len = lua_array_length(l, cfg, json);
-        if (len > 0 || (len == 0 && !cfg->encode_empty_table_as_object))
+
+        int as_array = 0;
+        int has_metatable = lua_getmetatable(l, -1);
+
+        if (has_metatable) {
+            lua_pushlightuserdata(l, &json_array);
+            lua_rawget(l, LUA_REGISTRYINDEX);
+            as_array = lua_rawequal(l, -1, -2);
+            lua_pop(l, 2);
+        }
+
+        if (as_array) {
+            len = lua_objlen(l, -1);
             json_append_array(l, cfg, current_depth, json, len);
-        else {
-            int as_array = 0;
-            if (lua_getmetatable(l, -1)) {
-                lua_pushlightuserdata(l, &json_empty_array);
-                lua_rawget(l, LUA_REGISTRYINDEX);
-                as_array = lua_rawequal(l, -1, -2);
-                lua_pop(l, 2);
-            }
+        } else {
+            len = lua_array_length(l, cfg, json);
 
-            if (as_array) {
-                json_append_array(l, cfg, current_depth, json, 0);
+            if (len > 0 || (len == 0 && !cfg->encode_empty_table_as_object)) {
+                json_append_array(l, cfg, current_depth, json, len);
             } else {
+                if (has_metatable) {
+                    lua_getmetatable(l, -1);
+                    lua_pushlightuserdata(l, &json_empty_array);
+                    lua_rawget(l, LUA_REGISTRYINDEX);
+                    as_array = lua_rawequal(l, -1, -2);
+                    lua_pop(l, 2); /* pop pointer + metatable */
+                    if (as_array) {
+                        json_append_array(l, cfg, current_depth, json, 0);
+                        break;
+                    }
+                }
                 json_append_object(l, cfg, current_depth, json);
             }
         }
@@ -738,7 +755,7 @@ static void json_append_data(lua_State *l, json_config_t *cfg,
     case LUA_TLIGHTUSERDATA:
         if (lua_touserdata(l, -1) == NULL) {
             strbuf_append_mem(json, "null", 4);
-        } else if (lua_touserdata(l, -1) == &json_empty_array) {
+        } else if (lua_touserdata(l, -1) == &json_array) {
             json_append_array(l, cfg, current_depth, json, 0);
         }
         break;
@@ -1412,20 +1429,27 @@ static int lua_cjson_new(lua_State *l)
     /* Initialise number conversions */
     fpconv_init();
 
-    /* Test if empty array metatable is in registry */
+    /* Test if array metatables are in registry */
     lua_pushlightuserdata(l, &json_empty_array);
     lua_rawget(l, LUA_REGISTRYINDEX);
     if (lua_isnil(l, -1)) {
-        /* Create empty array metatable.
+        /* Create array metatables.
          *
          * If multiple calls to lua_cjson_new() are made,
-         * this prevents overriding the table at the given
+         * this prevents overriding the tables at the given
          * registry's index with a new one.
          */
         lua_pop(l, 1);
+
+        /* empty_array_mt */
         lua_pushlightuserdata(l, &json_empty_array);
         lua_newtable(l);
         lua_rawset(l, LUA_REGISTRYINDEX);
+
+        /* array_mt */
+        lua_pushlightuserdata(l, &json_array);
+        lua_newtable(l);
+        lua_rawset(l, LUA_REGISTRYINDEX);
     }
 
     /* cjson module table */
@@ -1444,8 +1468,13 @@ static int lua_cjson_new(lua_State *l)
     lua_rawget(l, LUA_REGISTRYINDEX);
     lua_setfield(l, -2, "empty_array_mt");
 
+    /* Set cjson.array_mt */
+    lua_pushlightuserdata(l, &json_array);
+    lua_rawget(l, LUA_REGISTRYINDEX);
+    lua_setfield(l, -2, "array_mt");
+
     /* Set cjson.empty_array */
-    lua_pushlightuserdata(l, &json_empty_array);
+    lua_pushlightuserdata(l, &json_array);
     lua_setfield(l, -2, "empty_array");
 
     /* Set module name / version fields */
diff --git a/tests/agentzh.t b/tests/agentzh.t
index 0b281a1..dd70fb8 100644
--- a/tests/agentzh.t
+++ b/tests/agentzh.t
@@ -116,7 +116,78 @@ print(cjson.encode(data))
 
 
 
-=== TEST 9: multiple calls to lua_cjson_new (1/2)
+=== TEST 9: array_mt on empty tables
+--- lua
+local cjson = require "cjson"
+local data = {}
+setmetatable(data, cjson.array_mt)
+print(cjson.encode(data))
+--- out
+[]
+
+
+
+=== TEST 10: array_mt on non-empty tables
+--- lua
+local cjson = require "cjson"
+local data = { "foo", "bar" }
+setmetatable(data, cjson.array_mt)
+print(cjson.encode(data))
+--- out
+["foo","bar"]
+
+
+
+=== TEST 11: array_mt on non-empty tables with holes
+--- lua
+local cjson = require "cjson"
+local data = {}
+data[1] = "foo"
+data[2] = "bar"
+data[4] = "last"
+data[9] = "none"
+setmetatable(data, cjson.array_mt)
+print(cjson.encode(data))
+--- out
+["foo","bar",null,"last"]
+
+
+
+=== TEST 12: array_mt on tables with hash part
+--- lua
+local cjson = require "cjson"
+local data
+
+if jit and string.find(jit.version, "LuaJIT 2.1.0", nil, true) then
+    local new_tab = require "table.new"
+    data = new_tab(0, 2) -- allocating hash part only
+
+else
+    data = {}
+end
+
+data.foo = "bar"
+data[1] = "hello"
+setmetatable(data, cjson.array_mt)
+print(cjson.encode(data))
+--- out
+["hello"]
+
+
+
+=== TEST 13: multiple calls to lua_cjson_new (1/3)
+--- lua
+local cjson = require "cjson"
+package.loaded["cjson"] = nil
+require "cjson"
+local arr = setmetatable({}, cjson.array_mt)
+print(cjson.encode(arr))
+--- out
+[]
+
+
+
+=== TEST 14: multiple calls to lua_cjson_new (2/3)
 --- lua
 local cjson = require "cjson"
 package.loaded["cjson"] = nil
@@ -128,7 +199,7 @@ print(cjson.encode(arr))
 
 
 
-=== TEST 10: multiple calls to lua_cjson_new (2/2)
+=== TEST 15: multiple calls to lua_cjson_new (3/3)
 --- lua
 local cjson = require "cjson.safe"
 -- load another cjson instance (not in package.loaded)
@@ -140,7 +211,7 @@ print(cjson.encode(arr))
 
 
 
-=== TEST 11: & in JSON
+=== TEST 16: & in JSON
 --- lua
 local cjson = require "cjson"
 local a="[\"a=1&b=2\"]"
@@ -151,7 +222,7 @@ print(cjson.encode(b))
 
 
 
-=== TEST 12: default and max precision
+=== TEST 17: default and max precision
 --- lua
 local math = require "math"
 local cjson = require "cjson"