[mypyc] Add support for C string literals in the IR (#19383)

Previously only Python str and bytes literals were supported, but
sometimes we want zero-terminated C string literals instead. They
don't need to be allocated from the heap and are usually stored in a
read-only data section, so they are more efficient in some use cases.

These will be useful for a feature I'm working on.

Author: JukkaL

mypyc/codegen/emitfunc.py

@@ -33,6 +33,7 @@
     Cast,
     ComparisonOp,
     ControlOp,
+    CString,
     DecRef,
     Extend,
     Float,
@@ -850,6 +851,8 @@ def reg(self, reg: Value) -> str:
             elif r == "nan":
                 return "NAN"
             return r
+        elif isinstance(reg, CString):
+            return '"' + encode_c_string_literal(reg.value) + '"'
         else:
             return self.emitter.reg(reg)
 
@@ -911,3 +914,30 @@ def emit_unsigned_int_cast(self, type: RType) -> str:
             return "(uint64_t)"
         else:
             return ""
+
+
+_translation_table: Final[dict[int, str]] = {}
+
+
+def encode_c_string_literal(b: bytes) -> str:
+    """Convert bytestring to the C string literal syntax (with necessary escaping).
+
+    For example, b'foo\n' gets converted to 'foo\\n' (note that double quotes are not added).
+    """
+    if not _translation_table:
+        # Initialize the translation table on the first call.
+        d = {
+            ord("\n"): "\\n",
+            ord("\r"): "\\r",
+            ord("\t"): "\\t",
+            ord('"'): '\\"',
+            ord("\\"): "\\\\",
+        }
+        for i in range(256):
+            if i not in d:
+                if i < 32 or i >= 127:
+                    d[i] = "\\x%.2x" % i
+                else:
+                    d[i] = chr(i)
+        _translation_table.update(str.maketrans(d))
+    return b.decode("latin1").translate(_translation_table)

mypyc/ir/ops.py

@@ -39,6 +39,7 @@ class to enable the new behavior. Sometimes adding a new abstract
     RVoid,
     bit_rprimitive,
     bool_rprimitive,
+    cstring_rprimitive,
     float_rprimitive,
     int_rprimitive,
     is_bit_rprimitive,
@@ -230,6 +231,20 @@ def __init__(self, value: float, line: int = -1) -> None:
         self.line = line
 
 
+@final
+class CString(Value):
+    """C string literal (zero-terminated).
+
+    You can also include zero values in the value, but then you'll need to track
+    the length of the string separately.
+    """
+
+    def __init__(self, value: bytes, line: int = -1) -> None:
+        self.value = value
+        self.type = cstring_rprimitive
+        self.line = line
+
+
 class Op(Value):
     """Abstract base class for all IR operations.
 

mypyc/ir/pprint.py

@@ -21,6 +21,7 @@
     Cast,
     ComparisonOp,
     ControlOp,
+    CString,
     DecRef,
     Extend,
     Float,
@@ -327,6 +328,8 @@ def format(self, fmt: str, *args: Any) -> str:
                         result.append(str(arg.value))
                     elif isinstance(arg, Float):
                         result.append(repr(arg.value))
+                    elif isinstance(arg, CString):
+                        result.append(f"CString({arg.value!r})")
                     else:
                         result.append(self.names[arg])
                 elif typespec == "d":

mypyc/ir/rtypes.py

@@ -254,13 +254,11 @@ def __init__(
         elif ctype == "CPyPtr":
             # TODO: Invent an overlapping error value?
             self.c_undefined = "0"
-        elif ctype == "PyObject *":
-            # Boxed types use the null pointer as the error value.
+        elif ctype.endswith("*"):
+            # Boxed and pointer types use the null pointer as the error value.
             self.c_undefined = "NULL"
         elif ctype == "char":
             self.c_undefined = "2"
-        elif ctype in ("PyObject **", "void *"):
-            self.c_undefined = "NULL"
         elif ctype == "double":
             self.c_undefined = "-113.0"
         elif ctype in ("uint8_t", "uint16_t", "uint32_t", "uint64_t"):
@@ -445,6 +443,10 @@ def __hash__(self) -> int:
     "c_ptr", is_unboxed=False, is_refcounted=False, ctype="void *"
 )
 
+cstring_rprimitive: Final = RPrimitive(
+    "cstring", is_unboxed=True, is_refcounted=False, ctype="const char *"
+)
+
 # The type corresponding to mypyc.common.BITMAP_TYPE
 bitmap_rprimitive: Final = uint32_rprimitive
 

mypyc/test/test_emitfunc.py

@@ -19,6 +19,7 @@
     CallC,
     Cast,
     ComparisonOp,
+    CString,
     DecRef,
     Extend,
     GetAttr,
@@ -49,6 +50,7 @@
     RType,
     bool_rprimitive,
     c_int_rprimitive,
+    cstring_rprimitive,
     dict_rprimitive,
     int32_rprimitive,
     int64_rprimitive,
@@ -836,6 +838,30 @@ def test_inc_ref_int_literal(self) -> None:
             b = LoadLiteral(x, object_rprimitive)
             self.assert_emit([b, IncRef(b)], "CPy_INCREF(cpy_r_r0);")
 
+    def test_c_string(self) -> None:
+        s = Register(cstring_rprimitive, "s")
+        self.assert_emit(Assign(s, CString(b"foo")), """cpy_r_s = "foo";""")
+        self.assert_emit(Assign(s, CString(b'foo "o')), r"""cpy_r_s = "foo \"o";""")
+        self.assert_emit(Assign(s, CString(b"\x00")), r"""cpy_r_s = "\x00";""")
+        self.assert_emit(Assign(s, CString(b"\\")), r"""cpy_r_s = "\\";""")
+        for i in range(256):
+            b = bytes([i])
+            if b == b"\n":
+                target = "\\n"
+            elif b == b"\r":
+                target = "\\r"
+            elif b == b"\t":
+                target = "\\t"
+            elif b == b'"':
+                target = '\\"'
+            elif b == b"\\":
+                target = "\\\\"
+            elif i < 32 or i >= 127:
+                target = "\\x%.2x" % i
+            else:
+                target = b.decode("ascii")
+            self.assert_emit(Assign(s, CString(b)), f'cpy_r_s = "{target}";')
+
     def assert_emit(
         self,
         op: Op | list[Op],