/* pickle accelerator C extensor: _pickle module.

*

* It is built as a built-in module (Py_BUILD_CORE_BUILTIN define) on Windows

* and as an extension module (Py_BUILD_CORE_MODULE define) on other

* platforms. */

#ifndef Py_BUILD_CORE_BUILTIN

# define Py_BUILD_CORE_MODULE 1

#endif

#include "Python.h"

#include "pycore_bytesobject.h" // _PyBytesWriter

#include "pycore_ceval.h" // _Py_EnterRecursiveCall()

#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()

#include "pycore_long.h" // _PyLong_AsByteArray()

#include "pycore_moduleobject.h" // _PyModule_GetState()

#include "pycore_object.h" // _PyNone_Type

#include "pycore_pyerrors.h" // _PyErr_FormatNote

#include "pycore_pystate.h" // _PyThreadState_GET()

#include "pycore_runtime.h" // _Py_ID()

#include "pycore_setobject.h" // _PySet_NextEntry()

#include "pycore_sysmodule.h" // _PySys_GetSizeOf()

#include "pycore_unicodeobject.h" // _PyUnicode_EqualToASCIIString()

#include <stdlib.h> // strtol()

PyDoc_STRVAR(pickle_module_doc,

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

/*[clinic input]

module _pickle

class _pickle.Pickler "PicklerObject *" ""

class _pickle.PicklerMemoProxy "PicklerMemoProxyObject *" ""

class _pickle.Unpickler "UnpicklerObject *" ""

class _pickle.UnpicklerMemoProxy "UnpicklerMemoProxyObject *" ""

[clinic start generated code]*/

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

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

Bump DEFAULT_PROTOCOL only when the oldest still supported version of Python

already includes it. */

enum {

HIGHEST_PROTOCOL = 5,

DEFAULT_PROTOCOL = 5

};

#ifdef MS_WINDOWS

// These are already typedefs from windows.h, pulled in via pycore_runtime.h.

#define FLOAT FLOAT_

#define INT INT_

#define LONG LONG_

/* This can already be defined on Windows to set the character set

the Windows header files treat as default */

#ifdef UNICODE

#undef UNICODE

#endif

#endif

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

/* Protocol 5 */

BYTEARRAY8 = '\x96',

NEXT_BUFFER = '\x97',

READONLY_BUFFER = '\x98'

};

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_MIN = 4,

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;

/* Types */

PyTypeObject *Pickler_Type;

PyTypeObject *Unpickler_Type;

PyTypeObject *Pdata_Type;

PyTypeObject *PicklerMemoProxyType;

PyTypeObject *UnpicklerMemoProxyType;

} PickleState;

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

static struct PyModuleDef _picklemodule;

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

static inline PickleState *

_Pickle_GetState(PyObject *module)

{

void *state = _PyModule_GetState(module);

assert(state != NULL);

return (PickleState *)state;

}

static inline PickleState *

_Pickle_GetStateByClass(PyTypeObject *cls)

{

void *state = _PyType_GetModuleState(cls);

assert(state != NULL);

return (PickleState *)state;

}

static inline PickleState *

_Pickle_FindStateByType(PyTypeObject *tp)

{

PyObject *module = PyType_GetModuleByDef(tp, &_picklemodule);

assert(module != NULL);

return _Pickle_GetState(module);

}

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

Py_CLEAR(st->Pickler_Type);

Py_CLEAR(st->Unpickler_Type);

Py_CLEAR(st->Pdata_Type);

Py_CLEAR(st->PicklerMemoProxyType);

Py_CLEAR(st->UnpicklerMemoProxyType);

}

/* Initialize the given pickle module state. */

static int

_Pickle_InitState(PickleState *st)

