[mypyc] Fix int operations on Python 3.12 (#15470)

The representation of `int` objects was changed in Python 3.12.

Define some helper macros to make it possible to support all Python
versions we care about without too much code duplication.

Work on mypyc/mypyc#995.

Co-authored-by: Shantanu <[email protected]>

Author: JukkaL

mypyc/lib-rt/int_ops.c

@@ -308,10 +308,10 @@ PyObject *CPyBool_Str(bool b) {
 }
 
 static void CPyLong_NormalizeUnsigned(PyLongObject *v) {
-    Py_ssize_t i = v->ob_base.ob_size;
-    while (i > 0 && v->ob_digit[i - 1] == 0)
+    Py_ssize_t i = CPY_LONG_SIZE_UNSIGNED(v);
+    while (i > 0 && CPY_LONG_DIGIT(v, i - 1) == 0)
         i--;
-    v->ob_base.ob_size = i;
+    CPyLong_SetUnsignedSize(v, i);
 }
 
 // Bitwise op '&', '|' or '^' using the generic (slow) API
@@ -361,8 +361,8 @@ static digit *GetIntDigits(CPyTagged n, Py_ssize_t *size, digit *buf) {
         return buf;
     } else {
         PyLongObject *obj = (PyLongObject *)CPyTagged_LongAsObject(n);
-        *size = obj->ob_base.ob_size;
-        return obj->ob_digit;
+        *size = CPY_LONG_SIZE_SIGNED(obj);
+        return &CPY_LONG_DIGIT(obj, 0);
     }
 }
 
