Micro-optimize cache primitives in librt.internal (#20194)

This includes the following optimizations:
* Split Buffer into ReadBuffer and WriteBuffer which are more
specialized
* Only check buffer types in wrapper functions, not in C primitives
* Use pointers instead of integer indexes in the buffer objects

This improves the performance of a micro-benchmark that reads integers
in a loop by 5%, but this could help more if we'd inline some of the
smaller functions (in the future). By making the functions simpler,
inlining is more feasible.

I'm not sure what's the best way to merge this -- maybe we'll need to
have broken master for a while, and then we ca publish a new version of
`librt`, and then update mypy to work using the new `librt` version.

Author: JukkaL

mypy/typeshed/stubs/librt/librt/internal.pyi

@@ -1,19 +1,27 @@
 from mypy_extensions import u8
 
+# TODO: Remove Buffer -- right now we have hacky support for BOTH the old and new APIs
+
 class Buffer:
     def __init__(self, source: bytes = ...) -> None: ...
     def getvalue(self) -> bytes: ...
 
-def write_bool(data: Buffer, value: bool) -> None: ...
-def read_bool(data: Buffer) -> bool: ...
-def write_str(data: Buffer, value: str) -> None: ...
-def read_str(data: Buffer) -> str: ...
-def write_bytes(data: Buffer, value: bytes) -> None: ...
-def read_bytes(data: Buffer) -> bytes: ...
-def write_float(data: Buffer, value: float) -> None: ...
-def read_float(data: Buffer) -> float: ...
-def write_int(data: Buffer, value: int) -> None: ...
-def read_int(data: Buffer) -> int: ...
-def write_tag(data: Buffer, value: u8) -> None: ...
-def read_tag(data: Buffer) -> u8: ...
+class ReadBuffer:
+    def __init__(self, source: bytes) -> None: ...
+
+class WriteBuffer:
+    def getvalue(self) -> bytes: ...
+
+def write_bool(data: WriteBuffer | Buffer, value: bool) -> None: ...
+def read_bool(data: ReadBuffer | Buffer) -> bool: ...
+def write_str(data: WriteBuffer | Buffer, value: str) -> None: ...
+def read_str(data: ReadBuffer | Buffer) -> str: ...
+def write_bytes(data: WriteBuffer | Buffer, value: bytes) -> None: ...
+def read_bytes(data: ReadBuffer | Buffer) -> bytes: ...
+def write_float(data: WriteBuffer | Buffer, value: float) -> None: ...
+def read_float(data: ReadBuffer | Buffer) -> float: ...
+def write_int(data: WriteBuffer | Buffer, value: int) -> None: ...
+def read_int(data: ReadBuffer | Buffer) -> int: ...
+def write_tag(data: WriteBuffer | Buffer, value: u8) -> None: ...
+def read_tag(data: ReadBuffer | Buffer) -> u8: ...
 def cache_version() -> u8: ...

mypyc/codegen/emit.py

@@ -705,13 +705,25 @@ def emit_cast(
             self.emit_lines(f"    {dest} = {src};", "else {")
             self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
             self.emit_line("}")
-        elif is_object_rprimitive(typ) or is_native_rprimitive(typ):
+        elif is_object_rprimitive(typ):
             if declare_dest:
                 self.emit_line(f"PyObject *{dest};")
             self.emit_arg_check(src, dest, typ, "", optional)
             self.emit_line(f"{dest} = {src};")
             if optional:
                 self.emit_line("}")
+        elif is_native_rprimitive(typ):
+            # Native primitive types have type check functions of form "CPy<Name>_Check(...)".
+            if declare_dest:
+                self.emit_line(f"PyObject *{dest};")
+            short_name = typ.name.rsplit(".", 1)[-1]
+            check = f"(CPy{short_name}_Check({src}))"
+            if likely:
+                check = f"(likely{check})"
+            self.emit_arg_check(src, dest, typ, check, optional)
+            self.emit_lines(f"    {dest} = {src};", "else {")
+            self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
+            self.emit_line("}")
         elif isinstance(typ, RUnion):
             self.emit_union_cast(
                 src, dest, typ, declare_dest, error, optional, src_type, raise_exception

mypyc/ir/rtypes.py

@@ -514,7 +514,7 @@ def __hash__(self) -> int:
 
 KNOWN_NATIVE_TYPES: Final = {
     name: RPrimitive(name, is_unboxed=False, is_refcounted=True)
-    for name in ["librt.internal.Buffer"]
+    for name in ["librt.internal.WriteBuffer", "librt.internal.ReadBuffer"]
 }