#include "Python.h"

#include "cStringIO.h"

#include "structmember.h"

PyDoc_STRVAR(cPickle_module_documentation,

"C implementation and optimization of the Python pickle module.");

#ifndef Py_eval_input

#include <graminit.h>

#define Py_eval_input eval_input

#endif /* Py_eval_input */

#define DEL_LIST_SLICE(list, from, to) (PyList_SetSlice(list, from, to, NULL))

#define WRITE_BUF_SIZE 256

/* Bump this when new opcodes are added to the pickle protocol. */

#define HIGHEST_PROTOCOL 2

/*

 * Pickle opcodes.  These must be kept in synch with pickle.py.  Extensive

 * docs are in pickletools.py.

 */

#define MARK        '('

#define STOP        '.'

#define POP         '0'

#define POP_MARK    '1'

#define DUP         '2'

#define FLOAT       'F'

#define BINFLOAT    'G'

#define INT         'I'

#define BININT      'J'

#define BININT1     'K'

#define LONG        'L'

#define BININT2     'M'

#define NONE        'N'

#define PERSID      'P'

#define BINPERSID   'Q'

#define REDUCE      'R'

#define STRING      'S'

#define BINSTRING   'T'

#define SHORT_BINSTRING 'U'

#define UNICODE     'V'

#define BINUNICODE  'X'

#define APPEND      'a'

#define BUILD       'b'

#define GLOBAL      'c'

#define DICT        'd'

#define EMPTY_DICT  '}'

#define APPENDS     'e'

#define GET         'g'

#define BINGET      'h'

#define INST        'i'

#define LONG_BINGET 'j'

#define LIST        'l'

#define EMPTY_LIST  ']'

#define OBJ         'o'

#define PUT         'p'

#define BINPUT      'q'

#define LONG_BINPUT 'r'

#define SETITEM     's'

#define TUPLE       't'

#define EMPTY_TUPLE ')'

#define SETITEMS    'u'

/* Protocol 2. */

#define PROTO	 '\x80' /* identify pickle protocol */

#define NEWOBJ   '\x81' /* build object by applying cls.__new__ to argtuple */

#define EXT1     '\x82' /* push object from extension registry; 1-byte index */

#define EXT2     '\x83' /* ditto, but 2-byte index */

#define EXT4     '\x84' /* ditto, but 4-byte index */

#define TUPLE1   '\x85' /* build 1-tuple from stack top */

#define TUPLE2   '\x86' /* build 2-tuple from two topmost stack items */

#define TUPLE3   '\x87' /* build 3-tuple from three topmost stack items */

#define NEWTRUE  '\x88' /* push True */

#define NEWFALSE '\x89' /* push False */

#define LONG1    '\x8a' /* push long from < 256 bytes */

#define LONG4    '\x8b' /* push really big long */

/* There aren't opcodes -- they're ways to pickle bools before protocol 2,

 * so that unpicklers written before bools were introduced unpickle them

 * as ints, but unpicklers after can recognize that bools were intended.

 * Note that protocol 2 added direct ways to pickle bools.

 */

#undef TRUE

#define TRUE        "I01\n"

#undef FALSE

#define FALSE       "I00\n"

/* Keep in synch with pickle.Pickler._BATCHSIZE.  This is how many elements

 * batch_list/dict() pumps out before doing APPENDS/SETITEMS.  Nothing will

 * break if this gets out of synch with pickle.py, but it's unclear that

 * would help anything either.

 */

#define BATCHSIZE 1000

static char MARKv = MARK;

static PyObject *PickleError;

static PyObject *PicklingError;

static PyObject *UnpickleableError;

static PyObject *UnpicklingError;

static PyObject *BadPickleGet;

/* As the name says, an empty tuple. */

static PyObject *empty_tuple;

/* copy_reg.dispatch_table, {type_object: pickling_function} */

static PyObject *dispatch_table;

/* For EXT[124] opcodes. */

