Add PyLong_AsDigitArray.value

* Rename functions and structure: * PyLong_AsDigitArray() => PyLong_Export() * PyLong_FreeDigitArray() => PyLong_FreeExport() * PyLong_DigitArray => PyLongExport * 'array' => 'export_long'
vstinner · Sep 17, 2024 · 37b1d49 · 37b1d49
1 parent 4221a49
commit 37b1d49
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 48 deletions.
diff --git a/Doc/c-api/long.rst b/Doc/c-api/long.rst
@@ -686,7 +686,7 @@ Export API
       Read-only array of unsigned digits.
 
 
-.. c:function:: int PyLong_AsDigitArray(PyObject *obj, PyLong_DigitArray *array)
+.. c:function:: int PyLong_Export(PyObject *obj, PyLong_DigitArray *array)
 
    Export a Python :class:`int` object as an array of digits.
 
@@ -696,13 +696,13 @@ Export API
    This function always succeeds if *obj* is a Python :class:`int` object or a
    subclass.
 
-   :c:func:`PyLong_FreeDigitArray` must be called once done with using
+   :c:func:`PyLong_FreeExport` must be called once done with using
    *array*.
 
 
-.. c:function:: void PyLong_FreeDigitArray(PyLong_DigitArray *array)
+.. c:function:: void PyLong_FreeExport(PyLong_DigitArray *array)
 
-   Release the export *array* created by :c:func:`PyLong_AsDigitArray`.
+   Release the export *array* created by :c:func:`PyLong_Export`.
 
 
 PyLongWriter API

diff --git a/Include/cpython/longintrepr.h b/Include/cpython/longintrepr.h
@@ -163,19 +163,20 @@ typedef struct PyLongLayout {
 
 PyAPI_FUNC(const PyLongLayout*) PyLong_GetNativeLayout(void);
 
-typedef struct PyLong_DigitArray {
-    int negative;
+typedef struct PyLongExport {
+    int64_t value;
+    uint8_t negative;
     Py_ssize_t ndigits;
     const void *digits;
     // Member used internally, must not be used for other purpose.
     Py_uintptr_t _reserved;
-} PyLong_DigitArray;
+} PyLongExport;
 
-PyAPI_FUNC(int) PyLong_AsDigitArray(
+PyAPI_FUNC(int) PyLong_Export(
     PyObject *obj,
-    PyLong_DigitArray *array);
-PyAPI_FUNC(void) PyLong_FreeDigitArray(
-    PyLong_DigitArray *array);
+    PyLongExport *export_long);
+PyAPI_FUNC(void) PyLong_FreeExport(
+    PyLongExport *export_long);
 
 
 /* --- PyLongWriter API --------------------------------------------------- */

diff --git a/Lib/test/test_capi/test_long.py b/Lib/test/test_capi/test_long.py
@@ -10,6 +10,7 @@
 
 NULL = None
 
+
 class IntSubclass(int):
     pass
 
@@ -686,19 +687,23 @@ def test_long_export(self):
         layout = _testcapi.get_pylong_layout()
         base = 2 ** layout['bits_per_digit']
 
-        pylong_asdigitarray = _testcapi.pylong_asdigitarray
-        self.assertEqual(pylong_asdigitarray(0), (0, [0]))
-        self.assertEqual(pylong_asdigitarray(123), (0, [123]))
-        self.assertEqual(pylong_asdigitarray(-123), (1, [123]))
-        self.assertEqual(pylong_asdigitarray(base**2 * 3 + base * 2 + 1),
-                         (0, [1, 2, 3]))
+        pylong_export = _testcapi.pylong_export
 
-        with self.assertRaises(TypeError):
-            pylong_asdigitarray(1.0)
-        with self.assertRaises(TypeError):
-            pylong_asdigitarray(0+1j)
-        with self.assertRaises(TypeError):
-            pylong_asdigitarray("abc")
+        # value fits into int64_t
+        self.assertEqual(pylong_export(0), 0)
+        self.assertEqual(pylong_export(123), 123)
+        self.assertEqual(pylong_export(-123), -123)
+
+        # use an array, doesn't fit into int64_t
+        self.assertEqual(pylong_export(base**10 * 2 + 1),
+                         (0, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]))
+        self.assertEqual(pylong_export(-(base**10 * 2 + 1)),
+                         (1, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]))
+
+        for value in (1.0, 0+1j, "abc"):
+            with self.subTest(value=value):
+                with self.assertRaises(TypeError):
+                    pylong_export(value)
 
     def test_longwriter_create(self):
         # Test PyLong_Import()
@@ -726,13 +731,28 @@ def test_longwriter_create(self):
             self.assertIs(pylongwriter_create(bool(num < 0), [abs(num), 0]),
                           num)
 
+        def to_digits(num):
+            digits = []
+            while True:
+                num, digit = divmod(num, base)
+                digits.append(digit)
+                if not num:
+                    break
+            return digits
+
         # round trip: Python int -> export -> Python int
-        pylong_asdigitarray = _testcapi.pylong_asdigitarray
+        pylong_export = _testcapi.pylong_export
         numbers = [*range(0, 10), 12345, 0xdeadbeef, 2**100, 2**100-1]
         numbers.extend(-num for num in list(numbers))
         for num in numbers:
             with self.subTest(num=num):
