/* Write Python objects to files and read them back.

   This is intended for writing and reading compiled Python code only;

   a true persistent storage facility would be much harder, since

   it would have to take circular links and sharing into account. */

#define PY_SSIZE_T_CLEAN

#include "Python.h"

#include "longintrepr.h"

#include "code.h"

#include "marshal.h"

/* High water mark to determine when the marshalled object is dangerously deep

 * and risks coring the interpreter.  When the object stack gets this deep,

 * raise an exception instead of continuing.

 */

#define MAX_MARSHAL_STACK_DEPTH 2000

#define TYPE_NULL		'0'

#define TYPE_NONE		'N'

#define TYPE_FALSE		'F'

#define TYPE_TRUE		'T'

#define TYPE_STOPITER		'S'

#define TYPE_ELLIPSIS   	'.'

#define TYPE_INT		'i'

#define TYPE_INT64		'I'

#define TYPE_FLOAT		'f'

#define TYPE_BINARY_FLOAT	'g'

#define TYPE_COMPLEX		'x'

#define TYPE_BINARY_COMPLEX	'y'

#define TYPE_LONG		'l'

#define TYPE_STRING		's'

#define TYPE_INTERNED		't'

#define TYPE_STRINGREF		'R'

#define TYPE_TUPLE		'('

#define TYPE_LIST		'['

#define TYPE_DICT		'{'

#define TYPE_CODE		'c'

#define TYPE_UNICODE		'u'

#define TYPE_UNKNOWN		'?'

#define TYPE_SET		'<'

#define TYPE_FROZENSET  	'>'

typedef struct {

	FILE *fp;

	int error;

	int depth;

	/* If fp == NULL, the following are valid: */

	PyObject *str;

	char *ptr;

	char *end;

	PyObject *strings; /* dict on marshal, list on unmarshal */

	int version;

} WFILE;

#define w_byte(c, p) if (((p)->fp)) putc((c), (p)->fp); \

		      else if ((p)->ptr != (p)->end) *(p)->ptr++ = (c); \

			   else w_more(c, p)

static void

w_more(int c, WFILE *p)

{

	Py_ssize_t size, newsize;

	if (p->str == NULL)

		return; /* An error already occurred */

	size = PyString_Size(p->str);

	newsize = size + size + 1024;

	if (newsize > 32*1024*1024) {

		newsize = size + (size >> 3);	/* 12.5% overallocation */

	}

	if (_PyString_Resize(&p->str, newsize) != 0) {

		p->ptr = p->end = NULL;

	}

	else {

		p->ptr = PyString_AS_STRING((PyStringObject *)p->str) + size;

		p->end =

			PyString_AS_STRING((PyStringObject *)p->str) + newsize;

		*p->ptr++ = Py_SAFE_DOWNCAST(c, int, char);

	}

}

static void

w_string(char *s, int n, WFILE *p)

{

	if (p->fp != NULL) {

		fwrite(s, 1, n, p->fp);

	}

	else {

		while (--n >= 0) {

			w_byte(*s, p);

			s++;

		}

	}

}

static void

w_short(int x, WFILE *p)

{

	w_byte((char)( x      & 0xff), p);

	w_byte((char)((x>> 8) & 0xff), p);

}

static void

w_long(long x, WFILE *p)

{

	w_byte((char)( x      & 0xff), p);

	w_byte((char)((x>> 8) & 0xff), p);

	w_byte((char)((x>>16) & 0xff), p);

	w_byte((char)((x>>24) & 0xff), p);

}

#if SIZEOF_LONG > 4

static void

w_long64(long x, WFILE *p)

{

	w_long(x, p);

	w_long(x>>32, p);

}

#endif

static void

w_object(PyObject *v, WFILE *p)

