Optimize str.encode with specializations for common used encodings

Author: svalentin

mypyc/irbuild/specialize.py

@@ -90,6 +90,11 @@
 )
 from mypyc.primitives.list_ops import new_list_set_item_op
 from mypyc.primitives.tuple_ops import new_tuple_set_item_op
+from mypyc.primitives.str_ops import (
+    str_encode_utf8_strict,
+    str_encode_ascii_strict,
+    str_encode_latin1_strict,
+)
 
 # Specializers are attempted before compiling the arguments to the
 # function.  Specializers can return None to indicate that they failed
@@ -682,6 +687,45 @@ def translate_fstring(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Va
     return None
 
 
+@specialize_function("encode", str_rprimitive)
+def str_encode_fast_path(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Specialize common cases of str.encode for most used encodings and strict errors."""
+
+    if not isinstance(callee, MemberExpr):
+        return None
+
+    # We can only specialize strict errors
+    if (
+        len(expr.arg_kinds) > 1
+        and isinstance(expr.args[1], StrExpr)
+        and expr.args[1].value != "strict"
+    ):
+        return None
+
+    if (
+        len(expr.args) > 0
+        and expr.arg_kinds[0] == ARG_NAMED
+        and expr.arg_names[0] == "errors"
+        and isinstance(expr.args[0], StrExpr)
+        and expr.args[0].value != "strict"
+    ):
+        return None
+
+    encoding = "utf8"
+    if len(expr.args) > 0 and isinstance(expr.args[0], StrExpr):
+        encoding = expr.args[0].value.lower().replace("-", "_")
+
+    # Specialized encodings and their accepted aliases
+    if encoding in ['u8', 'utf', 'utf8', 'utf_8', 'cp65001']:
+        return builder.call_c(str_encode_utf8_strict, [builder.accept(callee.expr)], expr.line)
+    elif encoding in ["ascii", "646", "us_ascii"]:
+        return builder.call_c(str_encode_ascii_strict, [builder.accept(callee.expr)], expr.line)
+    elif encoding in ['iso_8859_1', 'iso8859_1', '8859', 'cp819', 'latin', 'latin1', 'latin_1', 'l1']:
+        return builder.call_c(str_encode_latin1_strict, [builder.accept(callee.expr)], expr.line)
+
+    return None
+
+
 @specialize_function("mypy_extensions.i64")
 def translate_i64(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
     if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:

mypyc/primitives/str_ops.py

@@ -219,6 +219,30 @@
     extra_int_constants=[(0, pointer_rprimitive)],
 )
 
+# str.encode(encoding) - utf8 strict specialization
+str_encode_utf8_strict = custom_op(
+    arg_types=[str_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="PyUnicode_AsUTF8String",
+    error_kind=ERR_MAGIC,
+)
+
+# str.encode(encoding) - ascii strict specialization
+str_encode_ascii_strict = custom_op(
+    arg_types=[str_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="PyUnicode_AsASCIIString",
+    error_kind=ERR_MAGIC,
+)
+
+# str.encode(encoding) - latin1 strict specialization
+str_encode_latin1_strict = custom_op(
+    arg_types=[str_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="PyUnicode_AsLatin1String",
+    error_kind=ERR_MAGIC,
+)
+
 # str.encode(encoding, errors)
 method_op(
     name="encode",

mypyc/test-data/fixtures/ir.py

@@ -110,7 +110,7 @@ def upper(self) -> str: ...
     def startswith(self, x: str, start: int=..., end: int=...) -> bool: ...
     def endswith(self, x: str, start: int=..., end: int=...) -> bool: ...
     def replace(self, old: str, new: str, maxcount: int=...) -> str: ...
-    def encode(self, x: str=..., y: str=...) -> bytes: ...
+    def encode(self, encoding: str=..., errors: str=...) -> bytes: ...
 
 class float:
     def __init__(self, x: object) -> None: pass

mypyc/test-data/irbuild-str.test

@@ -293,20 +293,30 @@ L0:
 def f(s: str) -> None:
     s.encode()
     s.encode('utf-8')
+    s.encode('utf-8', 'strict')
+    s.encode('utf-8', errors='strict')
+    s.encode('utf-8', 'backslashreplace')
+    s.encode(encoding='ascii')
     s.encode('ascii', 'backslashreplace')
 [out]
 def f(s):
     s :: str
-    r0 :: bytes
-    r1 :: str
-    r2 :: bytes
-    r3, r4 :: str
-    r5 :: bytes
+    r0, r1, r2, r3 :: bytes
+    r4, r5 :: str
+    r6, r7 :: bytes
+    r8, r9 :: str
+    r10 :: bytes
 L0:
-    r0 = CPy_Encode(s, 0, 0)
-    r1 = 'utf-8'
-    r2 = CPy_Encode(s, r1, 0)
-    r3 = 'ascii'
-    r4 = 'backslashreplace'
-    r5 = CPy_Encode(s, r3, r4)
+    r0 = PyUnicode_AsUTF8String(s)
+    r1 = PyUnicode_AsUTF8String(s)
+    r2 = PyUnicode_AsUTF8String(s)
+    r3 = PyUnicode_AsUTF8String(s)
+    r4 = 'utf-8'
+    r5 = 'backslashreplace'
+    r6 = CPy_Encode(s, r4, r5)
+    r7 = PyUnicode_AsASCIIString(s)
+    r8 = 'ascii'
+    r9 = 'backslashreplace'
+    r10 = CPy_Encode(s, r8, r9)
     return 1
+