nilmdb/nilmdb/server/rocket.c

#include <Python.h>
#include <structmember.h>
#include <endian.h>

/* Somewhat arbitrary, just so we can use fixed sizes for strings
   etc. */
static const int MAX_LAYOUT_COUNT = 64;

typedef union {
	int8_t i8[8];
	uint8_t u8[8];
	int16_t i16[4];
	uint16_t u16[4];
	int32_t i32[2];
	uint32_t u32[2];
	int64_t i64[1];
	uint64_t u64[1];
	float f[2];
	double d[1];
} union_t;

typedef enum {
	LAYOUT_TYPE_NONE,
	LAYOUT_TYPE_INT8,
	LAYOUT_TYPE_UINT8,
	LAYOUT_TYPE_INT16,
	LAYOUT_TYPE_UINT16,
	LAYOUT_TYPE_INT32,
	LAYOUT_TYPE_UINT32,
	LAYOUT_TYPE_INT64,
	LAYOUT_TYPE_UINT64,
	LAYOUT_TYPE_FLOAT32,
	LAYOUT_TYPE_FLOAT64,
} layout_type_t;

struct {
	char *string;
	layout_type_t layout;
	int size;
} type_lookup[] = {
	{ "int8",    LAYOUT_TYPE_INT8,    1 },
	{ "uint8",   LAYOUT_TYPE_UINT8,   1 },
	{ "int16",   LAYOUT_TYPE_INT16,   2 },
	{ "uint16",  LAYOUT_TYPE_UINT16,  2 },
	{ "int32",   LAYOUT_TYPE_INT32,   4 },
	{ "uint32",  LAYOUT_TYPE_UINT32,  4 },
	{ "int64",   LAYOUT_TYPE_INT64,   8 },
	{ "uint64",  LAYOUT_TYPE_UINT64,  8 },
	{ "float32", LAYOUT_TYPE_FLOAT32, 4 },
	{ "float64", LAYOUT_TYPE_FLOAT64, 8 },
	{ NULL }
};

typedef struct {
	PyObject_HEAD
	layout_type_t layout_type;
	int layout_count;
	int binary_size;
} Rocket;

static void Rocket_dealloc(Rocket *self)
{
	self->ob_type->tp_free((PyObject *)self);
}

static PyObject *Rocket_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
	Rocket *self;

	self = (Rocket *)type->tp_alloc(type, 0);
	if (!self)
		return NULL;
	self->layout_type = LAYOUT_TYPE_NONE;
	self->layout_count = 0;
	self->binary_size = 0;
	return (PyObject *)self;
}

static int Rocket_init(Rocket *self, PyObject *args, PyObject *kwds)
{
	const char *layout;
	static char *kwlist[] = { "layout", NULL };
	if (!PyArg_ParseTupleAndKeywords(args, kwds, "s", kwlist, &layout))
		return -1;
	if (!layout)
		return -1;

	const char *under;
	char *tmp;
	under = strchr(layout, '_');
	if (!under) {
		PyErr_SetString(PyExc_ValueError, "no such layout: "
				"badly formatted string");
		return -1;
	}
	self->layout_count = strtoul(under+1, &tmp, 10);
	if (self->layout_count < 1 || *tmp != '\0') {
		PyErr_SetString(PyExc_ValueError, "no such layout: "
				"bad count");
		return -1;
	}
	if (self->layout_count >= MAX_LAYOUT_COUNT) {
		PyErr_SetString(PyExc_ValueError, "no such layout: "
				"count too high");
		return -1;
	}

	int i;
	for (i = 0; type_lookup[i].string; i++)
		if (strncmp(layout, type_lookup[i].string, under-layout) == 0)
			break;
	if (!type_lookup[i].string) {
		PyErr_SetString(PyExc_ValueError, "no such layout: "
				"bad data type");
		return -1;
	}
	self->layout_type = type_lookup[i].layout;
	self->binary_size = 8 + (type_lookup[i].size * self->layout_count);

	return 0;
}

