#include "Python.h"

#include "structmember.h"

PyDoc_STRVAR(pickle_module_doc,

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

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

enum {

HIGHEST_PROTOCOL = 3,

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'

};

/* These aren't opcodes -- they're ways to pickle bools before protocol 2

* so that unpicklers written before bools were introduced unpickle them

* as ints, but unpicklers after can recognize that bools were intended.

* Note that protocol 2 added direct ways to pickle bools.

*/

#undef TRUE

#define TRUE "I01\n"

#undef FALSE

#define FALSE "I00\n"

enum {

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

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,

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

stream. This is ignored for in-memory pickling. */

MAX_WRITE_BUF_SIZE = 64 * 1024,

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

PREFETCH = 8192 * 16

};

/* Exception classes for pickle. These should override the ones defined in

pickle.py, when the C-optimized Pickler and Unpickler are used. */

static PyObject *PickleError = NULL;

static PyObject *PicklingError = NULL;

static PyObject *UnpicklingError = NULL;

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

static PyObject *dispatch_table = NULL;

/* For EXT[124] opcodes. */

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

static PyObject *extension_registry = NULL;

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

static PyObject *inverted_registry = NULL;

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

static PyObject *extension_cache = NULL;

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

static PyObject *name_mapping_2to3 = NULL;

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

static PyObject *import_mapping_2to3 = NULL;

/* Same, but with REVERSE_NAME_MAPPING / REVERSE_IMPORT_MAPPING */

static PyObject *name_mapping_3to2 = NULL;

static PyObject *import_mapping_3to2 = NULL;

/* XXX: Are these really nescessary? */

/* As the name says, an empty tuple. */

static PyObject *empty_tuple = NULL;

/* For looking up name pairs in copyreg._extension_registry. */

static PyObject *two_tuple = NULL;

static int

stack_underflow(void)

{

PyErr_SetString(UnpicklingError, "unpickling stack underflow");

return -1;

}

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

typedef struct {

PyObject_VAR_HEAD

PyObject **data;

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

} Pdata;

static void

Pdata_dealloc(Pdata *self)

{

Py_ssize_t i = Py_SIZE(self);

while (--i >= 0) {

Py_DECREF(self->data[i]);

}

PyMem_FREE(self->data);

PyObject_Del(self);

}

static PyTypeObject Pdata_Type = {

PyVarObject_HEAD_INIT(NULL, 0)

"_pickle.Pdata", /*tp_name*/

sizeof(Pdata), /*tp_basicsize*/

0, /*tp_itemsize*/

(destructor)Pdata_dealloc, /*tp_dealloc*/

};

static PyObject *

Pdata_New(void)

{

Pdata *self;

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

return NULL;

Py_SIZE(self) = 0;

self->allocated = 8;

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

if (self->data)

return (PyObject *)self;

Py_DECREF(self);

return PyErr_NoMemory();

}

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

if (clearto < 0)

return stack_underflow();

if (clearto >= i)

return 0;

while (--i >= clearto) {

Py_CLEAR(self->data[i]);

}

Py_SIZE(self) = clearto;

return 0;

}

static int

Pdata_grow(Pdata *self)

{

PyObject **data = self->data;

Py_ssize_t allocated = self->allocated;

Py_ssize_t new_allocated;

new_allocated = (allocated >> 3) + 6;

/* check for integer overflow */

if (new_allocated > PY_SSIZE_T_MAX - allocated)

goto nomemory;

new_allocated += allocated;

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

goto nomemory;

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

if (data == NULL)

goto nomemory;

self->data = data;

self->allocated = new_allocated;

return 0;

nomemory:

PyErr_NoMemory();

return -1;

}

/* 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) == 0) {

PyErr_SetString(UnpicklingError, "bad pickle data");

return NULL;

}

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

}

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

static int

Pdata_push(Pdata *self, PyObject *obj)

{

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

return -1;

}

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

return 0;

}

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

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

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

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;

PyObject *arg;

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;

/* Forward declarations */

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

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

static PyTypeObject Pickler_Type;

static PyTypeObject Unpickler_Type;

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

A custom hashtable mapping void* to longs. 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_MALLOC(self->mt_allocated * sizeof(PyMemoEntry));

if (new->mt_table == NULL) {

PyMem_FREE(new);

return NULL;

}

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

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

}

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

sizeof(PyMemoEntry) * self->mt_allocated);

return new;

}

static Py_ssize_t

PyMemoTable_Size(PyMemoTable *self)

{

return self->mt_used;

}

static int

PyMemoTable_Clear(PyMemoTable *self)

{

Py_ssize_t i = self->mt_allocated;

while (--i >= 0) {

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

}

self->mt_used = 0;

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

return 0;

}

static void

PyMemoTable_Del(PyMemoTable *self)