{

	Py_ssize_t i, n;

	p->depth++;

	if (p->depth > MAX_MARSHAL_STACK_DEPTH) {

		p->error = 2;

	}

	else if (v == NULL) {

		w_byte(TYPE_NULL, p);

	}

	else if (v == Py_None) {

		w_byte(TYPE_NONE, p);

	}

	else if (v == PyExc_StopIteration) {

		w_byte(TYPE_STOPITER, p);

	}

	else if (v == Py_Ellipsis) {

	        w_byte(TYPE_ELLIPSIS, p);

	}

	else if (v == Py_False) {

	        w_byte(TYPE_FALSE, p);

	}

	else if (v == Py_True) {

	        w_byte(TYPE_TRUE, p);

	}

	else if (PyInt_CheckExact(v)) {

		long x = PyInt_AS_LONG((PyIntObject *)v);

#if SIZEOF_LONG > 4

		long y = Py_ARITHMETIC_RIGHT_SHIFT(long, x, 31);

		if (y && y != -1) {

			w_byte(TYPE_INT64, p);

			w_long64(x, p);

		}

		else

#endif

			{

			w_byte(TYPE_INT, p);

			w_long(x, p);

		}

	}

	else if (PyLong_CheckExact(v)) {

		PyLongObject *ob = (PyLongObject *)v;

		w_byte(TYPE_LONG, p);

		n = ob->ob_size;

		w_long((long)n, p);

		if (n < 0)

			n = -n;

		for (i = 0; i < n; i++)

			w_short(ob->ob_digit[i], p);

	}

	else if (PyFloat_CheckExact(v)) {

		if (p->version > 1) {

			unsigned char buf[8];

			if (_PyFloat_Pack8(PyFloat_AsDouble(v), 

					   buf, 1) < 0) {

				p->error = 1;

				return;

			}

			w_byte(TYPE_BINARY_FLOAT, p);

			w_string((char*)buf, 8, p);

		}

		else {

			char buf[256]; /* Plenty to format any double */

			PyFloat_AsReprString(buf, (PyFloatObject *)v);

			n = strlen(buf);

			w_byte(TYPE_FLOAT, p);

			w_byte((int)n, p);

			w_string(buf, (int)n, p);

		}

	}

#ifndef WITHOUT_COMPLEX

	else if (PyComplex_CheckExact(v)) {

		if (p->version > 1) {

			unsigned char buf[8];

			if (_PyFloat_Pack8(PyComplex_RealAsDouble(v),

					   buf, 1) < 0) {

				p->error = 1;

				return;

			}

			w_byte(TYPE_BINARY_COMPLEX, p);

			w_string((char*)buf, 8, p);

			if (_PyFloat_Pack8(PyComplex_ImagAsDouble(v), 

					   buf, 1) < 0) {

				p->error = 1;

				return;

			}

			w_string((char*)buf, 8, p);

		}

		else {

			char buf[256]; /* Plenty to format any double */

			PyFloatObject *temp;

			w_byte(TYPE_COMPLEX, p);

			temp = (PyFloatObject*)PyFloat_FromDouble(

				PyComplex_RealAsDouble(v));

			if (!temp) {

				p->error = 1;

				return;

			}

			PyFloat_AsReprString(buf, temp);

			Py_DECREF(temp);

			n = strlen(buf);

			w_byte((int)n, p);

			w_string(buf, (int)n, p);

			temp = (PyFloatObject*)PyFloat_FromDouble(

				PyComplex_ImagAsDouble(v));

			if (!temp) {

				p->error = 1;

				return;

			}

			PyFloat_AsReprString(buf, temp);

			Py_DECREF(temp);

			n = strlen(buf);

			w_byte((int)n, p);

			w_string(buf, (int)n, p);

		}

	}

#endif

	else if (PyString_CheckExact(v)) {

		if (p->strings && PyString_CHECK_INTERNED(v)) {

			PyObject *o = PyDict_GetItem(p->strings, v);

			if (o) {

				long w = PyInt_AsLong(o);

				w_byte(TYPE_STRINGREF, p);

				w_long(w, p);

				goto exit;

			}

			else {

				int ok;

				o = PyInt_FromSsize_t(PyDict_Size(p->strings));

				ok = o &&

				     PyDict_SetItem(p->strings, v, o) >= 0;

				Py_XDECREF(o);

				if (!ok) {

					p->depth--;

					p->error = 1;

					return;

				}

				w_byte(TYPE_INTERNED, p);

			}

		}

		else {

			w_byte(TYPE_STRING, p);

		}

		n = PyString_GET_SIZE(v);

		if (n > INT_MAX) {

			/* huge strings are not supported */

			p->depth--;

			p->error = 1;

			return;

		}

		w_long((long)n, p);

		w_string(PyString_AS_STRING(v), (int)n, p);

	}

#ifdef Py_USING_UNICODE

	else if (PyUnicode_CheckExact(v)) {

	        PyObject *utf8;

		utf8 = PyUnicode_AsUTF8String(v);

		if (utf8 == NULL) {

			p->depth--;

			p->error = 1;

			return;

		}

		w_byte(TYPE_UNICODE, p);

		n = PyString_GET_SIZE(utf8);

		if (n > INT_MAX) {

			p->depth--;

			p->error = 1;

			return;

		}

		w_long((long)n, p);

		w_string(PyString_AS_STRING(utf8), (int)n, p);

		Py_DECREF(utf8);

	}

#endif

	else if (PyTuple_CheckExact(v)) {

		w_byte(TYPE_TUPLE, p);

		n = PyTuple_Size(v);

		w_long((long)n, p);

		for (i = 0; i < n; i++) {

			w_object(PyTuple_GET_ITEM(v, i), p);

		}

	}

	else if (PyList_CheckExact(v)) {

		w_byte(TYPE_LIST, p);

		n = PyList_GET_SIZE(v);

		w_long((long)n, p);

		for (i = 0; i < n; i++) {

			w_object(PyList_GET_ITEM(v, i), p);

		}

	}

	else if (PyDict_CheckExact(v)) {

		Py_ssize_t pos;

		PyObject *key, *value;

		w_byte(TYPE_DICT, p);

		/* This one is NULL object terminated! */

		pos = 0;

		while (PyDict_Next(v, &pos, &key, &value)) {

			w_object(key, p);

			w_object(value, p);

		}

		w_object((PyObject *)NULL, p);

	}

	else if (PyAnySet_CheckExact(v)) {

		PyObject *value, *it;

		if (PyObject_TypeCheck(v, &PySet_Type))

			w_byte(TYPE_SET, p);

		else

			w_byte(TYPE_FROZENSET, p);

		n = PyObject_Size(v);

		if (n == -1) {

			p->depth--;

			p->error = 1;

			return;

		}

		w_long((long)n, p);

		it = PyObject_GetIter(v);

		if (it == NULL) {

			p->depth--;

			p->error = 1;

			return;

		}

		while ((value = PyIter_Next(it)) != NULL) {

			w_object(value, p);

			Py_DECREF(value);

		}

		Py_DECREF(it);

		if (PyErr_Occurred()) {

			p->depth--;

			p->error = 1;

			return;

		}

	}

	else if (PyCode_Check(v)) {

		PyCodeObject *co = (PyCodeObject *)v;

		w_byte(TYPE_CODE, p);

		w_long(co->co_argcount, p);

		w_long(co->co_nlocals, p);

		w_long(co->co_stacksize, p);

		w_long(co->co_flags, p);

		w_object(co->co_code, p);

		w_object(co->co_consts, p);

		w_object(co->co_names, p);

		w_object(co->co_varnames, p);

		w_object(co->co_freevars, p);

		w_object(co->co_cellvars, p);

		w_object(co->co_filename, p);

		w_object(co->co_name, p);

		w_long(co->co_firstlineno, p);

		w_object(co->co_lnotab, p);

	}

	else if (PyObject_CheckReadBuffer(v)) {

		/* Write unknown buffer-style objects as a string */

		char *s;

		PyBufferProcs *pb = v->ob_type->tp_as_buffer;

		w_byte(TYPE_STRING, p);

		n = (*pb->bf_getreadbuffer)(v, 0, (void **)&s);

		if (n > INT_MAX) {

			p->depth--;

			p->error = 1;

			return;

		}

		w_long((long)n, p);

		w_string(s, (int)n, p);

	}

	else {

		w_byte(TYPE_UNKNOWN, p);

		p->error = 1;

	}

   exit:

	p->depth--;

}

