Add PyLong Import/Export API (#121)

skirpichev · vstinner · web-flow · commit 61709bfa512f · 2024-12-13T15:47:15.000+01:00
PyPy is not supported, as well as Python 2.  In the later case it's possible,
but hardly worth code complications: most real-word potential consumers
(e.g. Sage, gmpy2 or python-flint) support only Python 3.

Co-authored-by: Victor Stinner &lt;vstinner@python.org&gt;
diff --git a/docs/api.rst b/docs/api.rst
@@ -29,6 +29,42 @@ Latest version of the header file:
 Python 3.14
 -----------
 
+.. c:struct:: PyLongLayout
+
+   See `PyLongLayout documentation <https://docs.python.org/dev/c-api/long.html#c.PyLongLayout>`__.
+
+.. c:function:: const PyLongLayout* PyLong_GetNativeLayout(void)
+
+   See `PyLong_GetNativeLayout() documentation  <https://docs.python.org/dev/c-api/long.html#c.PyLong_GetNativeLayout>`__.
+
+.. c:struct:: PyLongExport
+
+   See `PyLongExport documentation <https://docs.python.org/dev/c-api/long.html#c.PyLongExport>`__.
+
+.. c:function:: int PyLong_Export(PyObject *obj, PyLongExport *export_long)
+
+   See `PyLong_Export() documentation  <https://docs.python.org/dev/c-api/long.html#c.PyLong_Export>`__.
+
+.. c:function:: void PyLong_FreeExport(PyLongExport *export_long)
+
+   See `PyLong_FreeExport() documentation  <https://docs.python.org/dev/c-api/long.html#c.PyLong_FreeExport>`__.
+
+.. c:struct:: PyLongWriter
+
+   See `PyLongWriter documentation <https://docs.python.org/dev/c-api/long.html#c.PyLongWriter>`__.
+
+.. c:function:: PyLongWriter* PyLongWriter_Create(int negative, Py_ssize_t ndigits, void **digits)
+
+   See `PyLongWriter_Create() documentation  <https://docs.python.org/dev/c-api/long.html#c.PyLongWriter_Create>`__.
+
+.. c:function:: PyObject* PyLongWriter_Finish(PyLongWriter *writer)
+
+   See `PyLongWriter_Finish() documentation  <https://docs.python.org/dev/c-api/long.html#c.PyLongWriter_Finish>`__.
+
+.. c:function:: void PyLongWriter_Discard(PyLongWriter *writer)
+
+   See `PyLongWriter_Discard() documentation  <https://docs.python.org/dev/c-api/long.html#c.PyLongWriter_Discard>`__.
+
 .. c:function:: int PyLong_IsPositive(PyObject *obj)
 
    See `PyLong_IsPositive() documentation  <https://docs.python.org/dev/c-api/long.html#c.PyLong_IsPositive>`__.
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -1,6 +1,18 @@
 Changelog
 =========
 
+* 2024-12-13: Add functions and structs:
+
+  * ``PyLongLayout``
+  * ``PyLong_GetNativeLayout()``
+  * ``PyLongExport``
+  * ``PyLong_Export()``
+  * ``PyLong_FreeExport()``
+  * ``PyLongWriter``
+  * ``PyLongWriter_Create()``
+  * ``PyLongWriter_Finish()``
+  * ``PyLongWriter_Discard()``
+
 * 2024-11-12: Add functions:
 
   * ``PyLong_IsPositive()``
diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h
@@ -1720,6 +1720,185 @@ static inline int PyLong_AsUInt64(PyObject *obj, uint64_t *pvalue)
 #endif
 
 
+// gh-102471 added import and export API for integers to 3.14.0a2.
+#if PY_VERSION_HEX < 0x030E00A2 && PY_VERSION_HEX >= 0x03000000 && !defined(PYPY_VERSION)
+// Helpers to access PyLongObject internals.
+static inline void
+_PyLong_SetSignAndDigitCount(PyLongObject *op, int sign, Py_ssize_t size)
+{
+#if PY_VERSION_HEX >= 0x030C0000
+    op->long_value.lv_tag = (uintptr_t)(1 - sign) | ((uintptr_t)(size) << 3);
+#elif PY_VERSION_HEX >= 0x030900A4
+    Py_SET_SIZE(op, sign * size);
+#else
+    Py_SIZE(op) = sign * size;
+#endif
+}
+
+static inline Py_ssize_t
+_PyLong_DigitCount(const PyLongObject *op)
+{
+#if PY_VERSION_HEX >= 0x030C0000
+    return (Py_ssize_t)(op->long_value.lv_tag >> 3);
+#else
+    return _PyLong_Sign((PyObject*)op) < 0 ? -Py_SIZE(op) : Py_SIZE(op);
+#endif
+}
+
+static inline digit*
+_PyLong_GetDigits(const PyLongObject *op)
+{
+#if PY_VERSION_HEX >= 0x030C0000
+    return (digit*)(op->long_value.ob_digit);
+#else
+    return (digit*)(op->ob_digit);
+#endif
+}
+
+typedef struct PyLongLayout {
+    uint8_t bits_per_digit;
+    uint8_t digit_size;
+    int8_t digits_order;
+    int8_t digit_endianness;
+} PyLongLayout;
+
+typedef struct PyLongExport {
+    int64_t value;
+    uint8_t negative;
+    Py_ssize_t ndigits;
+    const void *digits;
+    Py_uintptr_t _reserved;
+} PyLongExport;
+
+typedef struct PyLongWriter PyLongWriter;
+
+static inline const PyLongLayout*
+PyLong_GetNativeLayout(void)
+{
+    static const PyLongLayout PyLong_LAYOUT = {
+        PyLong_SHIFT,
+        sizeof(digit),
+        -1,  // least significant first
+        PY_LITTLE_ENDIAN ? -1 : 1,
+    };
+
+    return &PyLong_LAYOUT;
+}
+
+static inline int
+PyLong_Export(PyObject *obj, PyLongExport *export_long)
+{
+    if (!PyLong_Check(obj)) {
+        memset(export_long, 0, sizeof(*export_long));
+        PyErr_Format(PyExc_TypeError, "expected int, got %s",
+                     Py_TYPE(obj)->tp_name);
+        return -1;
+    }
+
+    // Fast-path: try to convert to a int64_t
+    PyLongObject *self = (PyLongObject*)obj;
+    int overflow;
+#if SIZEOF_LONG == 8
+    long value = PyLong_AsLongAndOverflow(obj, &overflow);
+#else
+    // Windows has 32-bit long, so use 64-bit long long instead
+    long long value = PyLong_AsLongLongAndOverflow(obj, &overflow);
+#endif
+    Py_BUILD_ASSERT(sizeof(value) == sizeof(int64_t));
+    // the function cannot fail since obj is a PyLongObject
+    assert(!(value == -1 && PyErr_Occurred()));
+
+    if (!overflow) {
+        export_long->value = value;
+        export_long->negative = 0;
+        export_long->ndigits = 0;
+        export_long->digits = 0;
+        export_long->_reserved = 0;
+    }
+    else {
+        export_long->value = 0;
+        export_long->negative = _PyLong_Sign(obj) < 0;
+        export_long->ndigits = _PyLong_DigitCount(self);
+        if (export_long->ndigits == 0) {
+            export_long->ndigits = 1;
+        }
+        export_long->digits = _PyLong_GetDigits(self);
+        export_long->_reserved = (Py_uintptr_t)Py_NewRef(obj);
+    }
+    return 0;
+}
+
+static inline void
+PyLong_FreeExport(PyLongExport *export_long)
+{
+    PyObject *obj = (PyObject*)export_long->_reserved;
+
+    if (obj) {
+        export_long->_reserved = 0;
+        Py_DECREF(obj);
+    }
+}
+
+static inline PyLongWriter*
+PyLongWriter_Create(int negative, Py_ssize_t ndigits, void **digits)
+{
+    if (ndigits <= 0) {
+        PyErr_SetString(PyExc_ValueError, "ndigits must be positive");
+        return NULL;
+    }
+    assert(digits != NULL);
+
+    PyLongObject *obj = _PyLong_New(ndigits);
+    if (obj == NULL) {
+        return NULL;
+    }
+    _PyLong_SetSignAndDigitCount(obj, negative?-1:1, ndigits);
+
+    *digits = _PyLong_GetDigits(obj);
+    return (PyLongWriter*)obj;
+}
+
+static inline void
+PyLongWriter_Discard(PyLongWriter *writer)
+{
+    PyLongObject *obj = (PyLongObject *)writer;
+
+    assert(Py_REFCNT(obj) == 1);
+    Py_DECREF(obj);
+}
+
+static inline PyObject*
+PyLongWriter_Finish(PyLongWriter *writer)
+{
+    PyObject *obj = (PyObject *)writer;
+    PyLongObject *self = (PyLongObject*)obj;
+    Py_ssize_t j = _PyLong_DigitCount(self);
+    Py_ssize_t i = j;
+    int sign = _PyLong_Sign(obj);
+
+    assert(Py_REFCNT(obj) == 1);
+
+    // Normalize and get singleton if possible
+    while (i > 0 && _PyLong_GetDigits(self)[i-1] == 0) {
+        --i;
+    }
+    if (i != j) {
+        if (i == 0) {
+            sign = 0;
+        }
+        _PyLong_SetSignAndDigitCount(self, sign, i);
+    }
+    if (i <= 1) {
+        long val = sign * (long)(_PyLong_GetDigits(self)[0]);
+        Py_DECREF(obj);
+        return PyLong_FromLong(val);
+    }
+
+    return obj;
+}
+#endif
+
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/tests/test_pythoncapi_compat_cext.c b/tests/test_pythoncapi_compat_cext.c
@@ -1426,6 +1426,50 @@ test_long_api(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
     assert(PyLong_IsNegative(obj) == 0);
     assert(PyLong_IsZero(obj) == 0);
 
+#if defined(PYTHON3) && !defined(PYPY_VERSION)
+    // test import/export API
+    digit *digits;
+    PyLongWriter *writer;
+    static PyLongExport long_export;
+
+    writer = PyLongWriter_Create(1, 1, (void**)&digits);
+    PyLongWriter_Discard(writer);
+
+    writer = PyLongWriter_Create(1, 1, (void**)&digits);
+    digits[0] = 123;
+    obj = PyLongWriter_Finish(writer);
+
+    check_int(obj, -123);
+    PyLong_Export(obj, &long_export);
+    assert(long_export.value == -123);
+    assert(long_export.digits == NULL);
+    PyLong_FreeExport(&long_export);
+    Py_DECREF(obj);
+
+    writer = PyLongWriter_Create(0, 5, (void**)&digits);
+    digits[0] = 1;
+    digits[1] = 0;
+    digits[2] = 0;
+    digits[3] = 0;
+    digits[4] = 1;
+    obj = PyLongWriter_Finish(writer);
+
+    PyLong_Export(obj, &long_export);
+    assert(long_export.value == 0);
+    digits = (digit*)long_export.digits;
+    assert(digits[0] == 1);
+    assert(digits[1] == 0);
+    assert(digits[2] == 0);
+    assert(digits[3] == 0);
+    assert(digits[4] == 1);
+    PyLong_FreeExport(&long_export);
+    Py_DECREF(obj);
+
+    const PyLongLayout *layout = PyLong_GetNativeLayout();
+    assert(layout->digits_order == -1);
+    assert(layout->digit_size == sizeof(digit));
+#endif // defined(PYTHON3) && !defined(PYPY_VERSION)
+
     Py_RETURN_NONE;
 }
 