{

if (self == NULL)

return;

PyMemoTable_Clear(self);

PyMem_FREE(self->mt_table);

PyMem_FREE(self);

}

/* 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_MALLOC(new_size * sizeof(PyMemoEntry));

if (self->mt_table == NULL) {

PyMem_FREE(oldtable);

PyErr_NoMemory();

return -1;

}

self->mt_allocated = new_size;

self->mt_mask = new_size - 1;

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

/* Copy entries from the old table. */

to_process = self->mt_used;

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

if (oldentry->me_key != NULL) {

to_process--;

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

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

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

/* Helpers for creating the argument tuple passed to functions. This has the

performance advantage of calling PyTuple_New() only once.

XXX(avassalotti): Inline directly in _Pickler_FastCall() and

_Unpickler_FastCall(). */

#define ARG_TUP(self, obj) do { \

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

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

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

} \

else { \

Py_DECREF((obj)); \

} \

} while (0)

#define FREE_ARG_TUP(self) do { \

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

Py_CLEAR((self)->arg); \

} while (0)

/* A temporary cleaner API for fast single argument function call.

XXX: Does caching the argument tuple provides any real performance benefits?

A quick benchmark, on a 2.0GHz Athlon64 3200+ running Linux 2.6.24 with

glibc 2.7, tells me that it takes roughly 20,000,000 PyTuple_New(1) calls

when the tuple is retrieved from the freelist (i.e, call PyTuple_New() then

immediately DECREF it) and 1,200,000 calls when allocating brand new tuples

(i.e, call PyTuple_New() and store the returned value in an array), to save

one second (wall clock time). Either ways, the loading time a pickle stream

large enough to generate this number of calls would be massively

overwhelmed by other factors, like I/O throughput, the GC traversal and

object allocation overhead. So, I really doubt these functions provide any

real benefits.

On the other hand, oprofile reports that pickle spends a lot of time in

these functions. But, that is probably more related to the function call

overhead, than the argument tuple allocation.

XXX: And, what is the reference behavior of these? Steal, borrow? At first

glance, it seems to steal the reference of 'arg' and borrow the reference

of 'func'. */

static PyObject *

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

{

PyObject *result = NULL;

ARG_TUP(self, arg);

if (self->arg) {

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

FREE_ARG_TUP(self);

}

return result;

}

static int

_Pickler_ClearBuffer(PicklerObject *self)

{

Py_CLEAR(self->output_buffer);

self->output_buffer =

PyBytes_FromStringAndSize(NULL, self->max_output_len);

if (self->output_buffer == NULL)

return -1;

self->output_len = 0;

return 0;

}

static PyObject *

_Pickler_GetString(PicklerObject *self)

{

PyObject *output_buffer = self->output_buffer;

assert(self->output_buffer != NULL);

self->output_buffer = NULL;

/* Resize down to exact size */

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

return NULL;

return output_buffer;

}

static int

_Pickler_FlushToFile(PicklerObject *self)

{

PyObject *output, *result;

assert(self->write != NULL);

output = _Pickler_GetString(self);

if (output == NULL)

return -1;

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

Py_XDECREF(result);

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

}

static Py_ssize_t

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

{

Py_ssize_t i, required;

char *buffer;

assert(s != NULL);

required = self->output_len + n;

if (required > self->max_output_len) {

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

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

wasteful. */

if (_Pickler_FlushToFile(self) < 0)

return -1;

if (_Pickler_ClearBuffer(self) < 0)

return -1;

}

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

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

PyObject *result;

/* XXX we could spare an intermediate copy and pass

a memoryview instead */

PyObject *output = PyBytes_FromStringAndSize(s, n);

if (s == NULL)

return -1;

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

Py_XDECREF(result);

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

}

else {

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

PyErr_NoMemory();

return -1;

}

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

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

return -1;

}

}

buffer = PyBytes_AS_STRING(self->output_buffer);

if (n < 8) {

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

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

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

}

}

else {

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

}

self->output_len += n;

return n;

}

static PicklerObject *

_Pickler_New(void)

{

PicklerObject *self;

self = PyObject_GC_New(PicklerObject, &Pickler_Type);

if (self == NULL)

return NULL;

self->pers_func = NULL;

self->dispatch_table = NULL;

self->arg = NULL;

self->write = NULL;

self->proto = 0;

self->bin = 0;

self->fast = 0;

self->fast_nesting = 0;

self->fix_imports = 0;

self->fast_memo = NULL;

self->memo = PyMemoTable_New();

if (self->memo == NULL) {

Py_DECREF(self);

return NULL;

}

self->max_output_len = WRITE_BUF_SIZE;

self->output_len = 0;

self->output_buffer = PyBytes_FromStringAndSize(NULL,

self->max_output_len);

if (self->output_buffer == NULL) {

Py_DECREF(self);

return NULL;

}

return self;

}

