From 4484c1ed50d8baf39b6559296a809430d0dd525f Mon Sep 17 00:00:00 2001
From: Cloud Wu <cloudwu@gmail.com>
Date: Sat, 10 Oct 2020 13:55:15 +0800
Subject: [PATCH] Add map for sproto

---
 lualib-src/sproto/README.md |  10 ++-
 lualib-src/sproto/lsproto.c | 109 +++++++++++++++++++------
 lualib-src/sproto/sproto.c  | 154 +++++++++++++++++++++++-------------
 lualib-src/sproto/sproto.h  |   9 ++-
 lualib/sprotoparser.lua     |  57 ++++++++++---
 5 files changed, 245 insertions(+), 94 deletions(-)

diff --git a/lualib-src/sproto/README.md b/lualib-src/sproto/README.md
index 237e4b5c..e5011fc9 100644
--- a/lualib-src/sproto/README.md
+++ b/lualib-src/sproto/README.md
@@ -165,6 +165,7 @@ The schema text is like this:
     phone 3 : *PhoneNumber	# *PhoneNumber means an array of PhoneNumber.
     height 4 : integer(2)	# (2) means a 1/100 fixed-point number.
     data 5 : binary		# Some binary data
+    weight 6 : double   # floating number
 }
 
 .AddressBook {
@@ -216,11 +217,14 @@ Types
 * **string** : string
 * **binary** : binary string (it's a sub type of string)
 * **integer** : integer, the max length of an integer is signed 64bit. It can be a fixed-point number with specified precision.
+* **double** : double, floating-point number.
 * **boolean** : true or false
 
-You can add * before the typename to declare an array. 
+You can add * before the typename to declare an array.
 
-You can also specify a main index, the array whould be encode as an unordered map.
+You can also specify a main index with the syntax likes `*array(id)`, the array would be encode as an unordered map with the `id` field as key.
+
+For empty main index likes `*array()`, the array would be encoded as an unordered map with the first field as key and the second field as value.
 
 User defined type can be any name in alphanumeric characters except the build-in typenames, and nested types are supported.
 
@@ -228,6 +232,8 @@ User defined type can be any name in alphanumeric characters except the build-in
 
 I have been using Google protocol buffers for many years in many projects, and I found the real types were seldom used. If you really need it, you can use string to serialize the double numbers. When you need decimal, you can specify the fixed-point presision.
 
+**NOTE** : `double` is supported now.
+
 * Where is enum?
 
 In lua, enum types are not very useful. You can use integer to define an enum table in lua.
diff --git a/lualib-src/sproto/lsproto.c b/lualib-src/sproto/lsproto.c
index 214ec8df..370ef719 100644
--- a/lualib-src/sproto/lsproto.c
+++ b/lualib-src/sproto/lsproto.c
@@ -128,6 +128,7 @@ struct encode_ud {
 	const char * array_tag;
 	int array_index;
 	int deep;
+	int map_entry;
 	int iter_func;
 	int iter_table;
 	int iter_key;
@@ -153,13 +154,11 @@ next_list(lua_State *L, struct encode_ud * self) {
 }
 
 static int
-encode(const struct sproto_arg *args) {
+get_encodefield(const struct sproto_arg *args) {
 	struct encode_ud *self = args->ud;
 	lua_State *L = self->L;
-	luaL_checkstack(L, 12, NULL);
-	if (self->deep >= ENCODE_DEEPLEVEL)
-		return luaL_error(L, "The table is too deep");
 	if (args->index > 0) {
+		int map = args->ktagname != NULL;
 		if (args->tagname != self->array_tag) {
 			// a new array
 			self->array_tag = args->tagname;
@@ -177,6 +176,13 @@ encode(const struct sproto_arg *args) {
 				self->array_index = lua_gettop(L);
 			}
 
+			if (map) {
+				if (!self->map_entry) {
+					lua_createtable(L, 0, 2); // key/value entry
+					self->map_entry = lua_gettop(L);
+				}
+			}
+
 			if (luaL_getmetafield(L, self->array_index, "__pairs")) {
 				lua_pushvalue(L, self->array_index);
 				lua_call(L,	1, 3);
@@ -194,26 +200,39 @@ encode(const struct sproto_arg *args) {
 				self->iter_key = lua_gettop(L);
 			}
 		}
-		if (args->mainindex >= 0) {
+		if (args->mainindex >= 0) { // *type(mainindex)
 			if (!next_list(L, self)) {
 				// iterate end
 				lua_pushnil(L);
 				lua_replace(L, self->iter_key);
 				return SPROTO_CB_NIL;
 			}
-			lua_insert(L, -2);
-			lua_replace(L, self->iter_key);
+			if (map) {
+				lua_pushvalue(L, -2);
+				lua_replace(L, self->iter_key);
+				lua_setfield(L, self->map_entry, args->vtagname);
+				lua_setfield(L, self->map_entry, args->ktagname);
+				lua_pushvalue(L, self->map_entry);
+			} else {
+				lua_insert(L, -2);
+				lua_replace(L, self->iter_key);
+			}
 		} else {
 			lua_geti(L, self->array_index, args->index);
 		}
 	} else {
 		lua_getfield(L, self->tbl_index, args->tagname);
 	}
-	if (lua_isnil(L, -1)) {
-		lua_pop(L,1);
-		return SPROTO_CB_NIL;
-	}
-	switch (args->type) {
+	return 0;
+}
+
+static int encode(const struct sproto_arg *args);
+
+static int
+encode_one(const struct sproto_arg *args, struct encode_ud *self) {
+	lua_State *L = self->L;
+	int type = args->type;
+	switch (type) {
 	case SPROTO_TINTEGER: {
 		int64_t v;
 		lua_Integer vh;
@@ -282,6 +301,7 @@ encode(const struct sproto_arg *args) {
 		sub.array_tag = NULL;
 		sub.array_index = 0;
 		sub.deep = self->deep + 1;
+		sub.map_entry = 0;
 		sub.iter_func = 0;
 		sub.iter_table = 0;
 		sub.iter_key = 0;
@@ -296,6 +316,25 @@ encode(const struct sproto_arg *args) {
 	}
 }
 
+static int
+encode(const struct sproto_arg *args) {
+	struct encode_ud *self = args->ud;
+	lua_State *L = self->L;
+	int code;
+	luaL_checkstack(L, 12, NULL);
+	if (self->deep >= ENCODE_DEEPLEVEL)
+		return luaL_error(L, "The table is too deep");
+	code = get_encodefield(args);
+	if (code < 0) {
+		return code;
+	}
+	if (lua_isnil(L, -1)) {
+		lua_pop(L,1);
+		return SPROTO_CB_NIL;
+	}
+	return encode_one(args, self);
+}
+
 static void *
 expand_buffer(lua_State *L, int osz, int nsz) {
 	void *output;
@@ -342,6 +381,7 @@ lencode(lua_State *L) {
 		self.deep = 0;
 
 		lua_settop(L, tbl_index);
+		self.map_entry = 0;
 		self.iter_func = 0;
 		self.iter_table = 0;
 		self.iter_key = 0;
@@ -365,6 +405,7 @@ struct decode_ud {
 	int deep;
 	int mainindex_tag;
 	int key_index;
+	int map_entry;
 };
 
 static int
@@ -422,14 +463,24 @@ decode(const struct sproto_arg *args) {
 		break;
 	}
 	case SPROTO_TSTRUCT: {
+		int map = args->ktagname != NULL;
 		struct decode_ud sub;
 		int r;
-		lua_newtable(L);
 		sub.L = L;
-		sub.result_index = lua_gettop(L);
+		if (map) {
+			if (!self->map_entry) {
+				lua_newtable(L);
+				self->map_entry = lua_gettop(L);
+			}
+			sub.result_index = self->map_entry;
+		} else {
+			lua_newtable(L);
+			sub.result_index = lua_gettop(L);
+		}
 		sub.deep = self->deep + 1;
 		sub.array_index = 0;
 		sub.array_tag = NULL;
+		sub.map_entry = 0;
 		if (args->mainindex >= 0) {
 			// This struct will set into a map, so mark the main index tag.
 			sub.mainindex_tag = args->mainindex;
@@ -441,13 +492,26 @@ decode(const struct sproto_arg *args) {
 				return SPROTO_CB_ERROR;
 			if (r != args->length)
 				return r;
-			lua_pushvalue(L, sub.key_index);
-			if (lua_isnil(L, -1)) {
-				luaL_error(L, "Can't find main index (tag=%d) in [%s]", args->mainindex, args->tagname);
+			if (map) {
+				lua_getfield(L, sub.result_index, args->ktagname);
+				if (lua_isnil(L, -1)) {
+					luaL_error(L, "Can't find key field in [%s]", args->tagname);
+				}
+				lua_getfield(L, sub.result_index, args->vtagname);
+				if (lua_isnil(L, -1)) {
+					luaL_error(L, "Can't find value field in [%s]", args->tagname);
+				}
+				lua_settable(L, self->array_index);
+				lua_settop(L, sub.result_index);
+			} else {
+				lua_pushvalue(L, sub.key_index);
+				if (lua_isnil(L, -1)) {
+					luaL_error(L, "Can't find main index (tag=%d) in [%s]", args->mainindex, args->tagname);
+				}
+				lua_pushvalue(L, sub.result_index);
+				lua_settable(L, self->array_index);
+				lua_settop(L, sub.result_index-1);
 			}
-			lua_pushvalue(L, sub.result_index);
-			lua_settable(L, self->array_index);
-			lua_settop(L, sub.result_index-1);
 			return 0;
 		} else {
 			sub.mainindex_tag = -1;
@@ -524,6 +588,7 @@ ldecode(lua_State *L) {
 	self.deep = 0;
 	self.mainindex_tag = -1;
 	self.key_index = 0;
+	self.map_entry = 0;
 	r = sproto_decode(st, buffer, (int)sz, decode, &self);
 	if (r < 0) {
 		return luaL_error(L, "decode error");
@@ -679,7 +744,7 @@ lloadproto(lua_State *L) {
 }
 
 static void
-push_default(const struct sproto_arg *args, int array) {
+push_default(const struct sproto_arg *args, int table) {
 	lua_State *L = args->ud;
 	switch(args->type) {
 	case SPROTO_TINTEGER:
@@ -698,7 +763,7 @@ push_default(const struct sproto_arg *args, int array) {
 		lua_pushliteral(L, "");
 		break;
 	case SPROTO_TSTRUCT:
-		if (array) {
+		if (table) {
 			lua_pushstring(L, sproto_name(args->subtype));
 		} else {
 			lua_createtable(L, 0, 1);
diff --git a/lualib-src/sproto/sproto.c b/lualib-src/sproto/sproto.c
index 62c336e9..0b7f4132 100644
--- a/lualib-src/sproto/sproto.c
+++ b/lualib-src/sproto/sproto.c
@@ -6,7 +6,6 @@
 
 #include "sproto.h"
 
-#define SPROTO_TARRAY 0x80
 #define CHUNK_SIZE 1000
 #define SIZEOF_LENGTH 4
 #define SIZEOF_HEADER 2
@@ -20,6 +19,7 @@ struct field {
 	const char * name;
 	struct sproto_type * st;
 	int key;
+	int map; // interpreted two fields struct as map
 	int extra;
 };
 
@@ -206,6 +206,7 @@ import_field(struct sproto *s, struct field *f, const uint8_t * stream) {
 	f->name = NULL;
 	f->st = NULL;
 	f->key = -1;
+	f->map = -1;
 	f->extra = 0;
 
 	sz = todword(stream);
@@ -262,6 +263,10 @@ import_field(struct sproto *s, struct field *f, const uint8_t * stream) {
 		case 5:	// key
 			f->key = value;
 			break;
+		case 6: // map
+			if (value)
+				f->map = 1;
+			break;
 		default:
 			return NULL;
 		}
@@ -281,6 +286,8 @@ import_field(struct sproto *s, struct field *f, const uint8_t * stream) {
 		type 2 : integer
 		tag 3 : integer
 		array 4 : boolean
+		key 5 : integer
+		map 6 : boolean // Interpreted two fields struct as map when decoding
 	}
 	name 0 : string
 	fields 1 : *field
@@ -487,6 +494,32 @@ sproto_release(struct sproto * s) {
 	pool_release(&s->memory);
 }
 
+static const char *
+get_typename(int type, struct field *f) {
+	if (type == SPROTO_TSTRUCT) {
+		return f->st->name;
+	} else {
+		switch (type) {
+		case SPROTO_TINTEGER:
+			if (f->extra)
+				return "decimal";
+			else
+				return "integer";
+		case SPROTO_TBOOLEAN:
+			return "boolean";
+		case SPROTO_TSTRING:
+			if (f->extra == SPROTO_TSTRING_BINARY)
+				return "binary";
+			else
+				return "string";
+		case SPROTO_TDOUBLE:
+			return "double";
+		default:
+			return "invalid";
+		}
+	}
+}
+
 void
 sproto_dump(struct sproto *s) {
 	int i,j;
@@ -495,50 +528,30 @@ sproto_dump(struct sproto *s) {
 		struct sproto_type *t = &s->type[i];
 		printf("%s\n", t->name);
 		for (j=0;j<t->n;j++) {
-			char array[2] = { 0, 0 };
-			const char * type_name = NULL;
+			char container[2] = { 0, 0 };
+			const char * typename = NULL;
 			struct field *f = &t->f[j];
 			int type = f->type & ~SPROTO_TARRAY;
 			if (f->type & SPROTO_TARRAY) {
-				array[0] = '*';
+				container[0] = '*';
 			} else {
-				array[0] = 0;
+				container[0] = 0;
 			}
-			if (type == SPROTO_TSTRUCT) {
-				type_name = f->st->name;
-			} else {
-				switch(type) {
-				case SPROTO_TINTEGER:
-					if (f->extra) {
-						type_name = "decimal";
-					} else {
-						type_name = "integer";
-					}
-					break;
-				case SPROTO_TBOOLEAN:
-					type_name = "boolean";
-					break;
-				case SPROTO_TSTRING:
-					if (f->extra == SPROTO_TSTRING_BINARY)
-						type_name = "binary";
-					else
-						type_name = "string";
-					break;
-				default:
-					type_name = "invalid";
-					break;
-				}
-			}
-			printf("\t%s (%d) %s%s", f->name, f->tag, array, type_name);
+			typename = get_typename(type, f);
+			printf("\t%s (%d) %s%s", f->name, f->tag, container, typename);
 			if (type == SPROTO_TINTEGER && f->extra > 0) {
 				printf("(%d)", f->extra);
 			}
 			if (f->key >= 0) {
-				printf("[%d]", f->key);
+				printf(" key[%d]", f->key);
+				if (f->map >= 0) {
+					printf(" value[%d]", f->st->f[1].tag);
+				}
 			}
 			printf("\n");
 		}
 	}
+
 	printf("=== %d protocol ===\n", s->protocol_n);
 	for (i=0;i<s->protocol_n;i++) {
 		struct protocol *p = &s->proto[i];
@@ -840,6 +853,36 @@ encode_integer_array(sproto_callback cb, struct sproto_arg *args, uint8_t *buffe
 	return buffer;
 }
 
+static uint8_t *
+encode_array_object(sproto_callback cb, struct sproto_arg *args, uint8_t *buffer, int size, int *noarray) {
+	int sz;
+	*noarray = 0;
+	args->index = 1;
+	for (;;) {
+		if (size < SIZEOF_LENGTH)
+			return NULL;
+		size -= SIZEOF_LENGTH;
+		args->value = buffer + SIZEOF_LENGTH;
+		args->length = size;
+		sz = cb(args);
+		if (sz < 0) {
+			if (sz == SPROTO_CB_NIL) {
+				break;
+			}
+			if (sz == SPROTO_CB_NOARRAY) {	// no array, don't encode it
+				*noarray = 1;
+				break;
+			}
+			return NULL;	// sz == SPROTO_CB_ERROR
+		}
+		fill_size(buffer, sz);
+		buffer += SIZEOF_LENGTH+sz;
+		size -= sz;
+		++args->index;
+	}
+	return buffer;
+}
+
 static int
 encode_array(sproto_callback cb, struct sproto_arg *args, uint8_t *data, int size) {
 	uint8_t * buffer;
@@ -883,30 +926,16 @@ encode_array(sproto_callback cb, struct sproto_arg *args, uint8_t *data, int siz
 			++args->index;
 		}
 		break;
-	default:
-		args->index = 1;
-		for (;;) {
-			if (size < SIZEOF_LENGTH)
-				return -1;
-			size -= SIZEOF_LENGTH;
-			args->value = buffer+SIZEOF_LENGTH;
-			args->length = size;
-			sz = cb(args);
-			if (sz < 0) {
-				if (sz == SPROTO_CB_NIL) {
-					break;
-				}
-				if (sz == SPROTO_CB_NOARRAY)	// no array, don't encode it
-					return 0;
-				return -1;	// sz == SPROTO_CB_ERROR
-			}
-			fill_size(buffer, sz);
-			buffer += SIZEOF_LENGTH+sz;
-			size -=sz;
-			++args->index;
-		}
+	default: {
+		int noarray;
+		buffer = encode_array_object(cb, args, buffer, size, &noarray);
+		if (buffer == NULL)
+			return -1;
+		if (noarray)
+			return 0;
 		break;
 	}
+	}
 	sz = buffer - (data + SIZEOF_LENGTH);
 	return fill_size(data, sz);
 }
@@ -938,8 +967,14 @@ sproto_encode(const struct sproto_type *st, void * buffer, int size, sproto_call
 		args.subtype = f->st;
 		args.mainindex = f->key;
 		args.extra = f->extra;
+		args.ktagname = NULL;
+		args.vtagname = NULL;
 		if (type & SPROTO_TARRAY) {
-			args.type = type & ~SPROTO_TARRAY;
+			args.type = type & (~SPROTO_TARRAY);
+			if (f->map > 0) {
+				args.ktagname = f->st->f[0].name;
+				args.vtagname = f->st->f[1].name;
+			}
 			sz = encode_array(cb, &args, data, size);
 		} else {
 			args.type = type;
@@ -1168,13 +1203,20 @@ sproto_decode(const struct sproto_type *st, const void * data, int size, sproto_
 			continue;
 		args.tagname = f->name;
 		args.tagid = f->tag;
-		args.type = f->type & ~SPROTO_TARRAY;
+		args.type = f->type;
 		args.subtype = f->st;
 		args.index = 0;
 		args.mainindex = f->key;
 		args.extra = f->extra;
+		args.ktagname = NULL;
+		args.vtagname = NULL;
 		if (value < 0) {
 			if (f->type & SPROTO_TARRAY) {
+				args.type = f->type & (~SPROTO_TARRAY);
+				if (f->map > 0) {
+					args.ktagname = f->st->f[0].name;
+					args.vtagname = f->st->f[1].name;
+				}
 				if (decode_array(cb, &args, currentdata)) {
 					return -1;
 				}
diff --git a/lualib-src/sproto/sproto.h b/lualib-src/sproto/sproto.h
index 0ee5255b..4ff4952d 100644
--- a/lualib-src/sproto/sproto.h
+++ b/lualib-src/sproto/sproto.h
@@ -16,6 +16,9 @@ struct sproto_type;
 #define SPROTO_TDOUBLE 3
 #define SPROTO_TSTRUCT 4
 
+// container type
+#define SPROTO_TARRAY 0x80
+
 // sub type of string (sproto_arg.extra)
 #define SPROTO_TSTRING_STRING 0
 #define SPROTO_TSTRING_BINARY 1
@@ -46,9 +49,13 @@ struct sproto_arg {
 	struct sproto_type *subtype;
 	void *value;
 	int length;
-	int index;	// array base 1
+	int index;	// array base 1, negative value indicates that it is a empty array
 	int mainindex;	// for map
 	int extra; // SPROTO_TINTEGER: decimal ; SPROTO_TSTRING 0:utf8 string 1:binary
+
+	// When interpretd two fields struct as map, the following fields must not be NULL.
+	const char *ktagname;
+	const char *vtagname;
 };
 
 typedef int (*sproto_callback)(const struct sproto_arg *args);
diff --git a/lualib/sprotoparser.lua b/lualib/sprotoparser.lua
index 59eb0260..4335af8e 100644
--- a/lualib/sprotoparser.lua
+++ b/lualib/sprotoparser.lua
@@ -4,7 +4,7 @@ local table = require "table"
 local packbytes
 local packvalue
 
-if _VERSION == "Lua 5.4" then
+if _VERSION == "Lua 5.3" then
 	function packbytes(str)
 		return string.pack("<s4",str)
 	end
@@ -69,7 +69,7 @@ local word = alpha * alnum ^ 0
 local name = C(word)
 local typename = C(word * ("." * word) ^ 0)
 local tag = R"09" ^ 1 / tonumber
-local mainkey = "(" * blank0 * name * blank0 * ")"
+local mainkey = "(" * blank0 * C((word ^ 0)) * blank0 * ")"
 local decimal = "(" * blank0 * C(tag) * blank0 * ")"
 
 local function multipat(pat)
@@ -82,7 +82,7 @@ end
 
 local typedef = P {
 	"ALL",
-	FIELD = namedpat("field", (name * blanks * tag * blank0 * ":" * blank0 * (C"*")^-1 * typename * (mainkey +  decimal)^0)),
+	FIELD = namedpat("field", name * blanks * tag * blank0 * ":" * blank0 * (C"*")^-1 * typename * (mainkey + decimal)^0),
 	STRUCT = P"{" * multipat(V"FIELD" + V"TYPE") * P"}",
 	TYPE = namedpat("type", P"." * name * blank0 * V"STRUCT" ),
 	SUBPROTO = Ct((C"request" + C"response") * blanks * (typename + V"STRUCT")),
@@ -118,6 +118,11 @@ function convert.protocol(all, obj)
 	return result
 end
 
+local map_keytypes = {
+	integer = true,
+	string = true,
+}
+
 function convert.type(all, obj)
 	local result = {}
 	local typename = obj[1]
@@ -258,7 +263,8 @@ end
 		type 2 : integer
 		tag	3 :	integer
 		array 4	: boolean
-		key 5 : integer # If key exists, array must be true, and it's a map.
+		key 5 : integer # If key exists, array must be true
+		map 6 : boolean # Interpret two fields struct as map when decoding
 	}
 	name 0 : string
 	fields 1 : *field
@@ -282,7 +288,11 @@ local function packfield(f)
 	local strtbl = {}
 	if f.array then
 		if f.key then
-			table.insert(strtbl, "\6\0")  -- 6 fields
+			if f.map then
+				table.insert(strtbl, "\7\0")  -- 7 fields
+			else
+				table.insert(strtbl, "\6\0")  -- 6 fields
+			end
 		else
 			table.insert(strtbl, "\5\0")  -- 5 fields
 		end
@@ -305,9 +315,12 @@ local function packfield(f)
 	end
 	if f.array then
 		table.insert(strtbl, packvalue(1))	-- array = true (tag = 4)
-	end
-	if f.key then
-		table.insert(strtbl, packvalue(f.key)) -- key tag (tag = 5)
+		if f.key then
+			table.insert(strtbl, packvalue(f.key)) -- key tag (tag = 5)
+			if f.map then
+				table.insert(strtbl, packvalue(f.map)) -- map tag (tag = 6)
+			end
+		end
 	end
 	table.insert(strtbl, packbytes(f.name)) -- external object (name)
 	return packbytes(table.concat(strtbl))
@@ -334,10 +347,27 @@ local function packtype(name, t, alltypes)
 			tmp.type = nil
 		end
 		if f.key then
-			tmp.key = subtype.fields[f.key]
-			if not tmp.key then
+			assert(f.array)
+			if f.key == "" then
+				tmp.map = 1
+				local c = 0
+				local min_t = math.maxinteger
+				for n, t in pairs(subtype.fields) do
+					c = c + 1
+					if t.tag < min_t then
+						min_t = t.tag
+						f.key = n
+					end
+				end
+				if c ~= 2 then
+					error(string.format("Invalid map definition: %s, must only have two fields", tmp.name))
+				end
+			end
+			local stfield = subtype.fields[f.key]
+			if not stfield or not stfield.buildin then
 				error("Invalid map index :" .. f.key)
 			end
+			tmp.key = stfield.tag
 		else
 			tmp.key = nil
 		end
@@ -415,9 +445,10 @@ local function packgroup(t,p)
 	for idx, name in ipairs(alltypes) do
 		local fields = {}
 		for _, type_fields in ipairs(t[name]) do
-			if buildin_types[type_fields.typename] then
-				fields[type_fields.name] = type_fields.tag
-			end
+			fields[type_fields.name] = {
+				tag = type_fields.tag,
+				buildin = buildin_types[type_fields.typename]
+			}
 		end
 		alltypes[name] = { id = idx - 1, fields = fields }
 	end