diff -r a0b2b5abf88c Python/marshal.c --- a/Python/marshal.c Wed Jan 28 17:11:12 2015 +0200 +++ b/Python/marshal.c Wed Jan 28 17:14:14 2015 +0200 @@ -72,7 +72,12 @@ typedef struct { char *end; char *buf; Py_ssize_t buf_size; + unsigned int refs_count; PyObject *refs; /* dict on marshal, list on unmarshal */ + PyObject *int_refs; + PyObject *float_refs; + PyObject *complex_refs; + PyObject *str_refs; int version; } WFILE; @@ -223,8 +228,9 @@ w_PyLong(const PyLongObject *ob, char fl static int w_ref(PyObject *v, char *flag, WFILE *p) { - PyObject *id; + PyObject *id = v; PyObject *idx; + PyObject *refs = NULL; if (p->version < 3 || p->refs == NULL) return 0; /* not writing object references */ @@ -233,14 +239,30 @@ w_ref(PyObject *v, char *flag, WFILE *p) if (Py_REFCNT(v) == 1) return 0; - id = PyLong_FromVoidPtr((void*)v); - if (id == NULL) - goto err; - idx = PyDict_GetItem(p->refs, id); + if (PyUnicode_CheckExact(v)) { + if (PyUnicode_GET_LENGTH(v) < 1000) + refs = p->str_refs; + } + else if (PyLong_CheckExact(v)) + refs = p->int_refs; + else if (PyFloat_CheckExact(v)) + refs = p->float_refs; + else if (PyComplex_CheckExact(v)) + refs = p->complex_refs; + + if (refs == NULL) { + refs = p->refs; + id = PyLong_FromVoidPtr((void*)v); + if (id == NULL) + goto err; + } + + idx = PyDict_GetItem(refs, id); if (idx != NULL) { /* write the reference index to the stream */ long w = PyLong_AsLong(idx); - Py_DECREF(id); + if (id != v) + Py_DECREF(id); if (w == -1 && PyErr_Occurred()) { goto err; } @@ -251,18 +273,20 @@ w_ref(PyObject *v, char *flag, WFILE *p) return 1; } else { int ok; - Py_ssize_t s = PyDict_Size(p->refs); + Py_ssize_t s = p->refs_count; /* we don't support long indices */ if (s >= 0x7fffffff) { PyErr_SetString(PyExc_ValueError, "too many objects"); goto err; } idx = PyLong_FromSsize_t(s); - ok = idx && PyDict_SetItem(p->refs, id, idx) == 0; - Py_DECREF(id); + ok = idx && PyDict_SetItem(refs, id, idx) == 0; + if (id != v) + Py_DECREF(id); Py_XDECREF(idx); if (!ok) goto err; + p->refs_count++; *flag |= FLAG_REF; return 0; } @@ -558,6 +582,41 @@ PyMarshal_WriteLongToFile(long x, FILE * w_long(x, &wf); } +static int +w_init_refs(WFILE *wf, int version) +{ + wf->refs = NULL; + wf->int_refs = NULL; + wf->float_refs = NULL; + wf->complex_refs = NULL; + wf->str_refs = NULL; + if (version >= 3) { + wf->refs_count = 0; + if ((wf->int_refs = PyDict_New()) == NULL || + (wf->float_refs = PyDict_New()) == NULL || + (wf->complex_refs = PyDict_New()) == NULL || + (wf->str_refs = PyDict_New()) == NULL || + (wf->refs = PyDict_New()) == NULL) { + Py_CLEAR(wf->int_refs); + Py_CLEAR(wf->float_refs); + Py_CLEAR(wf->complex_refs); + Py_CLEAR(wf->str_refs); + return -1; + } + } + return 0; +} + +static void +w_clear(WFILE *wf) +{ + Py_XDECREF(wf->refs); + Py_XDECREF(wf->int_refs); + Py_XDECREF(wf->float_refs); + Py_XDECREF(wf->complex_refs); + Py_XDECREF(wf->str_refs); +} + void PyMarshal_WriteObjectToFile(PyObject *x, FILE *fp, int version) { @@ -565,14 +624,11 @@ PyMarshal_WriteObjectToFile(PyObject *x, wf.fp = fp; wf.error = WFERR_OK; wf.depth = 0; - if (version >= 3) { - if ((wf.refs = PyDict_New()) == NULL) - return; /* caller mush check PyErr_Occurred() */ - } else - wf.refs = NULL; + if (w_init_refs(&wf, version)) + return; /* caller mush check PyErr_Occurred() */ wf.version = version; w_object(x, &wf); - Py_XDECREF(wf.refs); + w_clear(&wf); } typedef WFILE RFILE; /* Same struct with different invariants */ @@ -1519,15 +1575,12 @@ PyMarshal_WriteObjectToString(PyObject * wf.error = WFERR_OK; wf.depth = 0; wf.version = version; - if (version >= 3) { - if ((wf.refs = PyDict_New()) == NULL) { - Py_DECREF(wf.str); - return NULL; - } - } else - wf.refs = NULL; + if (w_init_refs(&wf, version)) { + Py_DECREF(wf.str); + return NULL; + } w_object(x, &wf); - Py_XDECREF(wf.refs); + w_clear(&wf); if (wf.str != NULL) { char *base = PyBytes_AS_STRING((PyBytesObject *)wf.str); if (wf.ptr - base > PY_SSIZE_T_MAX) {