[mypyc] Native int primitives

mypyc/lib-rt/CPy.h

@@ -122,6 +122,7 @@ static inline size_t CPy_FindAttrOffset(PyTypeObject *trait, CPyVTableItem *vtab
 
 CPyTagged CPyTagged_FromSsize_t(Py_ssize_t value);
 CPyTagged CPyTagged_FromVoidPtr(void *ptr);
+CPyTagged CPyTagged_FromInt64(int64_t value);
 CPyTagged CPyTagged_FromObject(PyObject *object);
 CPyTagged CPyTagged_StealFromObject(PyObject *object);
 CPyTagged CPyTagged_BorrowFromObject(PyObject *object);
@@ -150,6 +151,13 @@ PyObject *CPyLong_FromStrWithBase(PyObject *o, CPyTagged base);
 PyObject *CPyLong_FromStr(PyObject *o);
 PyObject *CPyLong_FromFloat(PyObject *o);
 PyObject *CPyBool_Str(bool b);
+int64_t CPyLong_AsInt64(PyObject *o);
+int64_t CPyInt64_Divide(int64_t x, int64_t y);
+int64_t CPyInt64_Remainder(int64_t x, int64_t y);
+int32_t CPyLong_AsInt32(PyObject *o);
+int32_t CPyInt32_Divide(int32_t x, int32_t y);
+int32_t CPyInt32_Remainder(int32_t x, int32_t y);
+void CPyInt32_Overflow(void);
 
 static inline int CPyTagged_CheckLong(CPyTagged x) {
     return x & CPY_INT_TAG;
@@ -193,6 +201,12 @@ static inline bool CPyTagged_TooBig(Py_ssize_t value) {
         && (value >= 0 || value < CPY_TAGGED_MIN);
 }
 
+static inline bool CPyTagged_TooBigInt64(int64_t value) {
+    // Micro-optimized for the common case where it fits.
+    return (uint64_t)value > CPY_TAGGED_MAX
+        && (value >= 0 || value < CPY_TAGGED_MIN);
+}
+
 static inline bool CPyTagged_IsAddOverflow(CPyTagged sum, CPyTagged left, CPyTagged right) {
     // This check was copied from some of my old code I believe that it works :-)
     return (Py_ssize_t)(sum ^ left) < 0 && (Py_ssize_t)(sum ^ right) < 0;
@@ -342,8 +356,11 @@ PyObject *CPyList_GetItemUnsafe(PyObject *list, CPyTagged index);
 PyObject *CPyList_GetItemShort(PyObject *list, CPyTagged index);
 PyObject *CPyList_GetItemBorrow(PyObject *list, CPyTagged index);
 PyObject *CPyList_GetItemShortBorrow(PyObject *list, CPyTagged index);
+PyObject *CPyList_GetItemInt64(PyObject *list, int64_t index);
+PyObject *CPyList_GetItemInt64Borrow(PyObject *list, int64_t index);
 bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value);
 bool CPyList_SetItemUnsafe(PyObject *list, CPyTagged index, PyObject *value);
+bool CPyList_SetItemInt64(PyObject *list, int64_t index, PyObject *value);
 PyObject *CPyList_PopLast(PyObject *obj);
 PyObject *CPyList_Pop(PyObject *obj, CPyTagged index);
 CPyTagged CPyList_Count(PyObject *obj, PyObject *value);

mypyc/lib-rt/int_ops.c

@@ -35,6 +35,15 @@ CPyTagged CPyTagged_FromVoidPtr(void *ptr) {
     }
 }
 
+CPyTagged CPyTagged_FromInt64(int64_t value) {
+    if (unlikely(CPyTagged_TooBigInt64(value))) {
+        PyObject *object = PyLong_FromLongLong(value);
+        return ((CPyTagged)object) | CPY_INT_TAG;
+    } else {
+        return value << 1;
+    }
+}
+
 CPyTagged CPyTagged_FromObject(PyObject *object) {
     int overflow;
     // The overflow check knows about CPyTagged's width
@@ -504,3 +513,129 @@ CPyTagged CPyTagged_Lshift(CPyTagged left, CPyTagged right) {
     }
     return CPyTagged_StealFromObject(result);
 }
+
+int64_t CPyLong_AsInt64(PyObject *o) {
+    if (likely(PyLong_Check(o))) {
+        PyLongObject *lobj = (PyLongObject *)o;
+        Py_ssize_t size = Py_SIZE(lobj);
+        if (likely(size == 1)) {
+            // Fast path
+            return lobj->ob_digit[0];
+        } else if (likely(size == 0)) {
+            return 0;
+        }
+    }
+    // Slow path
+    int overflow;
+    int64_t result = PyLong_AsLongLongAndOverflow(o, &overflow);
+    if (result == -1) {
+        if (PyErr_Occurred()) {
+            return CPY_LL_INT_ERROR;
+        } else if (overflow) {
+            PyErr_SetString(PyExc_OverflowError, "int too large to convert to i64");
+            return CPY_LL_INT_ERROR;
+        }
+    }
+    return result;
+}
+
+int64_t CPyInt64_Divide(int64_t x, int64_t y) {
+    if (y == 0) {
+        PyErr_SetString(PyExc_ZeroDivisionError, "integer division or modulo by zero");
+        return CPY_LL_INT_ERROR;
+    }
+    if (y == -1 && x == -1LL << 63) {
+        PyErr_SetString(PyExc_OverflowError, "integer division overflow");
+        return CPY_LL_INT_ERROR;
+    }
+    int64_t d = x / y;
+    // Adjust for Python semantics
+    if (((x < 0) != (y < 0)) && d * y != x) {
+        d--;
+    }
+    return d;
+}
+
+int64_t CPyInt64_Remainder(int64_t x, int64_t y) {
+    if (y == 0) {
+        PyErr_SetString(PyExc_ZeroDivisionError, "integer division or modulo by zero");
+        return CPY_LL_INT_ERROR;
+    }
+    // Edge case: avoid core dump
+    if (y == -1 && x == -1LL << 63) {
+        return 0;
+    }
+    int64_t d = x % y;
+    // Adjust for Python semantics
+    if (((x < 0) != (y < 0)) && d != 0) {
+        d += y;
+    }
+    return d;
+}
+
+int32_t CPyLong_AsInt32(PyObject *o) {
+    if (likely(PyLong_Check(o))) {
+        PyLongObject *lobj = (PyLongObject *)o;
+        Py_ssize_t size = lobj->ob_base.ob_size;
+        if (likely(size == 1)) {
+            // Fast path
+            return lobj->ob_digit[0];
+        } else if (likely(size == 0)) {
+            return 0;
+        }
+    }
+    // Slow path
+    int overflow;
+    long result = PyLong_AsLongAndOverflow(o, &overflow);
+    if (result > 0x7fffffffLL || result < -0x80000000LL) {
+        overflow = 1;
+        result = -1;
+    }
+    if (result == -1) {
+        if (PyErr_Occurred()) {
+            return CPY_LL_INT_ERROR;
+        } else if (overflow) {
+            PyErr_SetString(PyExc_OverflowError, "int too large to convert to i32");
+            return CPY_LL_INT_ERROR;
+        }
+    }
+    return result;
+}
+
+int32_t CPyInt32_Divide(int32_t x, int32_t y) {
+    if (y == 0) {
+        PyErr_SetString(PyExc_ZeroDivisionError, "integer division or modulo by zero");
+        return CPY_LL_INT_ERROR;
+    }
+    if (y == -1 && x == -1LL << 31) {
+        PyErr_SetString(PyExc_OverflowError, "integer division overflow");
+        return CPY_LL_INT_ERROR;
+    }
+    int32_t d = x / y;
+    // Adjust for Python semantics
+    if (((x < 0) != (y < 0)) && d * y != x) {
+        d--;
+    }
+    return d;
+}
+
+int32_t CPyInt32_Remainder(int32_t x, int32_t y) {
+    if (y == 0) {
+        PyErr_SetString(PyExc_ZeroDivisionError, "integer division or modulo by zero");
+        return CPY_LL_INT_ERROR;
+    }
+    // Edge case: avoid core dump
+    if (y == -1 && x == -1LL << 31) {
+        return 0;
+    }
+    int32_t d = x % y;
+    // Adjust for Python semantics
+    if (((x < 0) != (y < 0)) && d != 0) {
+        d += y;
+    }
+    return d;
+}
+
+void CPyInt32_Overflow() {
+    PyErr_SetString(PyExc_OverflowError, "int too large to convert to i32");
+}

mypyc/lib-rt/list_ops.c

@@ -118,6 +118,44 @@ PyObject *CPyList_GetItemBorrow(PyObject *list, CPyTagged index) {
     }
 }
 
