#include "Python.h"

#include "structmember.h"

PyDoc_STRVAR(pickle_module_doc,

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

/*[clinic input]

module _pickle

class _pickle.Pickler "PicklerObject *" "&Pickler_Type"

class _pickle.PicklerMemoProxy "PicklerMemoProxyObject *" "&PicklerMemoProxyType"

class _pickle.Unpickler "UnpicklerObject *" "&Unpickler_Type"

class _pickle.UnpicklerMemoProxy "UnpicklerMemoProxyObject *" "&UnpicklerMemoProxyType"

[clinic start generated code]*/

/*[clinic end generated code: output=da39a3ee5e6b4b0d input=4b3e113468a58e6c]*/

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

enum {

HIGHEST_PROTOCOL = 4,

DEFAULT_PROTOCOL = 3

};

/* Pickle opcodes. These must be kept updated with pickle.py.

Extensive docs are in pickletools.py. */

enum opcode {

MARK = '(',

STOP = '.',

POP = '0',

POP_MARK = '1',

DUP = '2',

FLOAT = 'F',

INT = 'I',

BININT = 'J',

BININT1 = 'K',

LONG = 'L',

BININT2 = 'M',

NONE = 'N',

PERSID = 'P',

BINPERSID = 'Q',

REDUCE = 'R',

STRING = 'S',

BINSTRING = 'T',

SHORT_BINSTRING = 'U',

UNICODE = 'V',

BINUNICODE = 'X',

APPEND = 'a',

BUILD = 'b',

GLOBAL = 'c',

DICT = 'd',

EMPTY_DICT = '}',

APPENDS = 'e',

GET = 'g',

BINGET = 'h',

INST = 'i',

LONG_BINGET = 'j',

LIST = 'l',

EMPTY_LIST = ']',

OBJ = 'o',

PUT = 'p',

BINPUT = 'q',

LONG_BINPUT = 'r',

SETITEM = 's',

TUPLE = 't',

EMPTY_TUPLE = ')',

SETITEMS = 'u',

BINFLOAT = 'G',

/* Protocol 2. */

PROTO = '\x80',

NEWOBJ = '\x81',

EXT1 = '\x82',

EXT2 = '\x83',

EXT4 = '\x84',

TUPLE1 = '\x85',

TUPLE2 = '\x86',

TUPLE3 = '\x87',

NEWTRUE = '\x88',

NEWFALSE = '\x89',

LONG1 = '\x8a',

LONG4 = '\x8b',

/* Protocol 3 (Python 3.x) */

BINBYTES = 'B',

SHORT_BINBYTES = 'C',

/* Protocol 4 */

SHORT_BINUNICODE = '\x8c',

BINUNICODE8 = '\x8d',

BINBYTES8 = '\x8e',

EMPTY_SET = '\x8f',

ADDITEMS = '\x90',

FROZENSET = '\x91',

NEWOBJ_EX = '\x92',

STACK_GLOBAL = '\x93',

MEMOIZE = '\x94',

FRAME = '\x95'

};

enum {

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

BATCHSIZE = 1000,

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

checking for self-referential data-structures. */

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

};

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

/* State of the pickle module, per PEP 3121. */

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,

{(oldmodule, oldname): (newmodule, newname)} */

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

2 and 3 */

PyObject *partial;

} PickleState;

/* Forward declaration of the _pickle module definition. */

static struct PyModuleDef _picklemodule;

/* Given a module object, get its per-module state. */

static PickleState *

_Pickle_GetState(PyObject *module)

{

return (PickleState *)PyModule_GetState(module);

}

/* Find the module instance imported in the currently running sub-interpreter

and get its state. */

static PickleState *

_Pickle_GetGlobalState(void)

{

return _Pickle_GetState(PyState_FindModule(&_picklemodule));

}

/* Clear the given pickle module state. */

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

}

/* Initialize the given pickle module state. */

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;

}

/* Helper for calling a function with a single argument quickly.

This function steals the reference of the given argument. */

static PyObject *

_Pickle_FastCall(PyObject *func, PyObject *obj)

{

PyObject *result;

PyObject *arg_tuple = PyTuple_New(1);

/* Note: this function used to reuse the argument tuple. This used to give

a slight performance boost with older pickle implementations where many

unbuffered reads occurred (thus needing many function calls).

However, this optimization was removed because it was too complicated

to get right. It abused the C API for tuples to mutate them which led

to subtle reference counting and concurrency bugs. Furthermore, the

introduction of protocol 4 and the prefetching optimization via peek()

significantly reduced the number of function calls we do. Thus, the

benefits became marginal at best. */

if (arg_tuple == NULL) {

Py_DECREF(obj);

return NULL;

}

PyTuple_SET_ITEM(arg_tuple, 0, obj);

result = PyObject_Call(func, arg_tuple, NULL);

Py_CLEAR(arg_tuple);

return result;

}

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