/* version currently has no effect for writing longs. */

void

PyMarshal_WriteLongToFile(long x, FILE *fp, int version)

{

	WFILE wf;

	wf.fp = fp;

	wf.error = 0;

	wf.depth = 0;

	wf.strings = NULL;

	wf.version = version;

	w_long(x, &wf);

}

void

PyMarshal_WriteObjectToFile(PyObject *x, FILE *fp, int version)

{

	WFILE wf;

	wf.fp = fp;

	wf.error = 0;

	wf.depth = 0;

	wf.strings = (version > 0) ? PyDict_New() : NULL;

	wf.version = version;

	w_object(x, &wf);

	Py_XDECREF(wf.strings);

}

typedef WFILE RFILE; /* Same struct with different invariants */

#define rs_byte(p) (((p)->ptr < (p)->end) ? (unsigned char)*(p)->ptr++ : EOF)

#define r_byte(p) ((p)->fp ? getc((p)->fp) : rs_byte(p))

static int

r_string(char *s, int n, RFILE *p)

{

	if (p->fp != NULL)

		/* The result fits into int because it must be <=n. */

		return (int)fread(s, 1, n, p->fp);

	if (p->end - p->ptr < n)

		n = (int)(p->end - p->ptr);

	memcpy(s, p->ptr, n);

	p->ptr += n;

	return n;

}