@@ -399,20 +399,20 @@ static CPyTagged BitwiseLongOp(CPyTagged a, CPyTagged b, char op) {
     Py_ssize_t i;
     if (op == '&') {
         for (i = 0; i < asize; i++) {
-            r->ob_digit[i] = adigits[i] & bdigits[i];
+            CPY_LONG_DIGIT(r, i) = adigits[i] & bdigits[i];
         }
     } else {
         if (op == '|') {
             for (i = 0; i < asize; i++) {
-                r->ob_digit[i] = adigits[i] | bdigits[i];
+                CPY_LONG_DIGIT(r, i) = adigits[i] | bdigits[i];
             }
         } else {
             for (i = 0; i < asize; i++) {
-                r->ob_digit[i] = adigits[i] ^ bdigits[i];
+                CPY_LONG_DIGIT(r, i) = adigits[i] ^ bdigits[i];
             }
         }
         for (; i < bsize; i++) {
-            r->ob_digit[i] = bdigits[i];
+            CPY_LONG_DIGIT(r, i) = bdigits[i];
         }
     }
     CPyLong_NormalizeUnsigned(r);
@@ -521,7 +521,7 @@ int64_t CPyLong_AsInt64(PyObject *o) {
         Py_ssize_t size = Py_SIZE(lobj);
         if (likely(size == 1)) {
             // Fast path
-            return lobj->ob_digit[0];
+            return CPY_LONG_DIGIT(lobj, 0);
         } else if (likely(size == 0)) {
             return 0;
         }
@@ -576,14 +576,25 @@ int64_t CPyInt64_Remainder(int64_t x, int64_t y) {
 
 int32_t CPyLong_AsInt32(PyObject *o) {
     if (likely(PyLong_Check(o))) {
+    #if CPY_3_12_FEATURES
+        PyLongObject *lobj = (PyLongObject *)o;
+        size_t tag = CPY_LONG_TAG(lobj);
+        if (likely(tag == (1 << CPY_NON_SIZE_BITS))) {
+            // Fast path
+            return CPY_LONG_DIGIT(lobj, 0);
+        } else if (likely(tag == CPY_SIGN_ZERO)) {
+            return 0;
+        }
+    #else
         PyLongObject *lobj = (PyLongObject *)o;
         Py_ssize_t size = lobj->ob_base.ob_size;
         if (likely(size == 1)) {
             // Fast path
-            return lobj->ob_digit[0];
+            return CPY_LONG_DIGIT(lobj, 0);
         } else if (likely(size == 0)) {
             return 0;
         }
+    #endif
     }
     // Slow path
     int overflow;

mypyc/lib-rt/mypyc_util.h

@@ -72,4 +72,44 @@ static inline CPyTagged CPyTagged_ShortFromSsize_t(Py_ssize_t x) {
 // Are we targeting Python 3.12 or newer?
 #define CPY_3_12_FEATURES (PY_VERSION_HEX >= 0x030c0000)
 
+#if CPY_3_12_FEATURES
+
+// Same as macros in CPython internal/pycore_long.h, but with a CPY_ prefix
+#define CPY_NON_SIZE_BITS 3
+#define CPY_SIGN_ZERO 1
+#define CPY_SIGN_NEGATIVE 2
+#define CPY_SIGN_MASK 3
+
+#define CPY_LONG_DIGIT(o, n) ((o)->long_value.ob_digit[n])
+
+// Only available on Python 3.12 and later
+#define CPY_LONG_TAG(o) ((o)->long_value.lv_tag)
+#define CPY_LONG_IS_NEGATIVE(o) (((o)->long_value.lv_tag & CPY_SIGN_MASK) == CPY_SIGN_NEGATIVE)
+// Only available on Python 3.12 and later
+#define CPY_LONG_SIZE(o) ((o)->long_value.lv_tag >> CPY_NON_SIZE_BITS)
+// Number of digits; negative for negative ints
+#define CPY_LONG_SIZE_SIGNED(o) (CPY_LONG_IS_NEGATIVE(o) ? -CPY_LONG_SIZE(o) : CPY_LONG_SIZE(o))
+// Number of digits, assuming int is non-negative
+#define CPY_LONG_SIZE_UNSIGNED(o) CPY_LONG_SIZE(o)
+
+static inline void CPyLong_SetUnsignedSize(PyLongObject *o, Py_ssize_t n) {
+    if (n == 0)
+        o->long_value.lv_tag = CPY_SIGN_ZERO;
+    else
+        o->long_value.lv_tag = n << CPY_NON_SIZE_BITS;
+}
+
+#else
+
+#define CPY_LONG_DIGIT(o, n) ((o)->ob_digit[n])
+#define CPY_LONG_IS_NEGATIVE(o) (((o)->ob_base.ob_size < 0)
+#define CPY_LONG_SIZE_SIGNED(o) ((o)->ob_base.ob_size)
+#define CPY_LONG_SIZE_UNSIGNED(o) ((o)->ob_base.ob_size)
+
+static inline void CPyLong_SetUnsignedSize(PyLongObject *o, Py_ssize_t n) {
+    o->ob_base.ob_size = n;
+}
+
+#endif
+
 #endif

mypyc/lib-rt/pythonsupport.h

@@ -129,6 +129,65 @@ init_subclass(PyTypeObject *type, PyObject *kwds)
     return 0;
 }
 
+#if CPY_3_12_FEATURES
+
+static inline Py_ssize_t
+CPyLong_AsSsize_tAndOverflow(PyObject *vv, int *overflow)
+{
+    /* This version by Tim Peters */
+    PyLongObject *v = (PyLongObject *)vv;
+    size_t x, prev;
+    Py_ssize_t res;
+    Py_ssize_t i;
+    int sign;
+
+    *overflow = 0;
+
+    res = -1;
+    i = CPY_LONG_TAG(v);
+
+    // TODO: Combine zero and non-zero cases helow?
+    if (likely(i == (1 << CPY_NON_SIZE_BITS))) {
+        res = CPY_LONG_DIGIT(v, 0);
+    } else if (likely(i == CPY_SIGN_ZERO)) {
+        res = 0;
+    } else if (i == ((1 << CPY_NON_SIZE_BITS) | CPY_SIGN_NEGATIVE)) {
+        res = -(sdigit)CPY_LONG_DIGIT(v, 0);
+    } else {
+        sign = 1;
+        x = 0;
+        if (i & CPY_SIGN_NEGATIVE) {
+            sign = -1;
+        }
+        i >>= CPY_NON_SIZE_BITS;
+        while (--i >= 0) {
+            prev = x;
+            x = (x << PyLong_SHIFT) + CPY_LONG_DIGIT(v, i);
+            if ((x >> PyLong_SHIFT) != prev) {
+                *overflow = sign;
+                goto exit;
+            }
+        }
+        /* Haven't lost any bits, but casting to long requires extra
+         * care (see comment above).
+         */
+        if (x <= (size_t)CPY_TAGGED_MAX) {
+            res = (Py_ssize_t)x * sign;
+        }
+        else if (sign < 0 && x == CPY_TAGGED_ABS_MIN) {
+            res = CPY_TAGGED_MIN;
+        }
+        else {
+            *overflow = sign;
+            /* res is already set to -1 */
+        }
+    }
+  exit:
+    return res;
+}
+
+#else
+
 // Adapted from longobject.c in Python 3.7.0
 
 /* This function adapted from PyLong_AsLongLongAndOverflow, but with
@@ -156,11 +215,11 @@ CPyLong_AsSsize_tAndOverflow(PyObject *vv, int *overflow)
     i = Py_SIZE(v);
 
     if (likely(i == 1)) {
-        res = v->ob_digit[0];
+        res = CPY_LONG_DIGIT(v, 0);
     } else if (likely(i == 0)) {
         res = 0;
     } else if (i == -1) {
-        res = -(sdigit)v->ob_digit[0];
+        res = -(sdigit)CPY_LONG_DIGIT(v, 0);
     } else {
         sign = 1;
         x = 0;
@@ -170,7 +229,7 @@ CPyLong_AsSsize_tAndOverflow(PyObject *vv, int *overflow)
         }
         while (--i >= 0) {
             prev = x;
-            x = (x << PyLong_SHIFT) + v->ob_digit[i];
+            x = (x << PyLong_SHIFT) + CPY_LONG_DIGIT(v, i);
             if ((x >> PyLong_SHIFT) != prev) {
                 *overflow = sign;
                 goto exit;
@@ -194,6 +253,8 @@ CPyLong_AsSsize_tAndOverflow(PyObject *vv, int *overflow)
     return res;
 }
 
+#endif
+
 // Adapted from listobject.c in Python 3.7.0
 static int
 list_resize(PyListObject *self, Py_ssize_t newsize)