#include "Python.h"
#include "cStringIO.h"
#include "structmember.h"

#if PY_MAJOR_VERSION >= 3
#error Python2 only
#endif

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 3
#define DEFAULT_PROTOCOL 3

/*
 * Note: The UNICODE macro controls the TCHAR meaning of the win32 API. Since
 * all headers have already been included here, we can safely redefine it.
 */
#ifdef UNICODE
#  undef UNICODE
#endif

/*
 * 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 */

/* Protocol 3. */
#define BINBYTES        'B'
#define SHORT_BINBYTES  'C'

/* 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,
  *dispatch_table_str;

static PyTypeObject BinaryType_t = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "zodbpickle.binary",        /* tp_name */
    /* the sizes are filled in from PyString_Type later */
    0,                          /* tp_basicsize */
    0,                          /* tp_itemsize */
    0,                          /* tp_dealloc */
    0,                          /* tp_print */
    0,                          /* tp_getattr */
    0,                          /* tp_setattr */
    0,                          /* tp_compare */
    0,                          /* tp_repr */
    0,                          /* tp_as_number */
    0,                          /* tp_as_sequence */
    0,                          /* tp_as_mapping */
    0,                          /* tp_hash */
    0,                          /* tp_call */
    0,                          /* tp_str */
    0,                          /* tp_getattro */
    0,                          /* tp_setattro */
    0,                          /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_CHECKTYPES |
        Py_TPFLAGS_BASETYPE | Py_TPFLAGS_STRING_SUBCLASS |
        Py_TPFLAGS_HAVE_NEWBUFFER,
    NULL,                       /* tp_doc */
    0,                          /* tp_traverse */
    0,                          /* tp_clear */
    0,                          /* tp_richcompare */
    0,                          /* tp_weaklistoffset */
    0,                          /* tp_iter */
    0,                          /* tp_iternext */
    0,                          /* tp_methods */
    0,                          /* tp_members */
    0,                          /* tp_getset */
    /* The Windows compiler compalins that this is not a constant. */
    /* So initialize later. */
    /* &PyString_Type, */       /* tp_base */
    0,                          /* tp_base */
    0,                          /* tp_dict */
    0,                          /* tp_descr_get */
    0,                          /* tp_descr_set */
    0,                          /* tp_dictoffset */
    0,                          /* tp_init */
    0,                          /* tp_alloc */
    0,                          /* tp_new */
    0,                          /* tp_free */
};

static PyObject* BinaryType = (PyObject*)&BinaryType_t;

/*************************************************************************
 Internal Data type for pickle data.                                     */

typedef struct {
    PyObject_HEAD
    Py_ssize_t length;  /* number of initial slots in data currently used */
    Py_ssize_t size;    /* number of slots in data allocated */
    PyObject **data;
} Pdata;

static void
Pdata_dealloc(Pdata *self)
{
    Py_ssize_t 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) "_pickle.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, Py_ssize_t clearto)
{
    Py_ssize_t 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)
{
    Py_ssize_t bigger;
    Py_ssize_t nbytes;

    PyObject **tmp;

    if (self->size > (PY_SSIZE_T_MAX >> 1))
        goto nomemory;
    bigger = self->size << 1;
    if (bigger > (PY_SSIZE_T_MAX / sizeof(PyObject *)))
        goto nomemory;
    nbytes = bigger * sizeof(PyObject *);
    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, Py_ssize_t start)
{
    PyObject *r;
    Py_ssize_t i, j, l;