feat: specialize

Author: BobTheBuidler

mypyc/irbuild/specialize.py

@@ -94,7 +94,12 @@
     isinstance_dict,
 )
 from mypyc.primitives.float_ops import isinstance_float
-from mypyc.primitives.int_ops import isinstance_int
+from mypyc.primitives.int_ops import (
+    isinstance_int,
+    int_to_big_endian_op,
+    int_to_bytes_op,
+    int_to_little_endian_op,
+)
 from mypyc.primitives.list_ops import isinstance_list, new_list_set_item_op
 from mypyc.primitives.misc_ops import isinstance_bool
 from mypyc.primitives.set_ops import isinstance_frozenset, isinstance_set
@@ -1000,3 +1005,24 @@ def translate_ord(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value
     if isinstance(arg, (StrExpr, BytesExpr)) and len(arg.value) == 1:
         return Integer(ord(arg.value))
     return None
+
+@specialize_function("to_bytes", int_rprimitive)
+def specialize_int_to_bytes(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    # int.to_bytes(length, byteorder, signed=False)
+    # args: [self, length, byteorder, (optional) signed]
+    if len(expr.args) < 3 or len(expr.args) > 4:
+        return None
+    self_arg = builder.accept(expr.args[0])
+    length_arg = builder.accept(expr.args[1])
+    byteorder_expr = expr.args[2]
+    signed_arg = builder.accept(expr.args[3]) if len(expr.args) == 4 else builder.false()
+    if isinstance(byteorder_expr, StrExpr):
+        if byteorder_expr.value == "little":
+            return builder.call_c(int_to_little_endian_op, [self_arg, length_arg, signed_arg], expr.line)
+        elif byteorder_expr.value == "big":
+            return builder.call_c(int_to_big_endian_op, [self_arg, length_arg, signed_arg], expr.line)
+    # Fallback to generic primitive op
+    byteorder_arg = builder.accept(byteorder_expr)
+    return builder.primitive_op(
+        int_to_bytes_op, [self_arg, length_arg, byteorder_arg, signed_arg], expr.line
+    )

mypyc/lib-rt/CPy.h

@@ -149,6 +149,8 @@ CPyTagged CPyTagged_BitwiseLongOp_(CPyTagged a, CPyTagged b, char op);
 CPyTagged CPyTagged_Rshift_(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_Lshift_(CPyTagged left, CPyTagged right);
 PyObject *CPyTagged_ToBytes(CPyTagged self, Py_ssize_t length, PyObject *byteorder, int signed_flag);
+PyObject *CPyTagged_ToBigEndianBytes(CPyTagged self, Py_ssize_t length, int signed_flag);
+PyObject *CPyTagged_ToLittleEndianBytes(CPyTagged self, Py_ssize_t length, int signed_flag);
 
 PyObject *CPyTagged_Str(CPyTagged n);
 CPyTagged CPyTagged_FromFloat(double f);

mypyc/lib-rt/int_ops.c

@@ -582,17 +582,8 @@ double CPyTagged_TrueDivide(CPyTagged x, CPyTagged y) {
     return 1.0;
 }
 
-static PyObject *CPyLong_ToBytes(PyObject *v, Py_ssize_t length, const char *byteorder, int signed_flag) {
+static PyObject *CPyLong_ToBytes(PyObject *v, Py_ssize_t length, int little_endian, int signed_flag) {
     // This is a wrapper for PyLong_AsByteArray and PyBytes_FromStringAndSize
-    int little_endian;
-    if (strcmp(byteorder, "big") == 0) {
-        little_endian = 0;
-    } else if (strcmp(byteorder, "little") == 0) {
-        little_endian = 1;
-    } else {
-        PyErr_SetString(PyExc_ValueError, "byteorder must be either 'little' or 'big'");
-        return NULL;
-    }
     PyObject *result = PyBytes_FromStringAndSize(NULL, length);
     if (!result) {
         return NULL;
@@ -623,7 +614,31 @@ PyObject *CPyTagged_ToBytes(CPyTagged self, Py_ssize_t length, PyObject *byteord
         Py_DECREF(pyint);
         return NULL;
     }
-    PyObject *result = CPyLong_ToBytes(pyint, length, order, signed_flag);
+    int little_endian;
+    if (strcmp(order, "big") == 0) {
+        little_endian = 0;
+    } else if (strcmp(order, "little") == 0) {
+        little_endian = 1;
+    } else {
+        PyErr_SetString(PyExc_ValueError, "byteorder must be either 'little' or 'big'");
+        return NULL;
+    }
+    PyObject *result = CPyLong_ToBytes(pyint, length, little_endian, signed_flag);
+    Py_DECREF(pyint);
+    return result;
+
+// int.to_bytes(length, byteorder="little", signed=False)
+PyObject *CPyTagged_ToLittleEndianBytes(CPyTagged self, Py_ssize_t length, int signed_flag) {
+    PyObject *pyint = CPyTagged_AsObject(self);
+    PyObject *result = CPyLong_ToBytes(pyint, length, 1, signed_flag);
+    Py_DECREF(pyint);
+    return result;
+}
+
+// int.to_bytes(length, "big", signed=False)
+PyObject *CPyTagged_ToBigEndianBytes(CPyTagged self, Py_ssize_t length, int signed_flag) {
+    PyObject *pyint = CPyTagged_AsObject(self);
+    PyObject *result = CPyLong_ToBytes(pyint, length, 0, signed_flag);
     Py_DECREF(pyint);
     return result;
 }

mypyc/primitives/int_ops.py

@@ -324,11 +324,31 @@ def int_unary_op(name: str, c_function_name: str) -> PrimitiveDescription:
     error_kind=ERR_MAGIC,
 )
 
-# int.to_bytes(length, byteorder, signed)
-method_op(
+# int.to_bytes(length, byteorder, signed=...)
+int_to_bytes_op = method_op(
     name="to_bytes",
     arg_types=[int_rprimitive, int_rprimitive, str_rprimitive, bool_rprimitive],
     return_type=bytes_rprimitive,
     c_function_name="CPyTagged_ToBytes",
     error_kind=ERR_MAGIC,
 )
+
+# specialized custom_op cases for int.to_bytes:
+
+# int.to_bytes(length, "big")
+# int.to_bytes(length, "big", signed=...)
+int_to_big_endian_op = custom_op(
+    arg_types=[int_rprimitive, int_rprimitive, bool_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="CPyTagged_ToBigEndianBytes",
+    error_kind=ERR_MAGIC,
+)
+
+# int.to_bytes(length, "little")
+# int.to_bytes(length, "little", signed=...)
+int_to_little_endian_op = custom_op(
+    arg_types=[int_rprimitive, int_rprimitive, bool_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="CPyTagged_ToLittleEndianBytes",
+    error_kind=ERR_MAGIC,
+)