cPickle support for TUPLE[123]. Incidentally plugged several undetected

overflow holes in Pdata_grow().

Lib/pickle.py

@@ -621,8 +621,11 @@ class Pickler:
         proto = self.proto
 
         n = len(obj)
-        if n == 0 and proto:
+        if n == 0:
+            if proto:
                 write(EMPTY_TUPLE)
+            else:
+                write(MARK + TUPLE)
             return
 
         save = self.save
@@ -639,13 +642,13 @@ class Pickler:
                 self.memoize(obj)
             return
 
-        # proto 0, or proto 1 and tuple isn't empty, or proto > 1 and tuple
+        # proto 0 or proto 1 and tuple isn't empty, or proto > 1 and tuple
         # has more than 3 elements.
         write(MARK)
         for element in obj:
             save(element)
 
-        if n and id(obj) in memo:
+        if id(obj) in memo:
             # Subtle.  d was not in memo when we entered save_tuple(), so
             # the process of saving the tuple's elements must have saved
             # the tuple itself:  the tuple is recursive.  The proper action
@@ -660,9 +663,8 @@ class Pickler:
                 write(POP * (n+1) + get)
             return
 
-        # No recursion (including the empty-tuple case for protocol 0).
+        # No recursion.
         self.write(TUPLE)
-        if obj:                      # No need to memoize empty tuple
         self.memoize(obj)
 
     dispatch[TupleType] = save_tuple

Lib/test/pickletester.py

@@ -484,6 +484,27 @@ class AbstractPickleTests(unittest.TestCase):
         self.assertEqual(x, y)
 
     def test_short_tuples(self):
+        # Map (proto, len(tuple)) to expected opcode.
+        expected_opcode = {(0, 0): pickle.TUPLE,
+                           (0, 1): pickle.TUPLE,
+                           (0, 2): pickle.TUPLE,
+                           (0, 3): pickle.TUPLE,
+                           (0, 4): pickle.TUPLE,
+
+                           (1, 0): pickle.EMPTY_TUPLE,
+                           (1, 1): pickle.TUPLE,
+                           (1, 2): pickle.TUPLE,
+                           (1, 3): pickle.TUPLE,
+                           (1, 4): pickle.TUPLE,
+
+                           (2, 0): pickle.EMPTY_TUPLE,
+                           (2, 1): pickle.TUPLE1,
+                           (2, 2): pickle.TUPLE2,
+                           (2, 3): pickle.TUPLE3,
+                           (2, 4): pickle.TUPLE,
+                          }
+        all_tuple_opcodes = (pickle.TUPLE, pickle.EMPTY_TUPLE,
+                             pickle.TUPLE1, pickle.TUPLE2, pickle.TUPLE3)
         a = ()
         b = (1,)
         c = (1, 2)
@@ -495,6 +516,16 @@ class AbstractPickleTests(unittest.TestCase):
                 y = self.loads(s)
                 self.assertEqual(x, y, (proto, x, s, y))
 
+                # Verify that the protocol-correct tuple-building opcode
+                # was generated.
+                expected = expected_opcode[proto, len(x)]
+                for opcode in s:
+                    if opcode in all_tuple_opcodes:
+                        self.assertEqual(expected, opcode)
+                        break
+                else:
+                    self.fail("didn't find a tuple-building opcode in pickle")
+
     def test_singletons(self):
         for proto in protocols:
             for x in None, False, True:

Modules/cPickle.c

@@ -110,7 +110,8 @@ static PyObject *__class___str, *__getinitargs___str, *__dict___str,
 
 typedef struct {
 	PyObject_HEAD
-	int length, size;
+	int length;	/* number of initial slots in data currently used */
+	int size;	/* number of slots in data allocated */
 	PyObject **data;
 } Pdata;
 
@@ -120,10 +121,11 @@ Pdata_dealloc(Pdata *self)
 	int i;
 	PyObject **p;
 
