#include "Python.h"

#include "structmember.h"

PyDoc_STRVAR(pickle_module_doc,

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

enum {

HIGHEST_PROTOCOL = 4,

DEFAULT_PROTOCOL = 3

};

enum opcode {

};

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,

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

PREFETCH = 8192 * 16,

FRAME_SIZE_TARGET = 64 * 1024,

FRAME_HEADER_SIZE = 9

};

typedef struct {

/* Exception classes for pickle. */

PyObject *PickleError;

PyObject *PicklingError;

PyObject *UnpicklingError;

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

PyObject *dispatch_table;

/* For the extension opcodes EXT1, EXT2 and EXT4. */

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

PyObject *extension_registry;

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

PyObject *extension_cache;

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

PyObject *inverted_registry;

/* Import mappings for compatibility with Python 2.x */

/* _compat_pickle.NAME_MAPPING,

PyObject *name_mapping_2to3;

/* _compat_pickle.IMPORT_MAPPING, {oldmodule: newmodule} */

PyObject *import_mapping_2to3;

/* Same, but with REVERSE_NAME_MAPPING / REVERSE_IMPORT_MAPPING */

PyObject *name_mapping_3to2;

PyObject *import_mapping_3to2;

/* codecs.encode, used for saving bytes in older protocols */

PyObject *codecs_encode;

/* builtins.getattr, used for saving nested names with protocol < 4 */

PyObject *getattr;

/* functools.partial, used for implementing __newobj_ex__ with protocols

PyObject *partial;

} PickleState;

static struct PyModuleDef _picklemodule;

static PickleState *

_Pickle_GetState(PyObject *module)

{

return (PickleState *)PyModule_GetState(module);

}

static PickleState *

_Pickle_GetGlobalState(void)

{

return _Pickle_GetState(PyState_FindModule(&_picklemodule));

}

static void

_Pickle_ClearState(PickleState *st)

{

Py_CLEAR(st->PickleError);

Py_CLEAR(st->PicklingError);

Py_CLEAR(st->UnpicklingError);

Py_CLEAR(st->dispatch_table);

Py_CLEAR(st->extension_registry);

Py_CLEAR(st->extension_cache);

Py_CLEAR(st->inverted_registry);

Py_CLEAR(st->name_mapping_2to3);

Py_CLEAR(st->import_mapping_2to3);

Py_CLEAR(st->name_mapping_3to2);

Py_CLEAR(st->import_mapping_3to2);

Py_CLEAR(st->codecs_encode);

Py_CLEAR(st->getattr);

Py_CLEAR(st->partial);

}

static int

_Pickle_InitState(PickleState *st)

{

PyObject *builtins;

PyObject *copyreg = NULL;

PyObject *compat_pickle = NULL;

PyObject *codecs = NULL;

PyObject *functools = NULL;

builtins = PyEval_GetBuiltins();

if (builtins == NULL)

goto error;

st->getattr = PyDict_GetItemString(builtins, "getattr");

if (st->getattr == NULL)

goto error;

Py_INCREF(st->getattr);

copyreg = PyImport_ImportModule("copyreg");

if (!copyreg)

goto error;

st->dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");

if (!st->dispatch_table)

goto error;

if (!PyDict_CheckExact(st->dispatch_table)) {

PyErr_Format(PyExc_RuntimeError,

"copyreg.dispatch_table should be a dict, not %.200s",

Py_TYPE(st->dispatch_table)->tp_name);

goto error;

}

st->extension_registry = \

PyObject_GetAttrString(copyreg, "_extension_registry");

if (!st->extension_registry)

goto error;

if (!PyDict_CheckExact(st->extension_registry)) {

PyErr_Format(PyExc_RuntimeError,

"copyreg._extension_registry should be a dict, "

"not %.200s", Py_TYPE(st->extension_registry)->tp_name);

goto error;

}

st->inverted_registry = \

PyObject_GetAttrString(copyreg, "_inverted_registry");

if (!st->inverted_registry)

goto error;

if (!PyDict_CheckExact(st->inverted_registry)) {

PyErr_Format(PyExc_RuntimeError,

"copyreg._inverted_registry should be a dict, "

"not %.200s", Py_TYPE(st->inverted_registry)->tp_name);

goto error;

}

st->extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");

if (!st->extension_cache)

goto error;

if (!PyDict_CheckExact(st->extension_cache)) {

PyErr_Format(PyExc_RuntimeError,

"copyreg._extension_cache should be a dict, "

"not %.200s", Py_TYPE(st->extension_cache)->tp_name);

goto error;

}

Py_CLEAR(copyreg);

/* Load the 2.x -> 3.x stdlib module mapping tables */

compat_pickle = PyImport_ImportModule("_compat_pickle");

if (!compat_pickle)

goto error;

st->name_mapping_2to3 = \

PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");

if (!st->name_mapping_2to3)

goto error;

if (!PyDict_CheckExact(st->name_mapping_2to3)) {

PyErr_Format(PyExc_RuntimeError,

"_compat_pickle.NAME_MAPPING should be a dict, not %.200s",

Py_TYPE(st->name_mapping_2to3)->tp_name);

goto error;

}

st->import_mapping_2to3 = \

PyObject_GetAttrString(compat_pickle, "IMPORT_MAPPING");

if (!st->import_mapping_2to3)

goto error;

if (!PyDict_CheckExact(st->import_mapping_2to3)) {

PyErr_Format(PyExc_RuntimeError,

"_compat_pickle.IMPORT_MAPPING should be a dict, "

"not %.200s", Py_TYPE(st->import_mapping_2to3)->tp_name);

goto error;

}

/* ... and the 3.x -> 2.x mapping tables */

st->name_mapping_3to2 = \

PyObject_GetAttrString(compat_pickle, "REVERSE_NAME_MAPPING");

if (!st->name_mapping_3to2)

goto error;

if (!PyDict_CheckExact(st->name_mapping_3to2)) {

PyErr_Format(PyExc_RuntimeError,

"_compat_pickle.REVERSE_NAME_MAPPING should be a dict, "

"not %.200s", Py_TYPE(st->name_mapping_3to2)->tp_name);

goto error;

}

st->import_mapping_3to2 = \

PyObject_GetAttrString(compat_pickle, "REVERSE_IMPORT_MAPPING");

if (!st->import_mapping_3to2)

goto error;

if (!PyDict_CheckExact(st->import_mapping_3to2)) {

PyErr_Format(PyExc_RuntimeError,

"_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "

"not %.200s", Py_TYPE(st->import_mapping_3to2)->tp_name);

goto error;

}

Py_CLEAR(compat_pickle);

codecs = PyImport_ImportModule("codecs");

if (codecs == NULL)

goto error;

st->codecs_encode = PyObject_GetAttrString(codecs, "encode");

if (st->codecs_encode == NULL) {

goto error;

}

if (!PyCallable_Check(st->codecs_encode)) {

PyErr_Format(PyExc_RuntimeError,

"codecs.encode should be a callable, not %.200s",

Py_TYPE(st->codecs_encode)->tp_name);

goto error;

}

Py_CLEAR(codecs);

functools = PyImport_ImportModule("functools");

if (!functools)

goto error;

st->partial = PyObject_GetAttrString(functools, "partial");

if (!st->partial)

goto error;

Py_CLEAR(functools);

return 0;

error:

Py_CLEAR(copyreg);

Py_CLEAR(compat_pickle);

Py_CLEAR(codecs);

Py_CLEAR(functools);

_Pickle_ClearState(st);

return -1;

}

