path: root/lmdb.c

#include <stdlib.h>
#include <stdbool.h>
#include <lmdb.h>
#include <lua.h>
#include <lauxlib.h>
#include <assert.h>

struct handle {
	bool closed;
	bool has_child;
	void *obj;
};

#define reterr(e) \
	do { \
		int __res = (e); if (__res != 0) { \
			lua_pushnil(L); \
			lua_pushstring(L, mdb_strerror(__res)); \
			lua_pushinteger(L, __res); \
			return 3; \
		} \
	} while (0)
#define asserr(e) \
	do { \
		int __res = (e); if (__res != 0) \
			luaL_error(L, "%s", mdb_strerror(__res)); \
	} while (0)

static void get_parent(lua_State *L, int index);
static void set_parent(lua_State *L, int index);


static void check_valid(lua_State *L, struct handle *h) {
	if (h->closed) luaL_error(L, "attempt to use expired handle");
}

static struct {const char *name; int flag;} env_flags[] = {
	{"rdonly", MDB_RDONLY},
	{"nosubdir", MDB_NOSUBDIR},
	{"nosync", MDB_NOSYNC},
	{"nometasync", MDB_NOMETASYNC},
	{"nordahead", MDB_NORDAHEAD},
	{"nolock", MDB_NOLOCK},
	{0},
};

static int version(lua_State *L) {
	int maj, min, patch;
	mdb_version(&maj, &min, &patch);
	lua_pushinteger(L, maj);
	lua_pushinteger(L, min);
	lua_pushinteger(L, patch);
	return 3;
}

#if LUA_VER == 51
lua_Integer lua_tointegerx(lua_State *L, int index, int *isnum) {
	if (isnum != NULL)
		*isnum = lua_isnumber(L,index);
	return lua_tointeger(L, index);
}
#endif

static int env_open(lua_State *L) {
	lua_settop(L, 2);
	luaL_argcheck(L, lua_isstring(L, 1), 1, "expected string");
	luaL_argcheck(L,
		lua_isnil(L, 2) || lua_istable(L, 2), 2, "expected table or nil");

	int flags = 0;
	int maxdbs = 0;
	int maxdbs_supplied = 0;
	int mapsize = 0;
	int mapsize_supplied = 0;
	if (!lua_isnil(L, 2)) {
		for (int i = 0; env_flags[i].name != NULL; i++) {
			lua_getfield(L, 2, env_flags[i].name);
			if (lua_toboolean(L, -1)) flags |= env_flags[i].flag;
			lua_pop(L, 1);
		}
		lua_getfield(L, 2, "maxdbs");
		maxdbs = lua_tointegerx(L, -1, &maxdbs_supplied);
		lua_getfield(L, 2, "mapsize");
		mapsize = lua_tointegerx(L, -1, &mapsize_supplied);
	}

	flags |= MDB_NOTLS;

	MDB_env *env;
	reterr(mdb_env_create(&env));
	if (maxdbs_supplied) mdb_env_set_maxdbs(env, maxdbs);
	if (mapsize_supplied)
		mdb_env_set_mapsize(env, mapsize);
	else
		// lmdb default mapsize is 1MiB which is tiny
		// so we use a default of 1GiB instead
		mdb_env_set_mapsize(env, 1073741824);

	int result = mdb_env_open(env, lua_tostring(L, 1), flags, 0664);
	if (result != 0) {
		mdb_env_close(env);
		reterr(result);
	}

	struct handle *ud = lua_newuserdata(L, sizeof(struct handle));
	ud->closed = false;
	ud->has_child = false;
	ud->obj = env;
	luaL_getmetatable(L, "lmdb.env");
	lua_setmetatable(L, -2);
	return 1;
};

static int env_close(lua_State *L) {
	lua_settop(L, 1);
	struct handle *ud = luaL_checkudata(L, 1, "lmdb.env");
	lua_pushboolean(L, true);
	if (ud->closed) return 1;
	mdb_env_close(ud->obj);
	ud->closed = true;
	return 1;
}