+PyObject *CPyList_GetItemInt64(PyObject *list, int64_t index) {
+    size_t size = PyList_GET_SIZE(list);
+    if (likely((uint64_t)index < size)) {
+        PyObject *result = PyList_GET_ITEM(list, index);
+        Py_INCREF(result);
+        return result;
+    }
+    if (index >= 0) {
+        PyErr_SetString(PyExc_IndexError, "list index out of range");
+        return NULL;
+    }
+    index += size;
+    if (index < 0) {
+        PyErr_SetString(PyExc_IndexError, "list index out of range");
+        return NULL;
+    }
+    PyObject *result = PyList_GET_ITEM(list, index);
+    Py_INCREF(result);
+    return result;
+}
+
+PyObject *CPyList_GetItemInt64Borrow(PyObject *list, int64_t index) {
+    size_t size = PyList_GET_SIZE(list);
+    if (likely((uint64_t)index < size)) {
+        return PyList_GET_ITEM(list, index);
+    }
+    if (index >= 0) {
+        PyErr_SetString(PyExc_IndexError, "list index out of range");
+        return NULL;
+    }
+    index += size;
+    if (index < 0) {
+        PyErr_SetString(PyExc_IndexError, "list index out of range");
+        return NULL;
+    }
+    return PyList_GET_ITEM(list, index);
+}
+
 bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value) {
     if (CPyTagged_CheckShort(index)) {
         Py_ssize_t n = CPyTagged_ShortAsSsize_t(index);
@@ -145,6 +183,26 @@ bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value) {
     }
 }
 