/* copy_reg._extension_registry, {(module_name, function_name): code} */

static PyObject *extension_registry;

/* copy_reg._inverted_registry, {code: (module_name, function_name)} */

static PyObject *inverted_registry;

/* copy_reg._extension_cache, {code: object} */

static PyObject *extension_cache;

/* For looking up name pairs in copy_reg._extension_registry. */

static PyObject *two_tuple;

static PyObject *__class___str, *__getinitargs___str, *__dict___str,

  *__getstate___str, *__setstate___str, *__name___str, *__reduce___str,

  *__reduce_ex___str,

  *write_str, *append_str,

  *read_str, *readline_str, *__main___str, 

  *copyreg_str, *dispatch_table_str;

/*************************************************************************

 Internal Data type for pickle data.                                     */

typedef struct {

	PyObject_HEAD

	int length;	/* number of initial slots in data currently used */

	int size;	/* number of slots in data allocated */

	PyObject **data;

} Pdata;

static void

Pdata_dealloc(Pdata *self)

{

	int i;

	PyObject **p;

	for (i = self->length, p = self->data; --i >= 0; p++) {

		Py_DECREF(*p);

	}

	if (self->data)

		free(self->data);

	PyObject_Del(self);

}

static PyTypeObject PdataType = {

	PyVarObject_HEAD_INIT(NULL, 0) "cPickle.Pdata", sizeof(Pdata), 0,

	(destructor)Pdata_dealloc,

	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0L,0L,0L,0L, ""

};

#define Pdata_Check(O) (Py_TYPE(O) == &PdataType)

static PyObject *

Pdata_New(void)

{

	Pdata *self;

	if (!(self = PyObject_New(Pdata, &PdataType)))

		return NULL;

	self->size = 8;

	self->length = 0;

	self->data = malloc(self->size * sizeof(PyObject*));

	if (self->data)

		return (PyObject*)self;

	Py_DECREF(self);

	return PyErr_NoMemory();

}

static int

stackUnderflow(void)

{

	PyErr_SetString(UnpicklingError, "unpickling stack underflow");

	return -1;

}

/* Retain only the initial clearto items.  If clearto >= the current

 * number of items, this is a (non-erroneous) NOP.

 */

static int

Pdata_clear(Pdata *self, int clearto)

{

	int i;

	PyObject **p;

	if (clearto < 0) return stackUnderflow();

	if (clearto >= self->length) return 0;

	for (i = self->length, p = self->data + clearto;

	     --i >= clearto;

	     p++) {

		Py_CLEAR(*p);

	}

	self->length = clearto;

	return 0;

}

static int

Pdata_grow(Pdata *self)

{

	int bigger;

	size_t nbytes;

	PyObject **tmp;

	bigger = self->size << 1;

	if (bigger <= 0)	/* was 0, or new value overflows */

		goto nomemory;

	if ((int)(size_t)bigger != bigger)

		goto nomemory;

	nbytes = (size_t)bigger * sizeof(PyObject *);

	if (nbytes / sizeof(PyObject *) != (size_t)bigger)

		goto nomemory;

	tmp = realloc(self->data, nbytes);

	if (tmp == NULL)

		goto nomemory;

	self->data = tmp;

	self->size = bigger;

	return 0;

  nomemory:

	PyErr_NoMemory();

	return -1;

}

/* D is a Pdata*.  Pop the topmost element and store it into V, which

 * must be an lvalue holding PyObject*.  On stack underflow, UnpicklingError

 * is raised and V is set to NULL.  D and V may be evaluated several times.

 */

#define PDATA_POP(D, V) {					\

	if ((D)->length)					\

		(V) = (D)->data[--((D)->length)];		\

	else {							\

		PyErr_SetString(UnpicklingError, "bad pickle data");	\

		(V) = NULL;					\

	}							\

}