static PyObject *

_Pickle_FastCall(PyObject *func, PyObject *obj)

{

PyObject *result;

result = PyObject_CallFunctionObjArgs(func, obj, NULL);

Py_DECREF(obj);

return result;

}

typedef struct {

PyObject_VAR_HEAD

PyObject **data;

int mark_set; /* is MARK set? */

Py_ssize_t fence; /* position of top MARK or 0 */

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

sizeof(PyObject *), /*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->mark_set = 0;

self->fence = 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);

assert(clearto >= self->fence);

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;

size_t allocated = (size_t)self->allocated;

size_t new_allocated;

new_allocated = (allocated >> 3) + 6;

/* check for integer overflow */

if (new_allocated > (size_t)PY_SSIZE_T_MAX - allocated)

goto nomemory;

new_allocated += allocated;

PyMem_RESIZE(data, PyObject *, new_allocated);

if (data == NULL)

goto nomemory;

self->data = data;

self->allocated = (Py_ssize_t)new_allocated;

return 0;

nomemory:

PyErr_NoMemory();

return -1;

}

static int

Pdata_stack_underflow(Pdata *self)

{

PickleState *st = _Pickle_GetGlobalState();

PyErr_SetString(st->UnpicklingError,

self->mark_set ?

"unexpected MARK found" :

"unpickling stack underflow");

return -1;

}