-                negative, digits = pylong_asdigitarray(num)
+                data = pylong_export(num)
+                if isinstance(data, tuple):
+                    negative, digits = data
+                else:
+                    value = data
+                    negative = int(value < 0)
+                    digits = to_digits(abs(value))
                 self.assertEqual(pylongwriter_create(negative, digits), num,
                                  (negative, digits))
 

diff --git a/Modules/_testcapi/long.c b/Modules/_testcapi/long.c
@@ -174,19 +174,24 @@ layout_to_dict(const PyLongLayout *layout)
 
 
 static PyObject *
-pylong_asdigitarray(PyObject *module, PyObject *obj)
+pylong_export(PyObject *module, PyObject *obj)
 {
-    PyLong_DigitArray array;
-    if (PyLong_AsDigitArray(obj, &array) < 0) {
+    PyLongExport export_long;
+    if (PyLong_Export(obj, &export_long) < 0) {
         return NULL;
     }
 
+    if (export_long.digits == NULL) {
+        return PyLong_FromInt64(export_long.value);
+        // PyLong_FreeExport() is not needed in this case
+    }
+
     assert(PyLong_GetNativeLayout()->digit_size == sizeof(Py_digit));
-    const Py_digit *array_digits = array.digits;
+    const Py_digit *export_long_digits = export_long.digits;
 
     PyObject *digits = PyList_New(0);
-    for (Py_ssize_t i=0; i < array.ndigits; i++) {
-        PyObject *digit = PyLong_FromUnsignedLong(array_digits[i]);
+    for (Py_ssize_t i=0; i < export_long.ndigits; i++) {
+        PyObject *digit = PyLong_FromUnsignedLong(export_long_digits[i]);
         if (digit == NULL) {
             goto error;
         }
@@ -198,16 +203,16 @@ pylong_asdigitarray(PyObject *module, PyObject *obj)
         Py_DECREF(digit);
     }
 
-    PyObject *res = Py_BuildValue("(iN)", array.negative, digits);
+    PyObject *res = Py_BuildValue("(iN)", export_long.negative, digits);
 
-    PyLong_FreeDigitArray(&array);
-    assert(array._reserved == 0);
+    PyLong_FreeExport(&export_long);
+    assert(export_long._reserved == 0);
 
     return res;
 
 error:
     Py_DECREF(digits);
-    PyLong_FreeDigitArray(&array);
+    PyLong_FreeExport(&export_long);
     return NULL;
 }
 
@@ -280,7 +285,7 @@ static PyMethodDef test_methods[] = {
     {"pylong_fromnativebytes",      pylong_fromnativebytes,     METH_VARARGS},
     {"pylong_getsign",              pylong_getsign,             METH_O},
     {"pylong_aspid",                pylong_aspid,               METH_O},
-    {"pylong_asdigitarray",         pylong_asdigitarray,        METH_O},
+    {"pylong_export",               pylong_export,              METH_O},
     {"pylongwriter_create",         pylongwriter_create,        METH_VARARGS},
     {"get_pylong_layout",           get_pylong_layout,          METH_NOARGS},
     {NULL},

diff --git a/Objects/longobject.c b/Objects/longobject.c
@@ -6791,31 +6791,43 @@ const PyLongLayout* PyLong_GetNativeLayout(void)
 
 
 int
-PyLong_AsDigitArray(PyObject *obj, PyLong_DigitArray *array)
+PyLong_Export(PyObject *obj, PyLongExport *export_long)
 {
     if (!PyLong_Check(obj)) {
         PyErr_Format(PyExc_TypeError, "expect int, got %T", obj);
         return -1;
     }
-    PyLongObject *self = (PyLongObject*)obj;
+    int64_t value;
+    if (PyLong_AsInt64(obj, &value) >= 0) {
+        export_long->value = value;
+        export_long->negative = 0;
+        export_long->ndigits = 0;
+        export_long->digits = 0;
+        export_long->_reserved = 0;
+    }
+    else {
+        PyErr_Clear();
 
-    array->negative = _PyLong_IsNegative(self);
-    array->ndigits = _PyLong_DigitCount(self);
-    if (array->ndigits == 0) {
-        array->ndigits = 1;
+        PyLongObject *self = (PyLongObject*)obj;
+        export_long->value = 0;
+        export_long->negative = _PyLong_IsNegative(self);
+        export_long->ndigits = _PyLong_DigitCount(self);
+        if (export_long->ndigits == 0) {
+            export_long->ndigits = 1;
+        }
+        export_long->digits = self->long_value.ob_digit;
+        export_long->_reserved = (Py_uintptr_t)Py_NewRef(obj);
     }
-    array->digits = self->long_value.ob_digit;
-    array->_reserved = (Py_uintptr_t)Py_NewRef(obj);
     return 0;
 }
 
 
 void
-PyLong_FreeDigitArray(PyLong_DigitArray *array)
+PyLong_FreeExport(PyLongExport *export_long)
 {
-    PyObject *obj = (PyObject*)array->_reserved;
-    array->_reserved = 0;
-    Py_DECREF(obj);
+    PyObject *obj = (PyObject*)export_long->_reserved;
+    export_long->_reserved = 0;
+    Py_XDECREF(obj);
 }