#include "Python.h"

#include "structmember.h"

PyDoc_STRVAR(pickle_module_doc,

"Optimized C implementation for the Python pickle module.");

enum {

HIGHEST_PROTOCOL = 3,

DEFAULT_PROTOCOL = 3

};

enum opcode {

};

#undef TRUE

#define TRUE "I01\n"

#undef FALSE

#define FALSE "I00\n"

enum {

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

BATCHSIZE = 1000,

/* Nesting limit until Pickler, when running in "fast mode", starts

FAST_NESTING_LIMIT = 50,

/* Initial size of the write buffer of Pickler. */

WRITE_BUF_SIZE = 4096,

/* Maximum size of the write buffer of Pickler when pickling to a

MAX_WRITE_BUF_SIZE = 64 * 1024,

/* Prefetch size when unpickling (disabled on unpeekable streams) */

PREFETCH = 8192 * 16

};

static PyObject *PickleError = NULL;

static PyObject *PicklingError = NULL;

static PyObject *UnpicklingError = NULL;

static PyObject *dispatch_table = NULL;

static PyObject *extension_registry = NULL;

static PyObject *inverted_registry = NULL;

static PyObject *extension_cache = NULL;

static PyObject *name_mapping_2to3 = NULL;

static PyObject *import_mapping_2to3 = NULL;

static PyObject *name_mapping_3to2 = NULL;

static PyObject *import_mapping_3to2 = NULL;

static PyObject *empty_tuple = NULL;

static PyObject *two_tuple = NULL;

static int

stack_underflow(void)

{

PyErr_SetString(UnpicklingError, "unpickling stack underflow");

return -1;

}

typedef struct {

PyObject_VAR_HEAD

PyObject **data;

Py_ssize_t allocated; /* number of slots in data allocated */

} Pdata;

static void

Pdata_dealloc(Pdata *self)

{

Py_ssize_t i = Py_SIZE(self);

while (--i >= 0) {

Py_DECREF(self->data[i]);

}

PyMem_FREE(self->data);

PyObject_Del(self);

}

static PyTypeObject Pdata_Type = {

PyVarObject_HEAD_INIT(NULL, 0)

"_pickle.Pdata", /*tp_name*/

sizeof(Pdata), /*tp_basicsize*/

0, /*tp_itemsize*/

(destructor)Pdata_dealloc, /*tp_dealloc*/

};

static PyObject *

Pdata_New(void)

{

Pdata *self;

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

return NULL;

Py_SIZE(self) = 0;

self->allocated = 8;

self->data = PyMem_MALLOC(self->allocated * sizeof(PyObject *));

if (self->data)

return (PyObject *)self;

Py_DECREF(self);

return PyErr_NoMemory();

}

static int

Pdata_clear(Pdata *self, Py_ssize_t clearto)

{

Py_ssize_t i = Py_SIZE(self);

if (clearto < 0)

return stack_underflow();

if (clearto >= i)

return 0;

while (--i >= clearto) {

Py_CLEAR(self->data[i]);

}

Py_SIZE(self) = clearto;

return 0;

}

static int

Pdata_grow(Pdata *self)

{

PyObject **data = self->data;

Py_ssize_t allocated = self->allocated;

Py_ssize_t new_allocated;

new_allocated = (allocated >> 3) + 6;

/* check for integer overflow */

if (new_allocated > PY_SSIZE_T_MAX - allocated)

goto nomemory;

new_allocated += allocated;

if (new_allocated > (PY_SSIZE_T_MAX / sizeof(PyObject *)))

goto nomemory;

data = PyMem_REALLOC(data, new_allocated * sizeof(PyObject *));

if (data == NULL)

goto nomemory;

self->data = data;

self->allocated = new_allocated;

return 0;

nomemory:

PyErr_NoMemory();

return -1;

}

static PyObject *

Pdata_pop(Pdata *self)

{

if (Py_SIZE(self) == 0) {

PyErr_SetString(UnpicklingError, "bad pickle data");

return NULL;

}

return self->data[--Py_SIZE(self)];

}