static int env_sync(lua_State *L) {
	lua_settop(L, 2);
	struct handle *ud = luaL_checkudata(L, 1, "lmdb.env");
	check_valid(L, ud);
	reterr(mdb_env_sync(ud->obj, lua_toboolean(L, 2)));
	lua_pushboolean(L, true);
	return 1;
}

static int env_copy(lua_State *L) {
	lua_settop(L, 2);
	struct handle *ud = luaL_checkudata(L, 1, "lmdb.env");
	check_valid(L, ud);
	luaL_argcheck(L, lua_isstring(L, 2), 2, "expected string");
	reterr(mdb_env_copy(ud->obj, lua_tostring(L, 2)));
	lua_pushboolean(L, true);
	return 1;
}

static MDB_env *get_env(lua_State *L, int idx) {
	lua_pushvalue(L, idx);
	struct handle *ud = NULL;
	while (get_parent(L, -1), lua_type(L, -1) == LUA_TUSERDATA) {
		ud = lua_touserdata(L, -1);
		if (ud->closed) return NULL;
		check_valid(L, ud);
		lua_remove(L, -2);
	}
	lua_pop(L, 2);
	return ud->obj;
}

static MDB_env *check_env(lua_State *L, int idx) {
	MDB_env *env = get_env(L, idx);
	if (env == NULL) luaL_error(L, "attempt to use expired handle");
	return env;
}

static int txn_begin(lua_State *L, bool nested) {
	lua_settop(L, 2);
	struct handle *ud = luaL_checkudata(L, 1, !nested ? "lmdb.env" : "lmdb.txn");
	check_valid(L, ud);
	if (ud->has_child)
		luaL_error(L, "cannot create new transaction while another exists");

	MDB_env *env;
	MDB_txn *parent = NULL;
	if (!nested) env = ud->obj;
	else {
		parent = ud->obj;
		env = check_env(L, 1);
	}
	
	MDB_txn *txn;
	reterr(mdb_txn_begin(
		env, parent, lua_toboolean(L, 2) || nested ? 0 : MDB_RDONLY, &txn));

	struct handle *tud = lua_newuserdata(L, sizeof(struct handle));
	tud->closed = false;
	tud->has_child = false;
	tud->obj = txn;
	luaL_getmetatable(L, "lmdb.txn");
	lua_setmetatable(L, -2);
	// store the supplied transaction/env as the parent
	lua_pushvalue(L, 1);
	set_parent(L, -2);

	ud->has_child = true;
	return 1;
}

static int env_txn_begin(lua_State *L) {
	return txn_begin(L, false);
}

static int txn_txn_begin(lua_State *L) {
	return txn_begin(L, true);
}

static int txn_abort(lua_State *L) {
	lua_settop(L, 1);
	struct handle *ud = luaL_checkudata(L, 1, "lmdb.txn");
	if (ud->closed || get_env(L, 1) == NULL) {
		lua_pushboolean(L, true);
		return 1;
	}
	mdb_txn_abort(ud->obj);
	ud->closed = true;

	get_parent(L, 1);
	struct handle *pud = lua_touserdata(L, -1);
	pud->has_child = false;

	lua_pushboolean(L, true);
	return 1;
}

static int txn_commit(lua_State *L) {
	lua_settop(L, 1);
	struct handle *ud = luaL_checkudata(L, 1, "lmdb.txn");
	check_valid(L, ud);
	check_env(L, 1);
	ud->closed = true;

	get_parent(L, 1);
	struct handle *pud = lua_touserdata(L, -1);
	pud->has_child = false;

	reterr(mdb_txn_commit(ud->obj));
	lua_pushboolean(L, true);
	return 1;
}

