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

#if !defined(Py_BUILD_CORE_BUILTIN) && !defined(Py_BUILD_CORE_MODULE)

# error "Py_BUILD_CORE_BUILTIN or Py_BUILD_CORE_MODULE must be defined"

#endif

#include "Python.h"

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

#include "structmember.h" // PyMemberDef

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

};

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

} 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 *copyreg = NULL;

PyObject *compat_pickle = NULL;

PyObject *codecs = NULL;

PyObject *functools = NULL;

_Py_IDENTIFIER(getattr);

st->getattr = _PyEval_GetBuiltinId(&PyId_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);

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;

result = PyObject_CallOneArg(func, obj);

Py_DECREF(obj);

return result;

}

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

/* Retrieve and deconstruct a method for avoiding a reference cycle

(pickler -> bound method of pickler -> pickler) */

static int

init_method_ref(PyObject *self, _Py_Identifier *name,

PyObject **method_func, PyObject **method_self)

{

PyObject *func, *func2;

int ret;

/* *method_func and *method_self should be consistent. All refcount decrements

should be occurred after setting *method_self and *method_func. */

ret = _PyObject_LookupAttrId(self, name, &func);

if (func == NULL) {

*method_self = NULL;

Py_CLEAR(*method_func);

return ret;

}

if (PyMethod_Check(func) && PyMethod_GET_SELF(func) == self) {

/* Deconstruct a bound Python method */

func2 = PyMethod_GET_FUNCTION(func);

Py_INCREF(func2);

*method_self = self; /* borrowed */

Py_XSETREF(*method_func, func2);

Py_DECREF(func);

return 0;

}

else {

*method_self = NULL;

Py_XSETREF(*method_func, func);

return 0;

}

}

/* Bind a method if it was deconstructed */

static PyObject *

reconstruct_method(PyObject *func, PyObject *self)

{

if (self) {

return PyMethod_New(func, self);

}

else {

Py_INCREF(func);

return func;

}

}

static PyObject *

call_method(PyObject *func, PyObject *self, PyObject *obj)

{

if (self) {

return PyObject_CallFunctionObjArgs(func, self, obj, NULL);

}

else {

return PyObject_CallOneArg(func, obj);

}

}

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

/* 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_Free(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_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)

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

}

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

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;

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

PyObject *pers_func_self; /* borrowed reference to self if pers_func

is an unbound method, NULL otherwise */

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 *pers_func; /* persistent_load() method, can be NULL. */

PyObject *pers_func_self; /* borrowed reference to self if pers_func

is an unbound method, NULL otherwise */

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;

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

{

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;

}

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

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