#define PDATA_POP(D, V) do { (V) = Pdata_pop((D)); } while (0)

static int

Pdata_push(Pdata *self, PyObject *obj)

{

if (Py_SIZE(self) == self->allocated && Pdata_grow(self) < 0) {

return -1;

}

self->data[Py_SIZE(self)++] = obj;

return 0;

}

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

if (Pdata_push((D), (O)) < 0) return (ER); } while(0)

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

Py_INCREF((O)); \

if (Pdata_push((D), (O)) < 0) return (ER); } while(0)

static PyObject *

Pdata_poptuple(Pdata *self, Py_ssize_t start)

{

PyObject *tuple;

Py_ssize_t len, i, j;

len = Py_SIZE(self) - start;

tuple = PyTuple_New(len);

if (tuple == NULL)

return NULL;

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

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

Py_SIZE(self) = start;

return tuple;

}

static PyObject *

Pdata_poplist(Pdata *self, Py_ssize_t start)

{

PyObject *list;

Py_ssize_t len, i, j;

len = Py_SIZE(self) - start;

list = PyList_New(len);

if (list == NULL)

return NULL;

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

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

Py_SIZE(self) = start;

return list;

}

typedef struct {

PyObject *me_key;

Py_ssize_t me_value;

} PyMemoEntry;

typedef struct {

Py_ssize_t mt_mask;

Py_ssize_t mt_used;

Py_ssize_t mt_allocated;

PyMemoEntry *mt_table;

} PyMemoTable;

typedef struct PicklerObject {

} PicklerObject;

typedef struct UnpicklerObject {

} UnpicklerObject;

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

static int save_reduce(PicklerObject *, PyObject *, PyObject *);

static PyTypeObject Pickler_Type;

static PyTypeObject Unpickler_Type;

#define MT_MINSIZE 8

#define PERTURB_SHIFT 5

static PyMemoTable *

PyMemoTable_New(void)

{

PyMemoTable *memo = PyMem_MALLOC(sizeof(PyMemoTable));

if (memo == NULL) {

PyErr_NoMemory();

return NULL;

}

memo->mt_used = 0;

memo->mt_allocated = MT_MINSIZE;

memo->mt_mask = MT_MINSIZE - 1;

memo->mt_table = PyMem_MALLOC(MT_MINSIZE * sizeof(PyMemoEntry));

if (memo->mt_table == NULL) {

PyMem_FREE(memo);

PyErr_NoMemory();

return NULL;

}

memset(memo->mt_table, 0, MT_MINSIZE * sizeof(PyMemoEntry));

return memo;

}

static PyMemoTable *

PyMemoTable_Copy(PyMemoTable *self)

{

Py_ssize_t i;

PyMemoTable *new = PyMemoTable_New();

if (new == NULL)

return NULL;

new->mt_used = self->mt_used;

new->mt_allocated = self->mt_allocated;

new->mt_mask = self->mt_mask;

/* The table we get from _New() is probably smaller than we wanted.

PyMem_FREE(new->mt_table);

new->mt_table = PyMem_MALLOC(self->mt_allocated * sizeof(PyMemoEntry));

if (new->mt_table == NULL) {

PyMem_FREE(new);

return NULL;

}

for (i = 0; i < self->mt_allocated; i++) {

Py_XINCREF(self->mt_table[i].me_key);

}

memcpy(new->mt_table, self->mt_table,

sizeof(PyMemoEntry) * self->mt_allocated);

return new;

}

static Py_ssize_t

PyMemoTable_Size(PyMemoTable *self)

{

return self->mt_used;

}

static int

PyMemoTable_Clear(PyMemoTable *self)

{

Py_ssize_t i = self->mt_allocated;

while (--i >= 0) {

Py_XDECREF(self->mt_table[i].me_key);

}

self->mt_used = 0;

memset(self->mt_table, 0, self->mt_allocated * sizeof(PyMemoEntry));

return 0;

}

static void

PyMemoTable_Del(PyMemoTable *self)

{

if (self == NULL)

return;

PyMemoTable_Clear(self);

PyMem_FREE(self->mt_table);

PyMem_FREE(self);

}

static PyMemoEntry *

_PyMemoTable_Lookup(PyMemoTable *self, PyObject *key)

{

size_t i;

size_t perturb;

size_t mask = (size_t)self->mt_mask;

PyMemoEntry *table = self->mt_table;

PyMemoEntry *entry;

Py_hash_t hash = (Py_hash_t)key >> 3;

i = hash & mask;

entry = &table[i];

if (entry->me_key == NULL || entry->me_key == key)

return entry;

for (perturb = hash; ; perturb >>= PERTURB_SHIFT) {

i = (i << 2) + i + perturb + 1;

entry = &table[i & mask];

if (entry->me_key == NULL || entry->me_key == key)

return entry;

}

assert(0); /* Never reached */

return NULL;

}