static int db_open(lua_State *L) {
	lua_settop(L, 3);
	struct handle *ud = luaL_checkudata(L, 1, "lmdb.txn");
	check_valid(L, ud);
	check_env(L, 1);
	luaL_argcheck(L,
		lua_isnil(L, 2) || lua_isstring(L, 2), 2, "expected string or nil");

	MDB_dbi dbi;
	reterr(mdb_dbi_open(ud->obj,
		lua_tostring(L, 2), lua_toboolean(L, 3) ? MDB_CREATE : 0, &dbi));
	*(MDB_dbi *) lua_newuserdata(L, sizeof(MDB_dbi)) = dbi;
	luaL_getmetatable(L, "lmdb.db");
	lua_setmetatable(L, -2);
	// transaction is parent of DB handle
	lua_pushvalue(L, 1);
	set_parent(L, -2);
	return 1;
}

static MDB_val toval(lua_State *L, int idx) {
	MDB_val val;
	val.mv_data = lua_tolstring(L, idx, &val.mv_size);
	return val;
}

static void fromval(lua_State *L, MDB_val val) {
	lua_pushlstring(L, val.mv_data, val.mv_size);
}

static int db_get(lua_State *L) {
	lua_settop(L, 2);
	MDB_dbi dbi = *(MDB_dbi *) luaL_checkudata(L, 1, "lmdb.db");
	check_env(L, 1);
	luaL_argcheck(L, lua_isstring(L, 2), 2, "expected string");

	get_parent(L, 1);
	struct handle *tud = lua_touserdata(L, -1);
	MDB_val key = toval(L, 2);
	MDB_val data;
	int result = mdb_get(tud->obj, dbi, &key, &data);
	if (result == MDB_NOTFOUND) {
		lua_pushnil(L);
		return 1;
	} else asserr(result);
	fromval(L, data);
	return 1;
}

static int db_put(lua_State *L) {
	lua_settop(L, 3);
	MDB_dbi dbi = *(MDB_dbi *) luaL_checkudata(L, 1, "lmdb.db");
	check_env(L, 1);
	luaL_argcheck(L, lua_isstring(L, 2), 2, "expected string");
	luaL_argcheck(
		L, lua_isnil(L, 3) || lua_isstring(L, 3), 3, "expected string or nil");

	get_parent(L, 1);
	struct handle *tud = lua_touserdata(L, -1);
	MDB_val key = toval(L, 2);
	if (!lua_isnil(L, 3)) {
		MDB_val data = toval(L, 3);
		asserr(mdb_put(tud->obj, dbi, &key, &data, 0));
	} else asserr(mdb_del(tud->obj, dbi, &key, NULL));
	return 0;
}

static int db_len(lua_State *L) {
	lua_settop(L, 1);
	MDB_dbi dbi = *(MDB_dbi *) luaL_checkudata(L, 1, "lmdb.db");
	check_env(L, 1);

	get_parent(L, 1);
	struct handle *tud = lua_touserdata(L, -1);
	struct MDB_stat stat;
	asserr(mdb_stat(tud->obj, dbi, &stat));
	lua_pushinteger(L, stat.ms_entries);
	return 1;
}

static int db_next(lua_State *L) {
	lua_settop(L, 2);
	MDB_dbi dbi = *(MDB_dbi *) luaL_checkudata(L, 1, "lmdb.db");
	check_env(L, 1);
	luaL_argcheck(
		L, lua_isnil(L, 2) || lua_isstring(L, 2), 2, "expected string or nil");

	get_parent(L, 1);
	struct handle *tud = lua_touserdata(L, -1);
	MDB_cursor *curs;
	asserr(mdb_cursor_open(tud->obj, dbi, &curs));

	MDB_val key;
	MDB_val data;
	int result;
	if (!lua_isnil(L, 2)) {
		key = toval(L, 2);
		asserr(mdb_cursor_get(curs, &key, &data, MDB_SET));
		result = mdb_cursor_get(curs, &key, &data, MDB_NEXT);
	} else
		result = mdb_cursor_get(curs, &key, &data, MDB_FIRST);
	if (result == MDB_NOTFOUND) {
		lua_pushnil(L);
		return 1;
	} else asserr(result);

	fromval(L, key);
	fromval(L, data);
	mdb_cursor_close(curs);
	return 2;
}