/* PDATA_PUSH and PDATA_APPEND both push rvalue PyObject* O on to Pdata*

 * D.  If the Pdata stack can't be grown to hold the new value, both

 * raise MemoryError and execute "return ER".  The difference is in ownership

 * of O after:  _PUSH transfers ownership of O from the caller to the stack

 * (no incref of O is done, and in case of error O is decrefed), while

 * _APPEND pushes a new reference.

 */

/* Push O on stack D, giving ownership of O to the stack. */

#define PDATA_PUSH(D, O, ER) {					\

	if (((Pdata*)(D))->length == ((Pdata*)(D))->size &&	\

	    Pdata_grow((Pdata*)(D)) < 0) {			\

		Py_DECREF(O);					\

		return ER;					\

	}							\

	((Pdata*)(D))->data[((Pdata*)(D))->length++] = (O);	\

}

/* Push O on stack D, pushing a new reference. */

#define PDATA_APPEND(D, O, ER) {				\

	if (((Pdata*)(D))->length == ((Pdata*)(D))->size &&	\

	    Pdata_grow((Pdata*)(D)) < 0)			\

		return ER;					\

	Py_INCREF(O);						\

	((Pdata*)(D))->data[((Pdata*)(D))->length++] = (O);	\

}

static PyObject *

Pdata_popTuple(Pdata *self, int start)

{

	PyObject *r;

	int i, j, l;

	l = self->length-start;

	r = PyTuple_New(l);

	if (r == NULL)

		return NULL;

	for (i = start, j = 0 ; j < l; i++, j++)

		PyTuple_SET_ITEM(r, j, self->data[i]);

	self->length = start;

	return r;

}

static PyObject *

Pdata_popList(Pdata *self, int start)

{

	PyObject *r;

	int i, j, l;

	l=self->length-start;

	if (!( r=PyList_New(l)))  return NULL;

	for (i=start, j=0 ; j < l; i++, j++)

		PyList_SET_ITEM(r, j, self->data[i]);

	self->length=start;

	return r;

}

/*************************************************************************/

#define ARG_TUP(self, o) {                          \

  if (self->arg || (self->arg=PyTuple_New(1))) {    \

      Py_XDECREF(PyTuple_GET_ITEM(self->arg,0));    \

      PyTuple_SET_ITEM(self->arg,0,o);              \

  }                                                 \

  else {                                            \

      Py_DECREF(o);                                 \

  }                                                 \

}

#define FREE_ARG_TUP(self) {                        \

    if (Py_REFCNT(self->arg) > 1) {                 \

      Py_DECREF(self->arg);                         \

      self->arg=NULL;                               \

    }                                               \

  }

typedef struct Picklerobject {

	PyObject_HEAD

	FILE *fp;

	PyObject *write;

	PyObject *file;

	PyObject *memo;

	PyObject *arg;

	PyObject *pers_func;

	PyObject *inst_pers_func;

	/* pickle protocol number, >= 0 */

	int proto;

	/* bool, true if proto > 0 */

	int bin;

	int fast; /* Fast mode doesn't save in memo, don't use if circ ref */

	int (*write_func)(struct Picklerobject *, const char *, Py_ssize_t);

	char *write_buf;

	int buf_size;

	PyObject *dispatch_table;

	int fast_container; /* count nested container dumps */

	PyObject *fast_memo;

} Picklerobject;

#ifndef PY_CPICKLE_FAST_LIMIT

#define PY_CPICKLE_FAST_LIMIT 50

#endif

static PyTypeObject Picklertype;

typedef struct Unpicklerobject {

	PyObject_HEAD

	FILE *fp;

	PyObject *file;

	PyObject *readline;

	PyObject *read;

	PyObject *memo;

	PyObject *arg;

	Pdata *stack;

	PyObject *mark;

	PyObject *pers_func;

	PyObject *last_string;

	int *marks;

	int num_marks;

	int marks_size;

	Py_ssize_t (*read_func)(struct Unpicklerobject *, char **, Py_ssize_t);

	Py_ssize_t (*readline_func)(struct Unpicklerobject *, char **);

	int buf_size;

	char *buf;

	PyObject *find_class;

} Unpicklerobject;