+bool CPyList_SetItemInt64(PyObject *list, int64_t index, PyObject *value) {
+    size_t size = PyList_GET_SIZE(list);
+    if (unlikely((uint64_t)index >= size)) {
+        if (index > 0) {
+            PyErr_SetString(PyExc_IndexError, "list assignment index out of range");
+            return false;
+        }
+        index += size;
+        if (index < 0) {
+            PyErr_SetString(PyExc_IndexError, "list assignment index out of range");
+            return false;
+        }
+    }
+    // PyList_SET_ITEM doesn't decref the old element, so we do
+    Py_DECREF(PyList_GET_ITEM(list, index));
+    // N.B: Steals reference
+    PyList_SET_ITEM(list, index, value);
+    return true;
+}
+
 // This function should only be used to fill in brand new lists.
 bool CPyList_SetItemUnsafe(PyObject *list, CPyTagged index, PyObject *value) {
     if (CPyTagged_CheckShort(index)) {

mypyc/lib-rt/mypyc_util.h

@@ -53,6 +53,9 @@ typedef PyObject CPyModule;
 // Tag bit used for long integers
 #define CPY_INT_TAG 1
 
+// Error value for fixed-width (low-level) integers
+#define CPY_LL_INT_ERROR -113
+
 typedef void (*CPyVTableItem)(void);
 
 static inline CPyTagged CPyTagged_ShortFromInt(int x) {

mypyc/primitives/int_ops.py

@@ -9,10 +9,11 @@
 """
 
 from typing import Dict, NamedTuple
-from mypyc.ir.ops import ERR_NEVER, ERR_MAGIC, ComparisonOp
+from mypyc.ir.ops import ERR_NEVER, ERR_MAGIC, ERR_MAGIC_OVERLAPPING, ERR_ALWAYS, ComparisonOp
 from mypyc.ir.rtypes import (
     int_rprimitive, bool_rprimitive, float_rprimitive, object_rprimitive,
-    str_rprimitive, bit_rprimitive, RType
+    str_rprimitive, bit_rprimitive, int64_rprimitive, int32_rprimitive, void_rtype, RType,
+    c_pyssize_t_rprimitive
 )
 from mypyc.primitives.registry import (
     load_address_op, unary_op, CFunctionDescription, function_op, binary_op, custom_op
@@ -165,3 +166,59 @@ def int_unary_op(name: str, c_function_name: str) -> CFunctionDescription:
     '>': IntComparisonOpDescription(ComparisonOp.SGT, int_less_than_, False, True),
     '>=': IntComparisonOpDescription(ComparisonOp.SGE, int_less_than_, True, False),
 }
+
+int64_divide_op = custom_op(
+    arg_types=[int64_rprimitive, int64_rprimitive],
+    return_type=int64_rprimitive,
+    c_function_name='CPyInt64_Divide',
+    error_kind=ERR_MAGIC_OVERLAPPING)
+
+int64_mod_op = custom_op(
+    arg_types=[int64_rprimitive, int64_rprimitive],
+    return_type=int64_rprimitive,
+    c_function_name='CPyInt64_Remainder',
+    error_kind=ERR_MAGIC_OVERLAPPING)
+
+int32_divide_op = custom_op(
+    arg_types=[int32_rprimitive, int32_rprimitive],
+    return_type=int32_rprimitive,
+    c_function_name='CPyInt32_Divide',
+    error_kind=ERR_MAGIC_OVERLAPPING)
+
+int32_mod_op = custom_op(
+    arg_types=[int32_rprimitive, int32_rprimitive],
+    return_type=int32_rprimitive,
+    c_function_name='CPyInt32_Remainder',
+    error_kind=ERR_MAGIC_OVERLAPPING)
+
+# Convert tagged int (as PyObject *) to i64
+int_to_int64_op = custom_op(
+    arg_types=[object_rprimitive],
+    return_type=int64_rprimitive,
+    c_function_name='CPyLong_AsInt64',
+    error_kind=ERR_MAGIC_OVERLAPPING)
+
+ssize_t_to_int_op = custom_op(
+    arg_types=[c_pyssize_t_rprimitive],
+    return_type=int_rprimitive,
+    c_function_name='CPyTagged_FromSsize_t',
+    error_kind=ERR_MAGIC)
+
+int64_to_int_op = custom_op(
+    arg_types=[int64_rprimitive],
+    return_type=int_rprimitive,
+    c_function_name='CPyTagged_FromInt64',
+    error_kind=ERR_MAGIC)
+
+# Convert tagged int (as PyObject *) to i32
+int_to_int32_op = custom_op(
+    arg_types=[object_rprimitive],
+    return_type=int32_rprimitive,
+    c_function_name='CPyLong_AsInt32',
+    error_kind=ERR_MAGIC_OVERLAPPING)
+
+int32_overflow = custom_op(
+    arg_types=[],
+    return_type=void_rtype,
+    c_function_name='CPyInt32_Overflow',
+    error_kind=ERR_ALWAYS)