FILE *PyFile_AsFile(PyObject *file)
{
	PyObject *result;
	int fd;
	result = PyObject_CallMethod(file, "fileno", NULL);
	if (result == NULL)
		return NULL;
	fd = PyInt_AsLong(result);
	if (fd < 0)
		return NULL;
	return fdopen(fd, "a+b");
}

static inline void write_pyobject(FILE *out, PyObject *val, layout_type_t type)
{
	union_t u;
	int ret = 0;

	switch (type) {
#define CASE(type, pyconvert, pytype, disktype, htole, bytes)	     \
		case LAYOUT_TYPE_##type:			     \
			u.pytype[0] = pyconvert(val);		     \
			if (PyErr_Occurred())			     \
				return;				     \
			u.disktype[0] = htole(u.disktype[0]);	     \
			ret = fwrite(&u.disktype[0], bytes, 1, out); \
			break
		CASE(INT8,   PyInt_AsLong, i8,  u8,         , 1);
		CASE(UINT8,  PyInt_AsLong, u8,  u8,         , 1);
		CASE(INT16,  PyInt_AsLong, i16, u16, htole16, 2);
		CASE(UINT16, PyInt_AsLong, u16, u16, htole16, 2);
		CASE(INT32,  PyInt_AsLong, i32, u32, htole32, 4);
		CASE(UINT32, PyInt_AsLong, u32, u32, htole32, 4);
		CASE(INT64,  PyInt_AsLong, i64, u64, htole64, 8);
		CASE(UINT64, PyInt_AsLong, u64, u64, htole64, 8);
		CASE(FLOAT32, PyFloat_AsDouble, f, u32, htole32, 4);
		CASE(FLOAT64, PyFloat_AsDouble, d, u64, htole64, 8);
#undef CASE
	default:
		PyErr_SetString(PyExc_TypeError, "unknown type");
		return;
	}
	if (ret <= 0)
		PyErr_SetFromErrno(PyExc_OSError);
}

static inline void *read_pyobject(FILE *in, layout_type_t type)
{
	union_t u;

	switch (type) {
#define CASE(type, pyconvert, pytype, disktype, letoh, bytes)		\
		case LAYOUT_TYPE_##type:				\
			if (fread(&u.disktype[0], bytes, 1, in) <= 0)	\
				break;					\
			u.disktype[0] = letoh(u.disktype[0]);		\
			return pyconvert(u.pytype[0]);			\
			break
		CASE(INT8,   PyInt_FromLong, i8,  u8,         , 1);
		CASE(UINT8,  PyInt_FromLong, u8,  u8,         , 1);
		CASE(INT16,  PyInt_FromLong, i16, u16, le16toh, 2);
		CASE(UINT16, PyInt_FromLong, u16, u16, le16toh, 2);
		CASE(INT32,  PyInt_FromLong, i32, u32, le32toh, 4);
		CASE(UINT32, PyInt_FromLong, u32, u32, le32toh, 4);
		CASE(INT64,  PyInt_FromLong, i64, u64, le64toh, 8);
		CASE(UINT64, PyInt_FromLong, u64, u64, le64toh, 8);
		CASE(FLOAT32, PyFloat_FromDouble, f, u32, le32toh, 4);
		CASE(FLOAT64, PyFloat_FromDouble, d, u64, le64toh, 8);
#undef CASE
	default:
		PyErr_SetString(PyExc_TypeError, "unknown type");
		return NULL;
	}
	PyErr_SetFromErrno(PyExc_OSError);
	return NULL;
}

