gh-143050: Add helper _PyLong_InitTag() (#147956)

With this we can assume, that _PyLong_SetSignAndDigitCount() and
_PyLong_SetDigitCount() operate on non-immortal integers.

Co-authored-by: Victor Stinner <vstinner@python.org>

Author: skirpichev

Include/internal/pycore_long.h

@@ -238,11 +238,12 @@ _PyLong_IsSmallInt(const PyLongObject *op)
 {
     assert(PyLong_Check(op));
     bool is_small_int = (op->long_value.lv_tag & IMMORTALITY_BIT_MASK) != 0;
-    assert(PyLong_CheckExact(op) || (!is_small_int));
-    assert(_Py_IsImmortal(op) || (!is_small_int));
-    assert((_PyLong_IsCompact(op)
-            && _PY_IS_SMALL_INT(_PyLong_CompactValue(op)))
-           || (!is_small_int));
+    if (is_small_int) {
+        assert(PyLong_CheckExact(op));
+        assert(_Py_IsImmortal(op));
+        assert((_PyLong_IsCompact(op)
+                && _PY_IS_SMALL_INT(_PyLong_CompactValue(op))));
+    }
     return is_small_int;
 }
 
@@ -285,6 +286,14 @@ _PyLong_SameSign(const PyLongObject *a, const PyLongObject *b)
     return (a->long_value.lv_tag & SIGN_MASK) == (b->long_value.lv_tag & SIGN_MASK);
 }
 
+/* Initialize the tag of a freshly-allocated int. */
+static inline void
+_PyLong_InitTag(PyLongObject *op)
+{
+    assert(PyLong_Check(op));
+    op->long_value.lv_tag = SIGN_ZERO; /* non-immortal zero */
+}
+
 #define TAG_FROM_SIGN_AND_SIZE(sign, size) \
     ((uintptr_t)(1 - (sign)) | ((uintptr_t)(size) << NON_SIZE_BITS))
 
@@ -294,13 +303,15 @@ _PyLong_SetSignAndDigitCount(PyLongObject *op, int sign, Py_ssize_t size)
     assert(size >= 0);
     assert(-1 <= sign && sign <= 1);
     assert(sign != 0 || size == 0);
+    assert(!_PyLong_IsSmallInt(op));
     op->long_value.lv_tag = TAG_FROM_SIGN_AND_SIZE(sign, size);
 }
 
 static inline void
 _PyLong_SetDigitCount(PyLongObject *op, Py_ssize_t size)
 {
     assert(size >= 0);
+    assert(!_PyLong_IsSmallInt(op));
     op->long_value.lv_tag = (((size_t)size) << NON_SIZE_BITS) | (op->long_value.lv_tag & SIGN_MASK);
 }
 

Objects/longobject.c

@@ -185,6 +185,7 @@ long_alloc(Py_ssize_t size)
             return NULL;
         }
         _PyObject_Init((PyObject*)result, &PyLong_Type);
+        _PyLong_InitTag(result);
     }
     _PyLong_SetSignAndDigitCount(result, size != 0, size);
     /* The digit has to be initialized explicitly to avoid
@@ -258,6 +259,7 @@ _PyLong_FromMedium(sdigit x)
             return NULL;
         }
         _PyObject_Init((PyObject*)v, &PyLong_Type);
+        _PyLong_InitTag(v);
     }
     digit abs_x = x < 0 ? -x : x;
     _PyLong_SetSignAndDigitCount(v, x<0?-1:1, 1);
@@ -337,6 +339,7 @@ medium_from_stwodigits(stwodigits x)
             return PyStackRef_NULL;
         }
         _PyObject_Init((PyObject*)v, &PyLong_Type);
+        _PyLong_InitTag(v);
     }
     digit abs_x = x < 0 ? (digit)(-x) : (digit)x;
     _PyLong_SetSignAndDigitCount(v, x<0?-1:1, 1);
@@ -6011,29 +6014,34 @@ static PyObject *
 long_subtype_new(PyTypeObject *type, PyObject *x, PyObject *obase)
 {
     PyLongObject *tmp, *newobj;
-    Py_ssize_t i, n;
+    Py_ssize_t size, ndigits;
+    int sign;
 
     assert(PyType_IsSubtype(type, &PyLong_Type));
     tmp = (PyLongObject *)long_new_impl(&PyLong_Type, x, obase);
     if (tmp == NULL)
         return NULL;
     assert(PyLong_Check(tmp));
-    n = _PyLong_DigitCount(tmp);
+    size = _PyLong_DigitCount(tmp);
     /* Fast operations for single digit integers (including zero)
      * assume that there is always at least one digit present. */
-    if (n == 0) {
-        n = 1;
-    }
-    newobj = (PyLongObject *)type->tp_alloc(type, n);
+    ndigits = size ? size : 1;
+    newobj = (PyLongObject *)type->tp_alloc(type, ndigits);
     if (newobj == NULL) {
         Py_DECREF(tmp);
         return NULL;
     }
     assert(PyLong_Check(newobj));
-    newobj->long_value.lv_tag = tmp->long_value.lv_tag & ~IMMORTALITY_BIT_MASK;
-    for (i = 0; i < n; i++) {
-        newobj->long_value.ob_digit[i] = tmp->long_value.ob_digit[i];
+    if (_PyLong_IsCompact(tmp)) {
+        sign = _PyLong_CompactSign(tmp);
+    }
+    else {
+        sign = _PyLong_NonCompactSign(tmp);
     }
+    _PyLong_InitTag(newobj);
+    _PyLong_SetSignAndDigitCount(newobj, sign, size);
+    memcpy(newobj->long_value.ob_digit, tmp->long_value.ob_digit,
+           ndigits * sizeof(digit));
     Py_DECREF(tmp);
     return (PyObject *)newobj;
 }