-	for (i=self->length, p=self->data; --i >= 0; p++) Py_DECREF(*p);
-
-	if (self->data) free(self->data);
-
+	for (i = self->length, p = self->data; --i >= 0; p++) {
+		Py_DECREF(*p);
+	}
+	if (self->data)
+		free(self->data);
 	PyObject_Del(self);
 }
 
@@ -140,11 +142,13 @@ Pdata_New(void)
 {
 	Pdata *self;
 
-	if (!( self = PyObject_New(Pdata, &PdataType)))  return NULL;
-	self->size=8;
-	self->length=0;
-	self->data=malloc(self->size * sizeof(PyObject*));
-	if (self->data) return (PyObject*)self;
+	if (!(self = PyObject_New(Pdata, &PdataType)))
+		return NULL;
+	self->size = 8;
+	self->length = 0;
+	self->data = malloc(self->size * sizeof(PyObject*));
+	if (self->data)
+		return (PyObject*)self;
 	Py_DECREF(self);
 	return PyErr_NoMemory();
 }
@@ -156,6 +160,9 @@ stackUnderflow(void)
 	return -1;
 }
 
+/* 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, int clearto)
 {
@@ -165,52 +172,71 @@ Pdata_clear(Pdata *self, int clearto)
 	if (clearto < 0) return stackUnderflow();
 	if (clearto >= self->length) return 0;
 
-	for (i=self->length, p=self->data+clearto; --i >= clearto; p++)
+	for (i = self->length, p = self->data + clearto;
+	     --i >= clearto;
+	     p++) {
 		Py_DECREF(*p);
-	self->length=clearto;
+	}
+	self->length = clearto;
 
 	return 0;
 }
 
-
 static int
 Pdata_grow(Pdata *self)
 {
-	if (! self->size) {
-		PyErr_NoMemory();
-		return -1;
-	}
-	self->size *= 2;
-	self->data = realloc(self->data, self->size*sizeof(PyObject*));
-	if (! self->data) {
+	int bigger;
+	size_t nbytes;
+
+	if (! self->size)
+		goto nomemory;
+	bigger = self->size << 1;
+	if (bigger <= 0)
+		goto nomemory;
+	if ((int)(size_t)bigger != bigger)
+		goto nomemory;
+	nbytes = (size_t)bigger * sizeof(PyObject *);
+	if (nbytes / sizeof(PyObject *) != (size_t)bigger)
+		goto nomemory;
+	self->data = realloc(self->data, nbytes);
+	if (self->data == NULL)
+		goto nomemory;
+	self->size = bigger;
+	return 0;
+
+  nomemory:
 	self->size = 0;
 	PyErr_NoMemory();
 	return -1;
-	}
-	return 0;
 }
 
-#define PDATA_POP(D,V) {                                      \
-    if ((D)->length) V=D->data[--((D)->length)];              \
+/* 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.  D and V may be evaluated several times.
+ */
+#define PDATA_POP(D, V) {					\
+    if ((D)->length)						\
+    	(V) = (D)->data[--((D)->length)];			\
     else {							\
         PyErr_SetString(UnpicklingError, "bad pickle data");	\
-        V=NULL;                                               \
+        (V) = NULL;						\
     }								\
 }
 
-
 static PyObject *
 Pdata_popTuple(Pdata *self, int start)
 {
 	PyObject *r;
 	int i, j, l;
 
-	l=self->length-start;
-	if (!( r=PyTuple_New(l)))  return NULL;
-	for (i=start, j=0 ; j < l; i++, j++)
+	l = self->length-start;
+	r = PyTuple_New(l);
+	if (r == NULL)
+		return NULL;
+	for (i = start, j = 0 ; j < l; i++, j++)
 		PyTuple_SET_ITEM(r, j, self->data[i]);
 
-	self->length=start;
+	self->length = start;
 	return r;
 }
 
@@ -1444,81 +1470,142 @@ save_unicode(Picklerobject *self, PyObject *args, int doput)
 }
 #endif
 