static int

r_short(RFILE *p)

{

	register short x;

	x = r_byte(p);

	x |= r_byte(p) << 8;

	/* Sign-extension, in case short greater than 16 bits */

	x |= -(x & 0x8000);

	return x;

}

static long

r_long(RFILE *p)

{

	register long x;

	register FILE *fp = p->fp;

	if (fp) {

		x = getc(fp);

		x |= (long)getc(fp) << 8;

		x |= (long)getc(fp) << 16;

		x |= (long)getc(fp) << 24;

	}

	else {

		x = rs_byte(p);

		x |= (long)rs_byte(p) << 8;

		x |= (long)rs_byte(p) << 16;

		x |= (long)rs_byte(p) << 24;

	}

#if SIZEOF_LONG > 4

	/* Sign extension for 64-bit machines */

	x |= -(x & 0x80000000L);

#endif

	return x;

}

/* r_long64 deals with the TYPE_INT64 code.  On a machine with

   sizeof(long) > 4, it returns a Python int object, else a Python long

   object.  Note that w_long64 writes out TYPE_INT if 32 bits is enough,

   so there's no inefficiency here in returning a PyLong on 32-bit boxes

   for everything written via TYPE_INT64 (i.e., if an int is written via

   TYPE_INT64, it *needs* more than 32 bits).

*/

static PyObject *

r_long64(RFILE *p)

{

	long lo4 = r_long(p);

	long hi4 = r_long(p);

#if SIZEOF_LONG > 4

	long x = (hi4 << 32) | (lo4 & 0xFFFFFFFFL);

	return PyInt_FromLong(x);

#else

	unsigned char buf[8];

	int one = 1;

	int is_little_endian = (int)*(char*)&one;

	if (is_little_endian) {

		memcpy(buf, &lo4, 4);

		memcpy(buf+4, &hi4, 4);

	}

	else {

		memcpy(buf, &hi4, 4);

		memcpy(buf+4, &lo4, 4);

	}

	return _PyLong_FromByteArray(buf, 8, is_little_endian, 1);

#endif

}

static PyObject *

r_object(RFILE *p)

{

	/* NULL is a valid return value, it does not necessarily means that

	   an exception is set. */

	PyObject *v, *v2;

	long i, n;

	int type = r_byte(p);

	PyObject *retval;

	p->depth++;

	if (p->depth > MAX_MARSHAL_STACK_DEPTH) {

		p->depth--;

		PyErr_SetString(PyExc_ValueError, "recursion limit exceeded");

		return NULL;

	}

	switch (type) {