static int

_PyMemoTable_ResizeTable(PyMemoTable *self, Py_ssize_t min_size)

{

PyMemoEntry *oldtable = NULL;

PyMemoEntry *oldentry, *newentry;

Py_ssize_t new_size = MT_MINSIZE;

Py_ssize_t to_process;

assert(min_size > 0);

/* Find the smallest valid table size >= min_size. */

while (new_size < min_size && new_size > 0)

new_size <<= 1;

if (new_size <= 0) {

PyErr_NoMemory();

return -1;

}

/* new_size needs to be a power of two. */

assert((new_size & (new_size - 1)) == 0);

/* Allocate new table. */

oldtable = self->mt_table;

self->mt_table = PyMem_MALLOC(new_size * sizeof(PyMemoEntry));

if (self->mt_table == NULL) {

PyMem_FREE(oldtable);

PyErr_NoMemory();

return -1;

}

self->mt_allocated = new_size;

self->mt_mask = new_size - 1;

memset(self->mt_table, 0, sizeof(PyMemoEntry) * new_size);

/* Copy entries from the old table. */

to_process = self->mt_used;

for (oldentry = oldtable; to_process > 0; oldentry++) {

if (oldentry->me_key != NULL) {

to_process--;

/* newentry is a pointer to a chunk of the new

newentry = _PyMemoTable_Lookup(self, oldentry->me_key);

newentry->me_key = oldentry->me_key;

newentry->me_value = oldentry->me_value;

}

}

/* Deallocate the old table. */

PyMem_FREE(oldtable);

return 0;

}

static Py_ssize_t *

PyMemoTable_Get(PyMemoTable *self, PyObject *key)

{

PyMemoEntry *entry = _PyMemoTable_Lookup(self, key);

if (entry->me_key == NULL)

return NULL;

return &entry->me_value;

}

static int

PyMemoTable_Set(PyMemoTable *self, PyObject *key, Py_ssize_t value)

{

PyMemoEntry *entry;

assert(key != NULL);

entry = _PyMemoTable_Lookup(self, key);

if (entry->me_key != NULL) {

entry->me_value = value;

return 0;

}

Py_INCREF(key);

entry->me_key = key;

entry->me_value = value;

self->mt_used++;

/* If we added a key, we can safely resize. Otherwise just return!

if (!(self->mt_used * 3 >= (self->mt_mask + 1) * 2))

return 0;

return _PyMemoTable_ResizeTable(self,

(self->mt_used > 50000 ? 2 : 4) * self->mt_used);

}

#undef MT_MINSIZE

#undef PERTURB_SHIFT

#define ARG_TUP(self, obj) do { \

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

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

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

} \

else { \

Py_DECREF((obj)); \

} \

} while (0)

#define FREE_ARG_TUP(self) do { \

if ((self)->arg->ob_refcnt > 1) \

Py_CLEAR((self)->arg); \

} while (0)

static PyObject *

_Pickler_FastCall(PicklerObject *self, PyObject *func, PyObject *arg)

{

PyObject *result = NULL;

ARG_TUP(self, arg);

if (self->arg) {

result = PyObject_Call(func, self->arg, NULL);

FREE_ARG_TUP(self);

}

return result;

}

static int

_Pickler_ClearBuffer(PicklerObject *self)

{

Py_CLEAR(self->output_buffer);

self->output_buffer =

PyBytes_FromStringAndSize(NULL, self->max_output_len);

if (self->output_buffer == NULL)

return -1;

self->output_len = 0;

return 0;

}

static PyObject *

_Pickler_GetString(PicklerObject *self)

{

PyObject *output_buffer = self->output_buffer;

assert(self->output_buffer != NULL);

self->output_buffer = NULL;

/* Resize down to exact size */

if (_PyBytes_Resize(&output_buffer, self->output_len) < 0)

return NULL;

return output_buffer;

}

static int

_Pickler_FlushToFile(PicklerObject *self)

{

PyObject *output, *result;

assert(self->write != NULL);

output = _Pickler_GetString(self);

if (output == NULL)

return -1;

result = _Pickler_FastCall(self, self->write, output);

Py_XDECREF(result);

return (result == NULL) ? -1 : 0;

}

static Py_ssize_t

_Pickler_Write(PicklerObject *self, const char *s, Py_ssize_t n)

{

Py_ssize_t i, required;

char *buffer;

assert(s != NULL);

required = self->output_len + n;

if (required > self->max_output_len) {

if (self->write != NULL && required > MAX_WRITE_BUF_SIZE) {

/* XXX This reallocates a new buffer every time, which is a bit

if (_Pickler_FlushToFile(self) < 0)

return -1;

if (_Pickler_ClearBuffer(self) < 0)

return -1;

}

if (self->write != NULL && n > MAX_WRITE_BUF_SIZE) {

/* we already flushed above, so the buffer is empty */

PyObject *result;

/* XXX we could spare an intermediate copy and pass

PyObject *output = PyBytes_FromStringAndSize(s, n);

if (s == NULL)

return -1;

result = _Pickler_FastCall(self, self->write, output);

Py_XDECREF(result);

return (result == NULL) ? -1 : 0;

}

else {

if (self->output_len >= PY_SSIZE_T_MAX / 2 - n) {

PyErr_NoMemory();

return -1;

}

self->max_output_len = (self->output_len + n) / 2 * 3;

if (_PyBytes_Resize(&self->output_buffer, self->max_output_len) < 0)

return -1;

}

}

buffer = PyBytes_AS_STRING(self->output_buffer);

if (n < 8) {

/* This is faster than memcpy when the string is short. */

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

buffer[self->output_len + i] = s[i];

}

}

else {

memcpy(buffer + self->output_len, s, n);

}

self->output_len += n;

return n;

}

static PicklerObject *

_Pickler_New(void)

{

PicklerObject *self;

self = PyObject_GC_New(PicklerObject, &Pickler_Type);

if (self == NULL)

return NULL;

self->pers_func = NULL;

self->arg = NULL;

self->write = NULL;

self->proto = 0;

self->bin = 0;

self->fast = 0;

self->fast_nesting = 0;

self->fix_imports = 0;

self->fast_memo = NULL;

self->memo = PyMemoTable_New();

if (self->memo == NULL) {

Py_DECREF(self);

return NULL;

}

self->max_output_len = WRITE_BUF_SIZE;

self->output_len = 0;

self->output_buffer = PyBytes_FromStringAndSize(NULL,

self->max_output_len);

if (self->output_buffer == NULL) {

Py_DECREF(self);

return NULL;

}

return self;

}

static int

_Pickler_SetProtocol(PicklerObject *self, PyObject *proto_obj,

PyObject *fix_imports_obj)

{

long proto = 0;

int fix_imports;

if (proto_obj == NULL || proto_obj == Py_None)

proto = DEFAULT_PROTOCOL;

else {

proto = PyLong_AsLong(proto_obj);

if (proto == -1 && PyErr_Occurred())

return -1;

}

if (proto < 0)

proto = HIGHEST_PROTOCOL;

if (proto > HIGHEST_PROTOCOL) {

PyErr_Format(PyExc_ValueError, "pickle protocol must be <= %d",

HIGHEST_PROTOCOL);

return -1;

}

fix_imports = PyObject_IsTrue(fix_imports_obj);

if (fix_imports == -1)

return -1;

self->proto = proto;

self->bin = proto > 0;

self->fix_imports = fix_imports && proto < 3;

return 0;

}

static int

_Pickler_SetOutputStream(PicklerObject *self, PyObject *file)

{

assert(file != NULL);

self->write = PyObject_GetAttrString(file, "write");

if (self->write == NULL) {

if (PyErr_ExceptionMatches(PyExc_AttributeError))

PyErr_SetString(PyExc_TypeError,

"file must have a 'write' attribute");

return -1;

}

return 0;

}

static PyObject *

_Unpickler_FastCall(UnpicklerObject *self, PyObject *func, PyObject *arg)

{

PyObject *result = NULL;

ARG_TUP(self, arg);

if (self->arg) {

result = PyObject_Call(func, self->arg, NULL);

FREE_ARG_TUP(self);

}

return result;

}

static Py_ssize_t

_Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)

{

if (self->buffer.buf != NULL)

PyBuffer_Release(&self->buffer);

if (PyObject_GetBuffer(input, &self->buffer, PyBUF_CONTIG_RO) < 0)

return -1;

self->input_buffer = self->buffer.buf;

self->input_len = self->buffer.len;

self->next_read_idx = 0;

self->prefetched_idx = self->input_len;

return self->input_len;

}

static int

_Unpickler_SkipConsumed(UnpicklerObject *self)

{

Py_ssize_t consumed = self->next_read_idx - self->prefetched_idx;

if (consumed > 0) {

PyObject *r;

assert(self->peek); /* otherwise we did something wrong */

/* This makes an useless copy... */

r = PyObject_CallFunction(self->read, "n", consumed);

if (r == NULL)

return -1;

Py_DECREF(r);

self->prefetched_idx = self->next_read_idx;

}

return 0;

}

static const Py_ssize_t READ_WHOLE_LINE = -1;

static Py_ssize_t

_Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)

