[mypyc] Add primitive for int to ascii conversion (#21036)

Added a primitive to convert `int`s into `bytes` objects suitable for
bytes format strings, eg. `b'%d' % my_int`. Works the same as the
existing `CPyTagged_Str` primitive used for formatting `str` objects,
except it returns a `bytes` object instead of `str`.

Currently only the bare `%d` specifier is supported, other integer
specifiers and `%d` with options go through the regular CPython handling
of format strings.

Results in around 23% reduced runtime in this microbenchmark:
```python
import time

t0 = time.time()
for i in range(10_000_000):
    val = b'%d' % i
print(time.time() - t0)
```

Results:
| setup | runtime (s) |
| - | - |
| interpreted | 0.90 |
| mypyc master | 0.92 |
| mypyc PR | 0.69 |

Author: p-sawicki

mypyc/irbuild/format_str_tokenizer.py

@@ -25,7 +25,7 @@
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.constant_fold import constant_fold_expr
 from mypyc.primitives.bytes_ops import bytes_build_op
-from mypyc.primitives.int_ops import int_to_str_op
+from mypyc.primitives.int_ops import int_to_ascii_op, int_to_str_op
 from mypyc.primitives.str_ops import str_build_op, str_op
 
 
@@ -225,8 +225,13 @@ def convert_format_expr_to_bytes(
                 var_bytes = builder.accept(x)
             else:
                 return None
-        else:
-            return None
+        elif format_op == FormatOp.INT:
+            if isinstance(folded := constant_fold_expr(builder, x), int):
+                var_bytes = builder.load_literal_value(str(folded).encode("ascii"))
+            elif is_int_rprimitive(node_type) or is_short_int_rprimitive(node_type):
+                var_bytes = builder.call_c(int_to_ascii_op, [builder.accept(x)], line)
+            else:
+                return None
         converted.append(var_bytes)
     return converted
 

mypyc/lib-rt/CPy.h

@@ -154,6 +154,7 @@ PyObject *CPyTagged_ToBigEndianBytes(CPyTagged self, Py_ssize_t length, int sign
 PyObject *CPyTagged_ToLittleEndianBytes(CPyTagged self, Py_ssize_t length, int signed_flag);
 
 PyObject *CPyTagged_Str(CPyTagged n);
+PyObject *CPyTagged_AsciiBytes(CPyTagged n);
 CPyTagged CPyTagged_FromFloat(double f);
 PyObject *CPyLong_FromStrWithBase(PyObject *o, CPyTagged base);
 PyObject *CPyLong_FromStr(PyObject *o);

mypyc/lib-rt/misc_ops.c

@@ -564,6 +564,24 @@ PyObject *CPyTagged_Str(CPyTagged n) {
     }
 }
 
+static PyObject *CPyTagged_ShortToAsciiBytes(Py_ssize_t n) {
+    PyObject *obj = PyBytes_FromStringAndSize(NULL, MAX_INT_CHARS);
+    if (!obj) return NULL;
+    int len = fmt_ssize_t(PyBytes_AsString(obj), n);
+    Py_SET_SIZE(obj, len);
+    return obj;
+}
+
+PyObject *CPyTagged_AsciiBytes(CPyTagged n) {
+    if (CPyTagged_CheckShort(n)) {
+        return CPyTagged_ShortToAsciiBytes(CPyTagged_ShortAsSsize_t(n));
+    }
+    PyObject *str = PyObject_Str(CPyTagged_AsObject(n));
+    PyObject *bytes = PyUnicode_AsASCIIString(str);
+    CPy_DECREF(str);
+    return bytes;
+}
+
 void CPyDebug_Print(const char *msg) {
     printf("%s\n", msg);
     fflush(stdout);

mypyc/primitives/int_ops.py

@@ -105,6 +105,13 @@
         priority=3,
     )
 
+int_to_ascii_op = custom_op(
+    arg_types=[int_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="CPyTagged_AsciiBytes",
+    error_kind=ERR_MAGIC,
+)
+
 
 def int_binary_primitive(
     op: str, primitive_name: str, return_type: RType = int_rprimitive, error_kind: int = ERR_NEVER

mypyc/test-data/irbuild-bytes.test

@@ -163,31 +163,39 @@ def f(var: bytes, num: int) -> None:
     b2 = b'aaaa%bbbbb%s%d' % (var, var, num)
     b3 = b'%b' % var
     b4 = b'%ssss' % var
+    b5 = b'%d' % num
+    b6 = b'%d' % 42
+    b7 = b'%d' % (67 + 2)
 [typing fixtures/typing-full.pyi]
 [out]
 def f(var, num):
     var :: bytes
     num :: int
-    r0, r1, r2, b1, r3 :: bytes
-    r4 :: tuple[bytes, bytes, int]
-    r5, r6 :: object
-    r7, b2, r8, b3, r9, r10, b4 :: bytes
+    r0, r1, r2, b1, r3, r4, r5, r6, b2, r7, b3, r8, r9, b4, r10, r11, b5, r12, r13, b6, r14, r15, b7 :: bytes
 L0:
     r0 = b'aaaa'
     r1 = b'bbbb'
     r2 = CPyBytes_Build(4, r0, var, r1, var)
     b1 = r2
-    r3 = b'aaaa%bbbbb%s%d'
-    r4 = (var, var, num)
-    r5 = box(tuple[bytes, bytes, int], r4)
-    r6 = PyNumber_Remainder(r3, r5)
-    r7 = cast(bytes, r6)
-    b2 = r7
-    r8 = CPyBytes_Build(1, var)
-    b3 = r8
-    r9 = b'sss'
-    r10 = CPyBytes_Build(2, var, r9)
-    b4 = r10
+    r3 = CPyTagged_AsciiBytes(num)
+    r4 = b'aaaa'
+    r5 = b'bbbb'
+    r6 = CPyBytes_Build(5, r4, var, r5, var, r3)
+    b2 = r6
+    r7 = CPyBytes_Build(1, var)
+    b3 = r7
+    r8 = b'sss'
+    r9 = CPyBytes_Build(2, var, r8)
+    b4 = r9
+    r10 = CPyTagged_AsciiBytes(num)
+    r11 = CPyBytes_Build(1, r10)
+    b5 = r11
+    r12 = b'42'
+    r13 = CPyBytes_Build(1, r12)
+    b6 = r13
+    r14 = b'69'
+    r15 = CPyBytes_Build(1, r14)
+    b7 = r15
     return 1
 
 [case testOptionalBytesEquality]

mypyc/test-data/run-bytes.test

@@ -412,6 +412,15 @@ def test_bytes_formatting() -> None:
     assert b'%4x' % val == b'   a'
     assert b'%#4x' % val == b' 0xa'
     assert b'%04X' % val == b'000A'
+    assert b'%d' % val == b'10'
+    assert b'%d' % (val + 2) == b'12'
+    assert b'%d' % 0 == b'0'
+    assert b'%d' % -3 == b'-3'
+    assert b'folded: %d' % (4 - 2) == b'folded: 2'
+
+    large_num = 2**65
+    assert b'number: %d' % large_num == b'number: 36893488147419103232'
+    assert b'negative integer: %d' % (-large_num) == b'negative integer: -36893488147419103232'
 
     assert b'%c' % 48 == b'0'
     assert b'%c' % b'a' == b'a'