Added new private API function _PyLong_NumBits. This will be used at the

start for the C implemention of new pickle LONG1 and LONG4 opcodes (the
linear-time way to pickle a long is to call _PyLong_AsByteArray, but
the caller has no idea how big an array to allocate, and correct
calculation is a bit subtle).

3 files changed, 85 insertions, 2 deletions

diff --git a/Include/longobject.h b/Include/longobject.h
index e452f516..3b808fb 100644
--- a/Include/longobject.h
+++ b/Include/longobject.h
@@ -44,6 +44,17 @@ PyAPI_FUNC(PyObject *) PyLong_FromString(char *, char **, int);
 PyAPI_FUNC(PyObject *) PyLong_FromUnicode(Py_UNICODE*, int, int);
 #endif
 
+/* _PyLong_NumBits.  Return the number of bits needed to represent a long
+   in contiguous 2's-complement form, including 1 for the sign bit.  For
+   example, this returns 1 for 0, and 2 for 1 and -1.  Note that the
+   ceiling of this divided by 8 is the number of bytes needed by
+   _PyLong_AsByteArray to store the long in 256's-complement form.
+   v must not be NULL, and must be a normalized long.
+   (size_t)-1 is returned and OverflowError set if the true result doesn't
+   fit in a size_t.
+*/
+PyAPI_FUNC(size_t) _PyLong_NumBits(PyObject *v);
+
 /* _PyLong_FromByteArray:  View the n unsigned bytes as a binary integer in
    base 256, and return a Python long with the same numeric value.
    If n is 0, the integer is 0.  Else:
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 2054c80..9359188 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -36,7 +36,7 @@ sizeof_error(const char* fatname, const char* typename,
         int expected, int got)
 {
 	char buf[1024];
-	PyOS_snprintf(buf, sizeof(buf), 
+	PyOS_snprintf(buf, sizeof(buf),
 		"%.200s #define == %d but sizeof(%.200s) == %d",
 		fatname, expected, typename, got);
 	PyErr_SetString(TestError, buf);
@@ -326,7 +326,7 @@ test_u_code(PyObject *self)
 	    len != PyUnicode_GET_SIZE(obj))
         	return raiseTestError("test_u_code",
 			"u# code returned wrong values for u'test'");
-	
+
 	Py_DECREF(tuple);
 	Py_INCREF(Py_None);
 	return Py_None;
@@ -334,6 +334,42 @@ test_u_code(PyObject *self)
 
 #endif
 
+/* Simple test of _PyLong_NumBits. */
+static PyObject *
+test_long_numbits(PyObject *self)
+{
+	struct pair {
+		long input;
+		size_t output;
+	} testcases[] = {{0, 1},
+			 {1L, 2},
+			 {-1L, 2},
+			 {2L, 3},
+			 {-2L, 3},
+			 {3L, 3},
+			 {-3L, 3},
+			 {4L, 4},
+			 {-4L, 4},
+			 {0x7fffL, 16},		/* one Python long digit */
+			 {-0x7fffL, 16},
+			 {0xfffffffL, 29},
+			 {-0xfffffffL, 29}};
+	int i;
+
+	for (i = 0; i < sizeof(testcases) / sizeof(struct pair); ++i) {
+		long input = testcases[i].input;
+		PyObject *plong = PyLong_FromLong(input);
+		size_t nbits = _PyLong_NumBits(plong);
+
+		Py_DECREF(plong);
+		if (nbits != testcases[i].output)
+			return raiseTestError("test_long_numbits",
+					      "wrong result");
+	}
+	Py_INCREF(Py_None);
+	return Py_None;
+}
+
 static PyObject *
 raise_exception(PyObject *self, PyObject *args)
 {
@@ -366,6 +402,7 @@ static PyMethodDef TestMethods[] = {
 	{"test_list_api",	(PyCFunction)test_list_api,	 METH_NOARGS},
 	{"test_dict_iteration",	(PyCFunction)test_dict_iteration,METH_NOARGS},
 	{"test_long_api",	(PyCFunction)test_long_api,	 METH_NOARGS},
+	{"test_long_numbits",	(PyCFunction)test_long_numbits,	 METH_NOARGS},
 #ifdef HAVE_LONG_LONG
 	{"test_longlong_api",	(PyCFunction)test_longlong_api,	 METH_NOARGS},
 	{"test_L_code",		(PyCFunction)test_L_code,	 METH_NOARGS},
diff --git a/Objects/longobject.c b/Objects/longobject.c
index cb27e79..1180ec2 100644
--- a/Objects/longobject.c
+++ b/Objects/longobject.c
@@ -260,6 +260,41 @@ PyLong_AsUnsignedLong(PyObject *vv)
 	return x;
 }
 
+size_t
+_PyLong_NumBits(PyObject *vv)
+{
+	PyLongObject *v = (PyLongObject *)vv;
+	size_t result = 1;	/* for the sign bit */
+	size_t ndigits = ABS(v->ob_size);
+
+	assert(v != NULL);
+	assert(PyLong_Check(v));
+	assert(ndigits == 0 || v->ob_digit[ndigits - 1] != 0);
+	if (ndigits > 0) {
+		size_t product;
+		digit msd = v->ob_digit[ndigits - 1];
+
+		product = (ndigits - 1) * SHIFT;
+		if (product / SHIFT != ndigits - 1)
+			goto Overflow;
+		result += product;
+		if (result < product)
+			goto Overflow;
+		do {
+			++result;
+			if (result == 0)
+				goto Overflow;
+			msd >>= 1;
+		} while (msd);
+	}
+	return result;
+
+Overflow:
+	PyErr_SetString(PyExc_OverflowError, "long has too many bits "
+			"to express in a platform size_t");
+	return (size_t)-1;
+}
+
 PyObject *
 _PyLong_FromByteArray(const unsigned char* bytes, size_t n,
 		      int little_endian, int is_signed)