Issue #17897: Optimized unpickle prefetching.

1 files changed, 22 insertions, 30 deletions

diff --git a/Modules/_pickle.c b/Modules/_pickle.c
index 6271c4b..a9fad0e 100644
--- a/Modules/_pickle.c
+++ b/Modules/_pickle.c
@@ -1121,7 +1121,7 @@ static Py_ssize_t
 _Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
 {
     PyObject *data;
-    Py_ssize_t read_size, prefetched_size = 0;
+    Py_ssize_t read_size;
 
     assert(self->read != NULL);
 
@@ -1134,46 +1134,38 @@ _Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
         Py_DECREF(empty_tuple);
     }
     else {
-        PyObject *len = PyLong_FromSsize_t(n);
-        if (len == NULL)
-            return -1;
-        data = _Pickle_FastCall(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 = _Pickle_FastCall(self->peek, len);
-        if (prefetched == NULL) {
-            if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {
+        PyObject *len;
+        /* Prefetch some data without advancing the file pointer, if possible */
+        if (self->peek && n < PREFETCH) {
+            len = PyLong_FromSsize_t(PREFETCH);
+            if (len == NULL)
+                return -1;
+            data = _Pickle_FastCall(self->peek, len);
+            if (data == NULL) {
+                if (!PyErr_ExceptionMatches(PyExc_NotImplementedError))
+                    return -1;
                 /* peek() is probably not supported by the given file object */
                 PyErr_Clear();
                 Py_CLEAR(self->peek);
             }
             else {
+                read_size = _Unpickler_SetStringInput(self, data);
                 Py_DECREF(data);
-                return -1;
+                self->prefetched_idx = 0;
+                if (n <= read_size)
+                    return n;
             }
         }
-        else {
-            assert(PyBytes_Check(prefetched));
-            prefetched_size = PyBytes_GET_SIZE(prefetched);
-            PyBytes_ConcatAndDel(&data, prefetched);
-            if (data == NULL)
-                return -1;
-        }
+        len = PyLong_FromSsize_t(n);
+        if (len == NULL)
+            return -1;
+        data = _Pickle_FastCall(self->read, len);
     }
+    if (data == NULL)
+        return -1;
 
-    read_size = _Unpickler_SetStringInput(self, data) - prefetched_size;
+    read_size = _Unpickler_SetStringInput(self, data);
     Py_DECREF(data);
-    self->prefetched_idx = read_size;
     return read_size;
 }