From 6fe39b76a9b353ad619108ad2877a155680270dc Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Sat, 30 Nov 2013 23:15:38 +0200
Subject: [PATCH] Issue #17897: Optimized unpickle prefetching.

---
 Misc/NEWS         |  2 ++
 Modules/_pickle.c | 58 ++++++++++++++++++++---------------------------
 2 files changed, 27 insertions(+), 33 deletions(-)

diff --git a/Misc/NEWS b/Misc/NEWS
index 09d2516b1a6..653b0f43260 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -18,6 +18,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #17897: Optimized unpickle prefetching.
+
 - Issue #3693: Make the error message more helpful when the array.array()
   constructor is given a str. Move the array module typecode documentation to
   the docstring of the constructor.
diff --git a/Modules/_pickle.c b/Modules/_pickle.c
index 6271c4b33a4..a9fad0e9495 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,7 +1134,29 @@ _Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
         Py_DECREF(empty_tuple);
     }
     else {
-        PyObject *len = PyLong_FromSsize_t(n);
+        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);
+                self->prefetched_idx = 0;
+                if (n <= read_size)
+                    return n;
+            }
+        }
+        len = PyLong_FromSsize_t(n);
         if (len == NULL)
             return -1;
         data = _Pickle_FastCall(self->read, len);
@@ -1142,38 +1164,8 @@ _Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
     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)) {
-                /* 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;
+    read_size = _Unpickler_SetStringInput(self, data);
     Py_DECREF(data);
-    self->prefetched_idx = read_size;
     return read_size;
 }