+/* A helper for save_tuple.  Push the len elements in tuple t on the stack. */
+static int
+store_tuple_elememts(Picklerobject *self, PyObject *t, int len)
+{
+	int i;
+	int res = -1;	/* guilty until proved innocent */
+
+	assert(PyTuple_Size(t) == len);
+
+	for (i = 0; i < len; i++) {
+		PyObject *element = PyTuple_GET_ITEM(t, i);
+
+		if (element == NULL)
+			goto finally;
+		if (save(self, element, 0) < 0)
+			goto finally;
+	}
+	res = 0;
+
+  finally:
+	return res;
+}
 
+/* Tuples are ubiquitous in the pickle protocols, so many techniques are
+ * used across protocols to minimize the space needed to pickle them.
+ * Tuples are also the only builtin immuatable type that can be recursive
+ * (a tuple can be reached from itself), and that requires some subtle
+ * magic so that it works in all cases.  IOW, this is a long routine.
+ */
 static int
 save_tuple(Picklerobject *self, PyObject *args)
 {
-	PyObject *element = 0, *py_tuple_id = 0;
-	int len, i, res = -1;
+	PyObject *py_tuple_id = NULL;
+	int len, i;
+	int res = -1;
 
 	static char tuple = TUPLE;
+	static char pop = POP;
+	static char pop_mark = POP_MARK;
+	static char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
 
-	if (self->write_func(self, &MARKv, 1) < 0)
+	if ((len = PyTuple_Size(args)) < 0)
 		goto finally;
 
-	if ((len = PyTuple_Size(args)) < 0)
+	if (len == 0) {
+		char c_str[2];
+
+		if (self->proto) {
+			c_str[0] = EMPTY_TUPLE;
+			len = 1;
+		}
+		else {
+			c_str[0] = MARK;
+			c_str[1] = TUPLE;
+			len = 2;
+		}
+		if (self->write_func(self, c_str, len) >= 0)
+			res = 0;
+		/* Don't memoize an empty tuple. */
 		goto finally;
+	}
 
-	for (i = 0; i < len; i++) {
-		if (!( element = PyTuple_GET_ITEM((PyTupleObject *)args, i)))
+	/* A non-empty tuple. */
+
+	/* id(tuple) isn't in the memo now.  If it shows up there after
+	 * saving the tuple elements, the tuple must be recursive, in
+	 * which case we'll pop everything we put on the stack, and fetch
+	 * its value from the memo.
+	 */
+	py_tuple_id = PyLong_FromVoidPtr(args);
+	if (py_tuple_id == NULL)
 		goto finally;
 
-		if (save(self, element, 0) < 0)
+	if (len <= 3 && self->proto >= 2) {
+		/* Use TUPLE{1,2,3} opcodes. */
+		if (store_tuple_elememts(self, args, len) < 0)
+			goto finally;
+		if (PyDict_GetItem(self->memo, py_tuple_id)) {
+			/* pop the len elements */
+			for (i = 0; i < len; ++i)
+				if (self->write_func(self, &pop, 1) < 0)
+					goto finally;
+			/* fetch from memo */
+			if (get(self, py_tuple_id) < 0)
+				goto finally;
+			res = 0;
 			goto finally;
 		}
+		/* Not recursive. */
+		if (self->write_func(self, len2opcode + len, 1) < 0)
+			goto finally;
+		goto memoize;
+	}
+
+	/* proto < 2 and len > 0, or proto >= 2 and len > 3.
+	 * Generate MARK elt1 elt2 ... TUPLE
+	 */
+	if (self->write_func(self, &MARKv, 1) < 0)
+		goto finally;
 
-	if (!( py_tuple_id = PyLong_FromVoidPtr(args)))
+	if (store_tuple_elememts(self, args, len) < 0)
 		goto finally;
 
-	if (len) {
 	if (PyDict_GetItem(self->memo, py_tuple_id)) {
+		/* pop the stack stuff we pushed */
 		if (self->bin) {
-				static char pop_mark = POP_MARK;
-
 			if (self->write_func(self, &pop_mark, 1) < 0)
 				goto finally;
 		}
 		else {
-				static char pop = POP;
-
-				for (i = 0; i <= len; i++) {
+			/* Note that we pop one more than len, to remove
+			 * the MARK too.
+			 */
+			for (i = 0; i <= len; i++)
 				if (self->write_func(self, &pop, 1) < 0)
 					goto finally;
 		}
-			}
-
-			if (get(self, py_tuple_id) < 0)
-				goto finally;
-
+		/* fetch from memo */
+		if (get(self, py_tuple_id) >= 0)
 			res = 0;
 		goto finally;
 	}
-	}
 
-	if (self->write_func(self, &tuple, 1) < 0) {
-		goto finally;
-	}
-
-	if (put(self, args) < 0)
+	/* Not recursive. */
+	if (self->write_func(self, &tuple, 1) < 0)
 		goto finally;
 
+  memoize:
+	if (put(self, args) >= 0)
 		res = 0;
 
   finally:
 	Py_XDECREF(py_tuple_id);
-
 	return res;
 }
 
-static int
-save_empty_tuple(Picklerobject *self, PyObject *args)
-{
-	static char tuple = EMPTY_TUPLE;
-
-	return self->write_func(self, &tuple, 1);
-}
-
-
 static int
 save_list(Picklerobject *self, PyObject *args)
 {
@@ -2028,9 +2115,8 @@ save(Picklerobject *self, PyObject *args, int  pers_save)
 		break;
 
         case 't':
-		if (type == &PyTuple_Type && PyTuple_Size(args)==0) {
-			if (self->bin) res = save_empty_tuple(self, args);
-			else          res = save_tuple(self, args);
+		if (type == &PyTuple_Type && PyTuple_Size(args) == 0) {
+			res = save_tuple(self, args);
 			goto finally;
 		}
 		break;
@@ -3193,11 +3279,21 @@ load_tuple(Unpicklerobject *self)
 }
 
 static int
-load_empty_tuple(Unpicklerobject *self)
+load_counted_tuple(Unpicklerobject *self, int len)
 {
-	PyObject *tup;
+	PyObject *tup = PyTuple_New(len);
 
-	if (!( tup=PyTuple_New(0)))  return -1;
+	if (tup == NULL)
+		return -1;
+
+	while (--len >= 0) {
+		PyObject *element;
+
+		PDATA_POP(self->stack, element);
+		if (element == NULL)
+			return -1;
+		PyTuple_SET_ITEM(tup, len, element);
+	}
 	PDATA_PUSH(self->stack, tup, -1);
 	return 0;
 }
@@ -4018,7 +4114,22 @@ load(Unpicklerobject *self)
 #endif
 
 		case EMPTY_TUPLE:
-			if (load_empty_tuple(self) < 0)
+			if (load_counted_tuple(self, 0) < 0)
+				break;
+			continue;
+
+		case TUPLE1:
+			if (load_counted_tuple(self, 1) < 0)
+				break;
+			continue;
+
+		case TUPLE2:
+			if (load_counted_tuple(self, 2) < 0)
+				break;
+			continue;
+
+		case TUPLE3:
+			if (load_counted_tuple(self, 3) < 0)
 				break;
 			continue;
 
@@ -4346,7 +4457,22 @@ noload(Unpicklerobject *self)
 #endif
 
 		case EMPTY_TUPLE:
-			if (load_empty_tuple(self) < 0)
+			if (load_counted_tuple(self, 0) < 0)
+				break;
+			continue;
+
+		case TUPLE1:
+			if (load_counted_tuple(self, 1) < 0)
+				break;
+			continue;
+
+		case TUPLE2:
+			if (load_counted_tuple(self, 2) < 0)
+				break;
+			continue;
+
+		case TUPLE3:
+			if (load_counted_tuple(self, 3) < 0)
 				break;
 			continue;