{

PyObject *copyreg = NULL;

PyObject *compat_pickle = NULL;

st->getattr = _PyEval_GetBuiltin(&_Py_ID(getattr));

if (st->getattr == NULL)

goto error;

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

st->codecs_encode = PyImport_ImportModuleAttrString("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;

}

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

if (!st->partial)

goto error;

return 0;

error:

Py_CLEAR(copyreg);

Py_CLEAR(compat_pickle);

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

result = PyObject_CallOneArg(func, obj);

Py_DECREF(obj);

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;

#define Pdata_CAST(op) ((Pdata *)(op))

static int

Pdata_traverse(PyObject *self, visitproc visit, void *arg)

{

Py_VISIT(Py_TYPE(self));

return 0;

}

static void

Pdata_dealloc(PyObject *op)

{

Pdata *self = Pdata_CAST(op);

PyTypeObject *tp = Py_TYPE(self);

PyObject_GC_UnTrack(self);

Py_ssize_t i = Py_SIZE(self);

while (--i >= 0) {

Py_DECREF(self->data[i]);

}

PyMem_Free(self->data);

tp->tp_free(self);

Py_DECREF(tp);

}

static PyType_Slot pdata_slots[] = {

{Py_tp_dealloc, Pdata_dealloc},

{Py_tp_traverse, Pdata_traverse},

{0, NULL},

};

static PyType_Spec pdata_spec = {

.name = "_pickle.Pdata",

.basicsize = sizeof(Pdata),

.itemsize = sizeof(PyObject *),

.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |

Py_TPFLAGS_IMMUTABLETYPE),

.slots = pdata_slots,

};

static PyObject *

Pdata_New(PickleState *state)

{

Pdata *self;

if (!(self = PyObject_GC_New(Pdata, state->Pdata_Type)))

return NULL;

Py_SET_SIZE(self, 0);

self->mark_set = 0;

self->fence = 0;

self->allocated = 8;

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

if (self->data) {

PyObject_GC_Track(self);

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_SET_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(PickleState *st, Pdata *self)

{

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(PickleState *state, Pdata *self)

{

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

Pdata_stack_underflow(state, self);

return NULL;

}

Py_SET_SIZE(self, Py_SIZE(self) - 1);

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

}

#define PDATA_POP(S, D, V) do { (V) = Pdata_pop(S, (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;

Py_SET_SIZE(self, Py_SIZE(self) + 1);

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(PickleState *state, Pdata *self, Py_ssize_t start)

{

PyObject *tuple;

Py_ssize_t len, i, j;

if (start < self->fence) {

Pdata_stack_underflow(state, 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_SET_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_SET_SIZE(self, start);

return list;

}

typedef struct {

PyObject *me_key;

Py_ssize_t me_value;

} PyMemoEntry;

typedef struct {

size_t mt_mask;

size_t mt_used;

size_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 *persistent_id; /* persistent_id() method, can be NULL */

PyObject *persistent_id_attr; /* instance attribute, can be NULL */

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

PyObject *reducer_override; /* hook for invoking user-defined callbacks

instead of save_global when pickling

functions and classes*/

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;

PyObject *buffer_callback; /* Callback for out-of-band buffers, or NULL */

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

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

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

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

PyObject *persistent_load_attr; /* instance attribute, 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 *readinto; /* readinto() method of the input stream. */

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

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

PyObject *buffers; /* iterable of out-of-band buffers, 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;

#define PicklerObject_CAST(op) ((PicklerObject *)(op))

#define UnpicklerObject_CAST(op) ((UnpicklerObject *)(op))

#define PicklerMemoProxyObject_CAST(op) ((PicklerMemoProxyObject *)(op))

#define UnpicklerMemoProxyObject_CAST(op) ((UnpicklerMemoProxyObject *)(op))

/* Forward declarations */

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

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

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

{

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 (size_t 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 = 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;

}

Py_UNREACHABLE();

}

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

static int

_PyMemoTable_ResizeTable(PyMemoTable *self, size_t min_size)

{

PyMemoEntry *oldtable = NULL;

PyMemoEntry *oldentry, *newentry;

size_t new_size = MT_MINSIZE;

size_t to_process;

assert(min_size > 0);

if (min_size > PY_SSIZE_T_MAX) {

PyErr_NoMemory();

return -1;

}

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

while (new_size < min_size) {

new_size <<= 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;

}

entry->me_key = Py_NewRef(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 (SIZE_MAX / 3 >= self->mt_used && self->mt_used * 3 < self->mt_allocated * 2) {

return 0;

}

// self->mt_used is always < PY_SSIZE_T_MAX, so this can't overflow.

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

return _PyMemoTable_ResizeTable(self, desired_size);

}

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

static_assert(sizeof(size_t) <= 8, "size_t is larger than 64-bit");

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

if (frame_len >= FRAME_SIZE_MIN) {

qdata[0] = FRAME;

_write_size64(qdata + 1, frame_len);

}

else {

memmove(qdata, qdata + FRAME_HEADER_SIZE, frame_len);

self->output_len -= FRAME_HEADER_SIZE;

}

self->frame_start = -1;

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;

}