Micro-optimize cache primitives in librt.internal

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"]
 }