static int

_Pickler_SetProtocol(PicklerObject *self, PyObject *proto_obj,

PyObject *fix_imports_obj)

{

long proto = 0;

int fix_imports;

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

proto = DEFAULT_PROTOCOL;

else {

proto = PyLong_AsLong(proto_obj);

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

return -1;

}

if (proto < 0)

proto = HIGHEST_PROTOCOL;

if (proto > HIGHEST_PROTOCOL) {

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

HIGHEST_PROTOCOL);

return -1;

}

fix_imports = PyObject_IsTrue(fix_imports_obj);

if (fix_imports == -1)

return -1;

self->proto = proto;

self->bin = proto > 0;

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

return 0;

}

/* Returns -1 (with an exception set) on failure, 0 on success. This may

be called once on a freshly created Pickler. */

static int

_Pickler_SetOutputStream(PicklerObject *self, PyObject *file)

{

_Py_IDENTIFIER(write);

assert(file != NULL);

self->write = _PyObject_GetAttrId(file, &PyId_write);

if (self->write == NULL) {

if (PyErr_ExceptionMatches(PyExc_AttributeError))

PyErr_SetString(PyExc_TypeError,

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

return -1;

}

return 0;

}

/* See documentation for _Pickler_FastCall(). */

static PyObject *

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

{

PyObject *result = NULL;

ARG_TUP(self, arg);

if (self->arg) {

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

FREE_ARG_TUP(self);

}

return result;

}

/* Returns the size of the input on success, -1 on failure. This takes its

own reference to `input`. */

static Py_ssize_t

_Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)

{

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

PyBuffer_Release(&self->buffer);

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

return -1;

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

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

self->next_read_idx = 0;

self->prefetched_idx = self->input_len;

return self->input_len;

}

static int

_Unpickler_SkipConsumed(UnpicklerObject *self)

{

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

if (consumed > 0) {

PyObject *r;

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

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

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

if (r == NULL)

return -1;

Py_DECREF(r);

self->prefetched_idx = self->next_read_idx;

}

return 0;

}

static const Py_ssize_t READ_WHOLE_LINE = -1;

/* If reading from a file, we need to only pull the bytes we need, since there

may be multiple pickle objects arranged contiguously in the same input

buffer.

If `n` is READ_WHOLE_LINE, read a whole line. Otherwise, read up to `n`

bytes from the input stream/buffer.

Update the unpickler's input buffer with the newly-read data. Returns -1 on

failure; on success, returns the number of bytes read from the file.

On success, self->input_len will be 0; this is intentional so that when

unpickling from a file, the "we've run out of data" code paths will trigger,

causing the Unpickler to go back to the file for more data. Use the returned

size to tell you how much data you can process. */

static Py_ssize_t

_Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)

{

PyObject *data;

Py_ssize_t read_size, prefetched_size = 0;

assert(self->read != NULL);

if (_Unpickler_SkipConsumed(self) < 0)

return -1;

if (n == READ_WHOLE_LINE)

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

else {

PyObject *len = PyLong_FromSsize_t(n);

if (len == NULL)

return -1;

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

}

if (data == NULL)

return -1;

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

if (self->peek) {

PyObject *len, *prefetched;

len = PyLong_FromSsize_t(PREFETCH);

if (len == NULL) {

Py_DECREF(data);

return -1;

}

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

if (prefetched == NULL) {

if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {

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

PyErr_Clear();

Py_CLEAR(self->peek);

}

else {

Py_DECREF(data);

return -1;

}

}

else {

assert(PyBytes_Check(prefetched));

prefetched_size = PyBytes_GET_SIZE(prefetched);

PyBytes_ConcatAndDel(&data, prefetched);

if (data == NULL)

return -1;

}

}

read_size = _Unpickler_SetStringInput(self, data) - prefetched_size;

Py_DECREF(data);

self->prefetched_idx = read_size;

return read_size;

}

/* Read `n` bytes from the unpickler's data source, storing the result in `*s`.

This should be used for all data reads, rather than accessing the unpickler's

input buffer directly. This method deals correctly with reading from input

streams, which the input buffer doesn't deal with.

Note that when reading from a file-like object, self->next_read_idx won't

be updated (it should remain at 0 for the entire unpickling process). You

should use this function's return value to know how many bytes you can

consume.

Returns -1 (with an exception set) on failure. On success, return the

number of chars read. */

static Py_ssize_t

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

{

Py_ssize_t num_read;

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

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

self->next_read_idx += n;

return n;

}

if (!self->read) {

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

return -1;

}

num_read = _Unpickler_ReadFromFile(self, n);

if (num_read < 0)

return -1;

if (num_read < n) {

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

return -1;

}

*s = self->input_buffer;

self->next_read_idx = n;