/* Internal data type used as the unpickling stack. */

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

}

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

}

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

*/

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;

}

/* Push an object on stack, transferring its ownership to the stack. */

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

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

/* Push an object on stack, adding a new reference to the object. */

#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 {

PyObject_HEAD

PyMemoTable *memo; /* Memo table, keep track of the seen

objects to support self-referential objects

pickling. */

PyObject *pers_func; /* persistent_id() method, can be NULL */

PyObject *dispatch_table; /* private dispatch_table, can be NULL */

PyObject *write; /* write() method of the output stream. */

PyObject *output_buffer; /* Write into a local bytearray buffer before

flushing to the stream. */

Py_ssize_t output_len; /* Length of output_buffer. */

Py_ssize_t max_output_len; /* Allocation size of output_buffer. */

int proto; /* Pickle protocol number, >= 0 */

int bin; /* Boolean, true if proto > 0 */

int framing; /* True when framing is enabled, proto >= 4 */

Py_ssize_t frame_start; /* Position in output_buffer where the

current frame begins. -1 if there

is no frame currently open. */

Py_ssize_t buf_size; /* Size of the current buffered pickle data */

int fast; /* Enable fast mode if set to a true value.

The fast mode disable the usage of memo,

therefore speeding the pickling process by

not generating superfluous PUT opcodes. It

should not be used if with self-referential

objects. */

int fast_nesting;

int fix_imports; /* Indicate whether Pickler should fix

the name of globals for Python 2.x. */

PyObject *fast_memo;

} PicklerObject;

typedef struct UnpicklerObject {

PyObject_HEAD

Pdata *stack; /* Pickle data stack, store unpickled objects. */

/* The unpickler memo is just an array of PyObject *s. Using a dict

is unnecessary, since the keys are contiguous ints. */

PyObject **memo;

Py_ssize_t memo_size; /* Capacity of the memo array */

Py_ssize_t memo_len; /* Number of objects in the memo */

PyObject *pers_func; /* persistent_load() method, can be NULL. */

Py_buffer buffer;

char *input_buffer;

char *input_line;

Py_ssize_t input_len;

Py_ssize_t next_read_idx;

Py_ssize_t prefetched_idx; /* index of first prefetched byte */

PyObject *read; /* read() method of the input stream. */

PyObject *readline; /* readline() method of the input stream. */

PyObject *peek; /* peek() method of the input stream, or NULL */

char *encoding; /* Name of the encoding to be used for

decoding strings pickled using Python

2.x. The default value is "ASCII" */

char *errors; /* Name of errors handling scheme to used when

decoding strings. The default value is

"strict". */

Py_ssize_t *marks; /* Mark stack, used for unpickling container

objects. */

Py_ssize_t num_marks; /* Number of marks in the mark stack. */

Py_ssize_t marks_size; /* Current allocated size of the mark stack. */

int proto; /* Protocol of the pickle loaded. */

int fix_imports; /* Indicate whether Unpickler should fix

the name of globals pickled by Python 2.x. */

} UnpicklerObject;

typedef struct {

PyObject_HEAD

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

} PicklerMemoProxyObject;

typedef struct {

PyObject_HEAD

UnpicklerObject *unpickler;

} UnpicklerMemoProxyObject;

/* Forward declarations */

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"

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

A custom hashtable mapping void* to Python ints. This is used by the pickler

for memoization. Using a custom hashtable rather than PyDict allows us to skip

a bunch of unnecessary object creation. This makes a huge performance

difference. */

#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.

Free it and allocate one that's the right size. */

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

}

/* Since entries cannot be deleted from this hashtable, _PyMemoTable_Lookup()

can be considerably simpler than dictobject.c's lookdict(). */

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;

}

/* Returns -1 on failure, 0 on success. */

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

mt_table, so we're setting the key:value pair

in-place. */

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;

}

/* Returns NULL on failure, a pointer to the value otherwise. */

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;

}

/* Returns -1 on failure, 0 on success. */

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 used >= 2/3 size, adjust size. Normally, this quaduples the size.

*

* Quadrupling the size improves average table sparseness

* (reducing collisions) at the cost of some memory. It also halves

* the number of expensive resize operations in a growing memo table.

*

* Very large memo tables (over 50K items) use doubling instead.

* This may help applications with severe memory constraints.

*/

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;

}