static PyObject *

Pdata_pop(Pdata *self)

{

if (Py_SIZE(self) <= self->fence) {

Pdata_stack_underflow(self);

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;

if (start < self->fence) {

Pdata_stack_underflow(self);

return NULL;

}

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;

typedef struct {

PyObject_HEAD

PicklerObject *pickler; /* Pickler whose memo table we're proxying. */

} PicklerMemoProxyObject;

typedef struct {

PyObject_HEAD

UnpicklerObject *unpickler;

} UnpicklerMemoProxyObject;

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

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

static PyTypeObject Pickler_Type;

static PyTypeObject Unpickler_Type;

#include "clinic/_pickle.c.h"

#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_NEW(PyMemoEntry, self->mt_allocated);

if (new->mt_table == NULL) {

PyMem_FREE(new);

PyErr_NoMemory();

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_NEW(PyMemoEntry, new_size);

if (self->mt_table == NULL) {

self->mt_table = 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

static int

_Pickler_ClearBuffer(PicklerObject *self)

{

Py_XSETREF(self->output_buffer,

PyBytes_FromStringAndSize(NULL, self->max_output_len));

if (self->output_buffer == NULL)

return -1;

self->output_len = 0;

self->frame_start = -1;

return 0;

}

static void

_write_size64(char *out, size_t value)

{

size_t i;

Py_BUILD_ASSERT(sizeof(size_t) <= 8);

for (i = 0; i < sizeof(size_t); i++) {

out[i] = (unsigned char)((value >> (8 * i)) & 0xff);

}

for (i = sizeof(size_t); i < 8; i++) {

out[i] = 0;

}

}

static void

_Pickler_WriteFrameHeader(PicklerObject *self, char *qdata, size_t frame_len)

{

qdata[0] = FRAME;

_write_size64(qdata + 1, frame_len);

}

static int

_Pickler_CommitFrame(PicklerObject *self)

{

size_t frame_len;

char *qdata;

if (!self->framing || self->frame_start == -1)

return 0;

frame_len = self->output_len - self->frame_start - FRAME_HEADER_SIZE;

qdata = PyBytes_AS_STRING(self->output_buffer) + self->frame_start;

_Pickler_WriteFrameHeader(self, qdata, frame_len);

self->frame_start = -1;

return 0;

}

static int

_Pickler_OpcodeBoundary(PicklerObject *self)

{

Py_ssize_t frame_len;

if (!self->framing || self->frame_start == -1)

return 0;

frame_len = self->output_len - self->frame_start - FRAME_HEADER_SIZE;

if (frame_len >= FRAME_SIZE_TARGET)

return _Pickler_CommitFrame(self);

else

return 0;

}

static PyObject *

_Pickler_GetString(PicklerObject *self)

{

PyObject *output_buffer = self->output_buffer;

assert(self->output_buffer != NULL);

if (_Pickler_CommitFrame(self))

return 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);

/* This will commit the frame first */

output = _Pickler_GetString(self);

if (output == NULL)

return -1;

result = _Pickle_FastCall(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 data_len)

{

Py_ssize_t i, n, required;

char *buffer;

int need_new_frame;

assert(s != NULL);

need_new_frame = (self->framing && self->frame_start == -1);

if (need_new_frame)

n = data_len + FRAME_HEADER_SIZE;

else

n = data_len;

required = self->output_len + n;

if (required > self->max_output_len) {

/* Make place in buffer for the pickle chunk */

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 (need_new_frame) {

/* Setup new frame */

Py_ssize_t frame_start = self->output_len;

self->frame_start = frame_start;

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

/* Write an invalid value, for debugging */

buffer[frame_start + i] = 0xFE;

}

self->output_len += FRAME_HEADER_SIZE;

}

if (data_len < 8) {

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

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

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

}

}

else {

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

}

self->output_len += data_len;

return data_len;

}

static PicklerObject *

_Pickler_New(void)