static int db_pairs(lua_State *L) {
	lua_settop(L, 1);
	luaL_checkudata(L, 1, "lmdb.db");
	lua_pushcfunction(L, db_next);
	lua_pushvalue(L, 1);
	lua_pushnil(L);
	return 3;
}

static int db_drop(lua_State *L) {
	lua_settop(L, 2);
	struct handle *ud = luaL_checkudata(L, 1, "lmdb.txn");
	check_valid(L, ud);
	check_env(L, 1);
	luaL_argcheck(L,
		lua_isnil(L, 2) || lua_isstring(L, 2), 2, "expected string or nil");

	MDB_dbi dbi;
	reterr(mdb_dbi_open(ud->obj, lua_tostring(L, 2), 0, &dbi));
	reterr(mdb_drop(ud->obj, dbi, !lua_isnil(L, 2)));
	lua_pushboolean(L, true);
	return 1;
}

static const struct luaL_Reg lmdb[] = {
	{"open", env_open},
	{"version", version},
	{"next", db_next},
	{"pairs",db_pairs},
	{NULL, NULL},
};

static const struct luaL_Reg lmdb_env[] = {
	{"txn_begin", env_txn_begin},
	{"copy", env_copy},
	{"sync", env_sync},
	{"close", env_close},
	{"abort", env_close},
	{"__close", env_close},
	{"__gc", env_close},
	{NULL, NULL},
};

static const struct luaL_Reg lmdb_txn[] = {
	{"open", db_open},
	{"abort", txn_abort},
	{"commit", txn_commit},
	{"txn_begin", txn_txn_begin},
	{"drop", db_drop},
	{"__close", txn_commit},
	{"__gc", txn_abort},
	{NULL, NULL},
};

static const struct luaL_Reg lmdb_db[] = {
	{"__index", db_get},
	{"__newindex", db_put},
	{"__pairs", db_pairs},
	{"__len", db_len},
	{NULL, NULL},
};

static void mt(lua_State *L, const char *name, const struct luaL_Reg l[]) {
	luaL_newmetatable(L, name);
#if LUA_VER == 51
	luaL_register(L, NULL, l);
#else
	luaL_setfuncs(L, l, 0);
#endif
	lua_getfield(L, -1, "__index");
	if (lua_isnil(L, -1)){
		lua_pop(L, 1);
		lua_setfield(L, -1, "__index");
	} else {
		lua_pop(L, 2);
	}
}

static void parentmap(lua_State *L) {
	lua_getfield(L, LUA_REGISTRYINDEX, "lmdb.parentmap");
}

static void get_parent(lua_State *L, int index) {
	lua_pushvalue(L, index); // k ... k 
	parentmap(L); // k ... k t
	lua_insert(L, -2); // k ... t k 
	lua_gettable(L, -2); // k ... t v
	lua_remove(L, -2); // k ... v 
}
static void set_parent(lua_State *L, int index) {
	// references to `index` need to be before anything changes the stack, 
	// otherwise indices will be off
	// k ... v 
	lua_pushvalue(L, index); // k ... v k
	parentmap(L); // k ... v k t
	lua_insert(L, -3); // k ... t v k 
	lua_insert(L, -2); // k ... t k v
	lua_settable(L, -3); // k ... t 
	lua_pop(L, 1); // k ...
}


int luaopen_lmdb(lua_State *L) {
	mt(L, "lmdb.env", lmdb_env);
	mt(L, "lmdb.txn", lmdb_txn);
	mt(L, "lmdb.db", lmdb_db);

	lua_newtable(L);
	lua_newtable(L);
	lua_pushliteral(L, "k");
	lua_setfield(L, -2, "__mode");
	lua_setmetatable(L, -2);
	lua_setfield(L, LUA_REGISTRYINDEX, "lmdb.parentmap");

#if LUA_VER == 51
	lua_newtable(L);
	luaL_register(L, NULL, lmdb);
#else
	luaL_newlib(L, lmdb);
#endif
	return 1;
}