{

PyObject *data;

Py_ssize_t read_size, prefetched_size = 0;

assert(self->read != NULL);

if (_Unpickler_SkipConsumed(self) < 0)

return -1;

if (n == READ_WHOLE_LINE)

data = PyObject_Call(self->readline, empty_tuple, NULL);

else {

PyObject *len = PyLong_FromSsize_t(n);

if (len == NULL)

return -1;

data = _Unpickler_FastCall(self, self->read, len);

}

if (data == NULL)

return -1;

/* Prefetch some data without advancing the file pointer, if possible */

if (self->peek) {

PyObject *len, *prefetched;

len = PyLong_FromSsize_t(PREFETCH);

if (len == NULL) {

Py_DECREF(data);

return -1;

}

prefetched = _Unpickler_FastCall(self, self->peek, len);

if (prefetched == NULL) {

if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {

/* peek() is probably not supported by the given file object */

PyErr_Clear();

Py_CLEAR(self->peek);

}

else {

Py_DECREF(data);

return -1;

}

}

else {

assert(PyBytes_Check(prefetched));

prefetched_size = PyBytes_GET_SIZE(prefetched);

PyBytes_ConcatAndDel(&data, prefetched);

if (data == NULL)

return -1;

}

}

read_size = _Unpickler_SetStringInput(self, data) - prefetched_size;

Py_DECREF(data);

self->prefetched_idx = read_size;

return read_size;

}

static Py_ssize_t

_Unpickler_Read(UnpicklerObject *self, char **s, Py_ssize_t n)

{

Py_ssize_t num_read;

if (self->next_read_idx + n <= self->input_len) {

*s = self->input_buffer + self->next_read_idx;

self->next_read_idx += n;

return n;

}

if (!self->read) {

PyErr_Format(PyExc_EOFError, "Ran out of input");

return -1;

}

num_read = _Unpickler_ReadFromFile(self, n);

if (num_read < 0)

return -1;

if (num_read < n) {

PyErr_Format(PyExc_EOFError, "Ran out of input");

return -1;

}

*s = self->input_buffer;

self->next_read_idx = n;

return n;

}