static PyObject *Rocket_append_list(Rocket *self, PyObject *args)
{
	PyObject *file, *list;
	FILE *out;
	if (!PyArg_ParseTuple(args, "OO:append_list", &file, &list))
		return NULL;
	if ((out = PyFile_AsFile(file)) == NULL)
		return NULL;
	if (!PyList_Check(list)) {
		PyErr_SetString(PyExc_TypeError, "need a list");
		return NULL;
	}

	Py_ssize_t count = PyList_Size(list);
	Py_ssize_t row;
	for (row = 0; row < count; row++) {
		PyObject *rowlist = PyList_GetItem(list, row);
		if (!PyList_Check(list)) {
			PyErr_SetString(PyExc_TypeError, "rows must be lists");
			fflush(out);
			return NULL;
		}
		if (PyList_Size(rowlist) != self->layout_count + 1) {
			PyErr_SetString(PyExc_TypeError, "short row list");
			fflush(out);
			return NULL;
		}

		/* Extract and write timestamp */
		write_pyobject(out, PyList_GetItem(rowlist, 0),
			       LAYOUT_TYPE_FLOAT64);
		if (PyErr_Occurred()) {
			fflush(out);
			return NULL;
		}

		/* Extract and write values */
		int i;
		for (i = 0; i < self->layout_count; i++) {
			write_pyobject(out, PyList_GetItem(rowlist, i+1),
				self->layout_type);
			if (PyErr_Occurred()) {
				fflush(out);
				return NULL;
			}
		}
	}
	fflush(out);
	/* All done */
	Py_INCREF(Py_None);
	return Py_None;
}

static int _extract_handle_params(PyObject *args, FILE **file, long *count)
{
	PyObject *pyfile, *pyoffset, *pycount;
	long offset;
	if (!PyArg_ParseTuple(args, "OOO",
			      &pyfile, &pyoffset, &pycount))
		return -1;
	if ((*file = PyFile_AsFile(pyfile)) == NULL)
		return -1;
	offset = PyLong_AsLong(pyoffset);
	if (PyErr_Occurred())
		return -1;
	*count = PyLong_AsLong(pycount);
	if (PyErr_Occurred())
		return -1;

	/* Seek to target location */
	if (fseek(*file, offset, SEEK_SET) < 0) {
		PyErr_SetFromErrno(PyExc_OSError);
		return -1;
	}
	return 0;
}

static PyObject *Rocket_extract_list(Rocket *self, PyObject *args)
{
	FILE *in;
	long count;
	if (_extract_handle_params(args, &in, &count) < 0)
		return NULL;

	/* Make a list to return */
	PyObject *retlist = PyList_New(0);
	if (!retlist)
		return NULL;

	/* Read data into new Python lists */
	int row;
	for (row = 0; row < count; row++)
	{
		PyObject *rowlist = PyList_New(self->layout_count + 1);
		if (!rowlist) {
			Py_DECREF(retlist);
			return NULL;
		}

		/* Timestamp */
		PyObject *entry = read_pyobject(in, LAYOUT_TYPE_FLOAT64);
		if (!entry || (PyList_SetItem(rowlist, 0, entry) < 0)) {
			Py_DECREF(rowlist);
			Py_DECREF(retlist);
			return NULL;
		}

		/* Data */
		int i;
		for (i = 0; i < self->layout_count; i++) {
			PyObject *ent = read_pyobject(in, self->layout_type);
			if (!ent || (PyList_SetItem(rowlist, i+1, ent) < 0)) {
				Py_DECREF(rowlist);
				Py_DECREF(retlist);
				return NULL;
			}
		}

		/* Add row to return value */
		if (PyList_Append(retlist, rowlist) < 0) {
			Py_DECREF(rowlist);
			Py_DECREF(retlist);
			return NULL;
		}

		Py_DECREF(rowlist);
	}
	return retlist;
}

