1 files changed, 343 insertions, 0 deletions
diff --git a/src/lua-squarkdb.c b/src/lua-squarkdb.c
new file mode 100644
index 0000000..5a30848
--- /dev/null
+++ b/src/lua-squarkdb.c
@@ -0,0 +1,343 @@
+#include <string.h>
+
+#include <lua.h>
+#include <lualib.h>
+#include <lauxlib.h>
+
+#include <cmph.h>
+
+#include "filterdb.h"
+
+#define SQUARKDB_META	"squarkdb"
+
+struct sqdb *Lsqdb_checkarg(lua_State *L, int index)
+{
+	struct sqdb *db;
+
+	luaL_checktype(L, index, LUA_TUSERDATA);
+	db = (struct sqdb *) luaL_checkudata(L, index, SQUARKDB_META);
+	if (db == NULL)
+		luaL_typerror(L, index, SQUARKDB_META);
+	return db;
+}
+
+static int Lsqdb_new(lua_State *L)
+{
+	struct sqdb *db;
+	const char *fn;
+
+	fn = luaL_checklstring(L, 1, NULL);
+
+	db = (struct sqdb *) lua_newuserdata(L, sizeof(struct sqdb));
+	luaL_getmetatable(L, SQUARKDB_META);
+	lua_setmetatable(L, -2);
+
+	if (sqdb_create(db, fn) < 0)
+		luaL_error(L, "Failed to create SquarkDB file '%s'", fn);
+
+	return 1;
+}
+
+static int Lsqdb_destroy(lua_State *L)
+{
+	struct sqdb *db;
+
+	db = Lsqdb_checkarg(L, 1);
+	sqdb_close(db);
+
+	return 1;
+}
+
+
+struct ioa_data {
+	lua_State *main;
+	lua_State *thread;
+};
+
+static void ioa_rewind(void *data)
+{
+	struct ioa_data *ioa = (struct ioa_data *) data;
+
+	/* pop previous thread */
+	lua_pop(ioa->main, 1);
+
+	/* create a new lua thread */
+	ioa->thread = lua_newthread(ioa->main);
+	lua_pushvalue(ioa->main, -2);		/* copy function to top */
+	lua_xmove(ioa->main, ioa->thread, 1);	/* move function from L to NL */
+}
+
+static cmph_uint32 ioa_count_keys(void *data)
+{
+	struct ioa_data *ioa = (struct ioa_data *) data;
+	lua_State *NL;
+	cmph_uint32 cnt = 0;
+
+	NL = lua_newthread(ioa->main);
+	lua_pushvalue(ioa->main, -2);	/* copy function to top */
+	lua_xmove(ioa->main, NL, 1);	/* move function from L to NL */
+
+	do {
+		cnt++;
+		lua_settop(NL, 1);
+	} while (lua_resume(NL, 0) == LUA_YIELD);
+	ioa_rewind(data);
+
+	return cnt - 1;
+}
+
+static int ioa_read(void *data, char **key, cmph_uint32 *len)
+{
+	struct ioa_data *ioa = (struct ioa_data *) data;
+	lua_State *L = ioa->thread;
+	size_t l;
+
+	/* get next key from lua thread */
+	lua_settop(L, 1);
+	if (lua_resume(L, 0) != LUA_YIELD ||
+	    !lua_isstring(L, 1)) {
+		*key = NULL;
+		*len = 0;
+		return -1;
+	}
+
+	*key = (char *) lua_tolstring(L, 1, &l);
+	*len = l;
+
+	return l;
+}
+
+static void ioa_dispose(void *data, char *key, cmph_uint32 len)
+{
+	/* LUA takes care of garbage collection */
+}
+
+static int Lsqdb_hash(lua_State *L)
+{
+	struct sqdb *db;
+	void *ptr;
+	cmph_uint32 hash;
+	const char *key;
+	size_t keylen;
+
+	db = Lsqdb_checkarg(L, 1);
+	key = luaL_checklstring(L, 2, &keylen);
+
+	ptr = sqdb_section_get(db, SQDB_SECTION_INDEX_MPH, NULL);
+	hash = cmph_search_packed(ptr, key, keylen);
+
+	lua_pushinteger(L, hash);
+
+	return 1;
+}
+
+static int Lsqdb_generate_hash(lua_State *L)
+{
+	struct sqdb *db;
+	struct ioa_data ioa;
+	cmph_config_t *cfg;
+	cmph_t *cmph;
+	cmph_io_adapter_t io;
+
+	char *ptr;
+	cmph_uint32 packed;
+
+	db = Lsqdb_checkarg(L, 1);
+	luaL_argcheck(L, lua_isfunction(L, 2) && !lua_iscfunction(L, 2),
+		      2, "Lua function expected");
+
+	ioa.main	= L;
+	io.data		= &ioa;
+	io.nkeys	= ioa_count_keys(io.data);
+	io.read		= ioa_read;
+	io.dispose	= ioa_dispose;
+	io.rewind	= ioa_rewind;
+
+	cfg = cmph_config_new(&io);
+	if (cfg == NULL)
+		luaL_error(L, "Failed to create CMPH config");
+
+	cmph_config_set_algo(cfg, CMPH_CHD);
+	cmph = cmph_new(cfg);
+	cmph_config_destroy(cfg);
+
+	if (cmph == NULL)
+		luaL_error(L, "Failed to create minimal perfect hash");
+
+	packed = cmph_packed_size(cmph);
+	ptr = sqdb_section_create(db, SQDB_SECTION_INDEX_MPH, packed);
+	if (ptr == NULL) {
+		cmph_destroy(cmph);
+		luaL_error(L, "Unable to allocation MPH section from SquarkDB");
+	}
+
+	cmph_pack(cmph, ptr);
+	cmph_destroy(cmph);
+
+	lua_pushinteger(L, io.nkeys);
+	lua_pushinteger(L, packed);
+
+	return 2;
+}
+
+static int Lsqdb_create_index(lua_State *L)
+{
+	struct sqdb *db;
+	lua_Integer num_entries;
+	void *ptr;
+
+	db = Lsqdb_checkarg(L, 1);
+	num_entries = luaL_checkinteger(L, 2);
+
+	ptr = sqdb_section_create(db, SQDB_SECTION_INDEX, sizeof(struct sqdb_index_entry) * num_entries);
+	if (ptr == NULL)
+		luaL_error(L, "Failed to create INDEX section");
+
+	return 0;
+}
+
+static int Lsqdb_assign_index(lua_State *L)
+{
+	struct sqdb *db;
+	size_t size;
+	lua_Integer idx;
+	struct sqdb_index_entry *ptr;
+
+	db = Lsqdb_checkarg(L, 1);
+	idx = luaL_checkinteger(L, 2);
+
+	ptr = sqdb_section_get(db, SQDB_SECTION_INDEX, &size);
+	if (size <= 0 || idx * sizeof(struct sqdb_index_entry) >= size)
+		luaL_error(L, "Bad index assignment (idx=%d, section size=%d)", idx, size);
+
+	ptr += idx;
+	if (ptr->component != 0)
+		luaL_error(L, "Index entry %d has been already assigned", idx);
+
+	ptr->category = luaL_checkinteger(L, 3);
+	ptr->has_subdomains = lua_toboolean(L, 4);
+	ptr->has_paths = lua_toboolean(L, 5);
+	ptr->component = luaL_checkinteger(L, 6);
+	ptr->parent = luaL_checkinteger(L, 7);
+
+	return 0;
+}
+
+static int Lsqdb_map_strings(lua_State *L)
+{
+	struct sqdb *db;
+	const char *str;
+	unsigned char *ptr;
+	size_t len, total, pos;
+
+	db = Lsqdb_checkarg(L, 1);
+	luaL_checktype(L, 2, LUA_TTABLE);
+
+	/* go through the table and count total amount of data */
+	total = 0;
+	lua_pushnil(L);
+	while (lua_next(L, 2) != 0) {
+		str = luaL_checklstring(L, -2, &len);
+		total += len;
+		if (len >= (1 << SQDB_LENGTH_BITS))
+			total++;
+		lua_pop(L, 1);
+	}
+
+	/* create string literal section */
+	ptr = sqdb_section_create(db, SQDB_SECTION_STRINGS, total);
+	if (ptr == NULL)
+		luaL_error(L, "Failed to create string literal section (%d bytes)", total);
+
+	/* populate string literals and return their indices */
+	pos = 0;
+	lua_pushnil(L);
+	while (lua_next(L, 2) != 0) {
+		str = lua_tolstring(L, -2, &len);
+		lua_pop(L, 1);
+
+		/* table[key] = encoded_string_pointer */
+		lua_pushvalue(L, -1);
+		if (len >= (1 << SQDB_LENGTH_BITS)) {
+			lua_pushinteger(L, pos << SQDB_LENGTH_BITS);
+			ptr[pos++] = len;
+		} else {
+			lua_pushinteger(L, (pos << SQDB_LENGTH_BITS) + len);
+		}
+		memcpy(&ptr[pos], str, len);
+		pos += len;
+
+		lua_rawset(L, 2);
+	}
+
+	return 0;
+}
+
+static int Lsqdb_write_section(lua_State *L)
+{
+	struct sqdb *db;
+	uint32_t *ptr;
+	const char *section;
+	int i, tbllen, si = -1;
+
+	db = Lsqdb_checkarg(L, 1);
+	section = luaL_checkstring(L, 2);
+	luaL_checktype(L, 3, LUA_TTABLE);
+	tbllen = lua_objlen(L, 3);
+
+	for (i = 0; sqdb_section_names[i] && i < SQDB_SECTION_MAX; i++) {
+		if (strcmp(sqdb_section_names[i], section) == 0) {
+			si = 0;
+			break;
+		}
+	}
+	if (si < 0)
+		luaL_error(L, "Section name '%s' is invalid", section);
+
+	ptr = sqdb_section_create(db, i, sizeof(uint32_t) * tbllen);
+	if (ptr == NULL)
+		luaL_error(L, "Failed to create section '%s'", section);
+
+	for (i = 0; i < tbllen; i++) {
+		lua_rawgeti(L, 3, i + 1);
+		ptr[i] = lua_tointeger(L, -1);
+		lua_pop(L, 1);
+	}
+
+	return 0;
+}
+
+static const luaL_reg sqdb_meta_methods[] = {
+	{ "__gc",		Lsqdb_destroy },
+	{ NULL,			NULL }
+};
+
+static const luaL_reg squarkdb_methods[] = {
+	{ "new",		Lsqdb_new },
+	{ "hash",		Lsqdb_hash },
+	{ "generate_hash",	Lsqdb_generate_hash },
+	{ "create_index",	Lsqdb_create_index },
+	{ "assign_index",	Lsqdb_assign_index },
+	{ "map_strings",	Lsqdb_map_strings },
+	{ "write_section",	Lsqdb_write_section },
+	{ NULL,			NULL }
+};
+
+LUALIB_API int luaopen_squarkdb(lua_State *L)
+{
+	/* Register squarkdb library */
+	luaI_openlib(L, "squarkdb", squarkdb_methods, 0);
+
+	/* And metatable for it */
+	luaL_newmetatable(L, SQUARKDB_META);
+	luaI_openlib(L, NULL, sqdb_meta_methods, 0);
+	lua_pushliteral(L, "__index");
+	lua_pushvalue(L, -3);
+	lua_rawset(L, -3);
+	lua_pushliteral(L, "__metatable");
+	lua_pushvalue(L, -3);
+	lua_rawset(L, -3);
+	lua_pop(L, 1);
+
+	return 1;
+}