static PyTypeObject Unpicklertype;

/* Forward decls that need the above structs */

static int save(Picklerobject *, PyObject *, int);

static int put2(Picklerobject *, PyObject *);

static

PyObject *

cPickle_ErrFormat(PyObject *ErrType, char *stringformat, char *format, ...)

{

	va_list va;

	PyObject *args=0, *retval=0;

	va_start(va, format);

	if (format) args = Py_VaBuildValue(format, va);

	va_end(va);

	if (format && ! args) return NULL;

	if (stringformat && !(retval=PyString_FromString(stringformat)))

		return NULL;

	if (retval) {

		if (args) {

			PyObject *v;

			v=PyString_Format(retval, args);

			Py_DECREF(retval);

			Py_DECREF(args);

			if (! v) return NULL;

			retval=v;

		}

	}

	else

		if (args) retval=args;

		else {

			PyErr_SetObject(ErrType,Py_None);

			return NULL;

		}

	PyErr_SetObject(ErrType,retval);

	Py_DECREF(retval);

	return NULL;

}

static int

write_file(Picklerobject *self, const char *s, Py_ssize_t  n)

{

	size_t nbyteswritten;

	if (s == NULL) {

		return 0;

	}

	if (n > INT_MAX) {

		/* String too large */

		return -1;

	}

	PyFile_IncUseCount((PyFileObject *)self->file);

	Py_BEGIN_ALLOW_THREADS

	nbyteswritten = fwrite(s, sizeof(char), n, self->fp);

	Py_END_ALLOW_THREADS

	PyFile_DecUseCount((PyFileObject *)self->file);

	if (nbyteswritten != (size_t)n) {

		PyErr_SetFromErrno(PyExc_IOError);

		return -1;

	}

	return (int)n;

}

static int

write_cStringIO(Picklerobject *self, const char *s, Py_ssize_t  n)

{

	if (s == NULL) {

		return 0;

	}

	if (PycStringIO->cwrite((PyObject *)self->file, s, n) != n) {

		return -1;

	}

	return (int)n;

}

static int

write_none(Picklerobject *self, const char *s, Py_ssize_t  n)

{

	if (s == NULL) return 0;

	if (n > INT_MAX) return -1;

	return (int)n;

}

static int

write_other(Picklerobject *self, const char *s, Py_ssize_t  _n)

{

	PyObject *py_str = 0, *junk = 0;

	int n;

	if (_n > INT_MAX)

		return -1;

	n = (int)_n;

	if (s == NULL) {

		if (!( self->buf_size ))  return 0;

		py_str = PyString_FromStringAndSize(self->write_buf,

						    self->buf_size);

		if (!py_str)

			return -1;

	}

	else {

		if (self->buf_size && (n + self->buf_size) > WRITE_BUF_SIZE) {

			if (write_other(self, NULL, 0) < 0)

				return -1;

		}

		if (n > WRITE_BUF_SIZE) {

			if (!( py_str =

			       PyString_FromStringAndSize(s, n)))

				return -1;

		}

		else {

			memcpy(self->write_buf + self->buf_size, s, n);

			self->buf_size += n;

			return n;

		}

	}

	if (self->write) {

		/* object with write method */

		ARG_TUP(self, py_str);

		if (self->arg) {

			junk = PyObject_Call(self->write, self->arg, NULL);

			FREE_ARG_TUP(self);

		}

		if (junk) Py_DECREF(junk);

		else return -1;

	}

	else

	    PDATA_PUSH(self->file, py_str, -1);

	self->buf_size = 0;

	return n;

}

static Py_ssize_t

read_file(Unpicklerobject *self, char **s, Py_ssize_t n)

{

	size_t nbytesread;

	if (self->buf_size == 0) {

		int size;

		size = ((n < 32) ? 32 : n);