static PyObject *Rocket_extract_string(Rocket *self, PyObject *args)
{
	FILE *in;
	long count;
	if (_extract_handle_params(args, &in, &count) < 0)
		return NULL;

	char *str = NULL, *new;
	long len_alloc = 0;
	long len = 0;
	int ret;

	/* min space free in string (and the maximum length of one
	   line); this is generous */
	const int min_free = 32 * MAX_LAYOUT_COUNT;

	/* how much to allocate at once */
	const int alloc_size = 1048576;

	int row, i;
	union_t u;
	for (row = 0; row < count; row++) {
		/* Make sure there's space for a line */
		if ((len_alloc - len) < min_free) {
			/* grow by 1 meg at a time */
			len_alloc += alloc_size;
			new = realloc(str, len_alloc);
			if (new == NULL)
				goto err;
			str = new;
		}

		/* Read and print timestamp */
		if (fread(&u.u64[0], 8, 1, in) != 1)
			goto err;
		/* Timestamps are always printed to the microsecond */
		ret = sprintf(&str[len], "%.6f", u.d[0]);
		if (ret <= 0)
			goto err;
		len += ret;

		/* Read and print values */
		switch (self->layout_type) {
#define CASE(type, fmt, fmttype, disktype, letoh, bytes)		\
		case LAYOUT_TYPE_##type:				\
			/* read and format in a loop */			\
			for (i = 0; i < self->layout_count; i++) {	\
				if (fread(&u.disktype[0], bytes,	\
					  1, in) < 0)			\
					goto err;			\
				u.disktype[0] = letoh(u.disktype[0]);	\
				ret = sprintf(&str[len], " " fmt,	\
					      u.fmttype[0]);		\
				if (ret <= 0)				\
					goto err;			\
				len += ret;				\
			}						\
			break
			CASE(INT8,   "%hhd",   i8,  u8,         , 1);
			CASE(UINT8,  "%hhu",   u8,  u8,         , 1);
			CASE(INT16,  "%hd",    i16, u16, le16toh, 2);
			CASE(UINT16, "%hu",    u16, u16, le16toh, 2);
			CASE(INT32,  "%d",     i32, u32, le32toh, 4);
			CASE(UINT32, "%u",     u32, u32, le32toh, 4);
			CASE(INT64,  "%ld",    i64, u64, le64toh, 8);
			CASE(UINT64, "%lu",    u64, u64, le64toh, 8);
			/* These next two are a bit debatable.  floats
			   are 6-9 significant figures, doubles are
			   15-19.  This matches the old prep format,
			   for float32. */
			CASE(FLOAT32, "%.6e",  f,   u32, le32toh, 4);
			CASE(FLOAT64, "%.16e", d,   u64, le64toh, 8);
#undef CASE
		default:
			PyErr_SetString(PyExc_TypeError, "unknown type");
			if (str) free(str);
			return NULL;
		}
		str[len++] = '\n';
	}

	PyObject *pystr = PyString_FromStringAndSize(str, len);
	free(str);
	return pystr;
err:
	if (str) free(str);
	PyErr_SetFromErrno(PyExc_OSError);
	return NULL;
}

static PyMemberDef Rocket_members[] = {
	{ "binary_size", T_INT, offsetof(Rocket, binary_size), 0,
	  "binary size per row" },
	{ NULL },
};

static PyMethodDef Rocket_methods[] = {
	{ "append_list", (PyCFunction)Rocket_append_list, METH_VARARGS,
	  "Append the list data to the file" },
	{ "extract_list", (PyCFunction)Rocket_extract_list, METH_VARARGS,
	  "Extract count rows of data from the file at offset offset. "
	  "Return a list of lists [[row],[row],...]" },
	{ "extract_string", (PyCFunction)Rocket_extract_string, METH_VARARGS,
	  "Extract count rows of data from the file at offset offset. "
	  "Return an ascii formatted string according to the layout" },
	{ NULL },
};

static PyTypeObject RocketType = {
	PyObject_HEAD_INIT(NULL)

	.tp_name	= "rocket.Rocket",
	.tp_basicsize	= sizeof(Rocket),
	.tp_flags	= Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,

	.tp_new		= Rocket_new,
	.tp_dealloc	= (destructor)Rocket_dealloc,
	.tp_init	= (initproc)Rocket_init,
	.tp_methods	= Rocket_methods,
	.tp_members	= Rocket_members,

	.tp_doc		= ("C implementation of the \"rocket\" data parsing "
			   "interface, which translates between the binary "
			   "format on disk and the ASCII or Python list "
			   "format used when communicating with the rest of "
			   "the system.")
};

static PyMethodDef module_methods[] = {
	{ NULL },
};

PyMODINIT_FUNC
initrocket(void)
{
	PyObject *module;

	RocketType.tp_new = PyType_GenericNew;
	if (PyType_Ready(&RocketType) < 0)
		return;

	module = Py_InitModule3("rocket", module_methods,
				"Rocket data parsing and formatting module");
	Py_INCREF(&RocketType);
	PyModule_AddObject(module, "Rocket", (PyObject *)&RocketType);
	return;
}