[Minor] Allow to store buckets of numbers in lua mempool

1 files changed, 71 insertions, 1 deletions

diff --git a/src/lua/lua_mempool.c b/src/lua/lua_mempool.c
index 569fe0eb6..a6ef9f2c2 100644
--- a/src/lua/lua_mempool.c
+++ b/src/lua/lua_mempool.c
@@ -67,6 +67,15 @@ LUA_FUNCTION_DEF (mempool, suggest_size);
  */
 LUA_FUNCTION_DEF (mempool, set_variable);
 /***
+ * @method mempool:set_bucket(name, num_values, [value1...valuen]|[table])
+ * Stores a variable bucket of numbers where the first number is number of elements to pack
+ * and then there should be either n numeric values or a plain table of numeric values
+ * @param {string} name variable's name to set
+ * @param {number} num_values number of variables in the bucket
+ * @param {table|list} values values
+ */
+LUA_FUNCTION_DEF (mempool, set_bucket);
+/***
  * @method mempool:get_variable(name[, type])
  * Unpacks mempool variable to lua If `type` is not specified, then a variable is
  * assumed to be zero-terminated C string. Otherwise, `type` is a comma separated (spaces are ignored)
@@ -78,6 +87,7 @@ LUA_FUNCTION_DEF (mempool, set_variable);
  * - `int`: unpack a single integer
  * - `int64`: unpack 64-bits integer
  * - `boolean`: unpack boolean
+ * - `bucket`: bucket of numbers represented as a table
  * @param {string} name variable's name to get
  * @param {string} type list of types to be extracted
  * @return {variable list} list of variables extracted (but **not** a table)
@@ -104,6 +114,7 @@ static const struct luaL_reg mempoollib_m[] = {
 	LUA_INTERFACE_DEF (mempool, stat),
 	LUA_INTERFACE_DEF (mempool, suggest_size),
 	LUA_INTERFACE_DEF (mempool, set_variable),
+	LUA_INTERFACE_DEF (mempool, set_bucket),
 	LUA_INTERFACE_DEF (mempool, get_variable),
 	LUA_INTERFACE_DEF (mempool, has_variable),
 	LUA_INTERFACE_DEF (mempool, delete_variable),
@@ -245,6 +256,47 @@ lua_mempool_suggest_size (lua_State *L)
 	return 1;
 }
 
+struct lua_numbers_bucket {
+	guint nelts;
+	gdouble elts[0];
+};
+
+static int
+lua_mempool_set_bucket (lua_State *L)
+{
+	struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
+	const gchar *var = luaL_checkstring (L, 2);
+	struct lua_numbers_bucket *bucket;
+	gint nelts = luaL_checknumber (L, 3), i;
+
+	if (var && nelts > 0) {
+		bucket = rspamd_mempool_alloc (mempool,
+				sizeof (*bucket) + sizeof (gdouble) * nelts);
+		bucket->nelts = nelts;
+
+		if (lua_type (L, 4) == LUA_TTABLE) {
+			/* Table version */
+			for (i = 1; i <= nelts; i ++) {
+				lua_rawgeti (L, 4, i);
+				bucket->elts[i - 1] = lua_tonumber (L, -1);
+				lua_pop (L, 1);
+			}
+		}
+		else {
+			for (i = 0; i <= nelts; i ++) {
+				bucket->elts[i] = lua_tonumber (L, 4 + i);
+			}
+		}
+
+		rspamd_mempool_set_variable (mempool, var, bucket, NULL);
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 0;
+}
+
 static int
 lua_mempool_set_variable (lua_State *L)
 {
@@ -276,6 +328,10 @@ lua_mempool_set_variable (lua_State *L)
 				(void)lua_tolstring (L, i, &slen);
 				len += slen + 1;
 			}
+			else if (type == LUA_TTABLE) {
+				/* We assume it as a bucket of numbers so far */
+
+			}
 			else {
 				msg_err ("cannot handle lua type %s", lua_typename (L, type));
 			}
@@ -330,8 +386,9 @@ lua_mempool_get_variable (lua_State *L)
 	struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
 	const gchar *var = luaL_checkstring (L, 2);
 	const gchar *type = NULL, *pt;
+	struct lua_numbers_bucket *bucket;
 	gchar *value, *pv;
-	guint len, nvar, slen;
+	guint len, nvar, slen, i;
 
 	if (mempool && var) {
 		value = rspamd_mempool_get_variable (mempool, var);
@@ -380,6 +437,19 @@ lua_mempool_get_variable (lua_State *L)
 						lua_pushlstring (L, st->str, st->len);
 						pv += sizeof (GString *);
 					}
+					else if (len == sizeof ("bucket") - 1 &&
+							g_ascii_strncasecmp (pt, "bucket", len) == 0) {
+						bucket = (struct lua_numbers_bucket *)pv;
+						lua_createtable (L, bucket->nelts, 0);
+
+						for (i = 0; i < bucket->nelts; i ++) {
+							lua_pushnumber (L, bucket->elts[i]);
+							lua_rawseti (L, -1, i + 1);
+						}
+
+						pv += sizeof (struct lua_numbers_bucket) +
+								bucket->nelts * sizeof (gdouble);
+					}
 					else {
 						msg_err ("unknown type for get_variable: %s", pt);
 						lua_pushnil (L);