Add inline specialized ops for ord(str[i])

python · JukkaL · Jan 13, 2026 · Jan 13, 2026 · Jan 13, 2026 · Jan 13, 2026
commit 5e0ff86b60aea278cc5c6a4c6466c6b2e53ebed6
diff --git a/mypyc/ir/deps.py b/mypyc/ir/deps.py
@@ -52,3 +52,4 @@ def get_header(self) -> str:
 BYTES_EXTRA_OPS: Final = SourceDep("bytes_extra_ops.c")
 BYTES_WRITER_EXTRA_OPS: Final = SourceDep("byteswriter_extra_ops.c")
 BYTEARRAY_EXTRA_OPS: Final = SourceDep("bytearray_extra_ops.c")
+STR_EXTRA_OPS: Final = SourceDep("str_extra_ops.c")
diff --git a/mypyc/irbuild/specialize.py b/mypyc/irbuild/specialize.py
@@ -126,9 +126,12 @@
     bytes_decode_latin1_strict,
     bytes_decode_utf8_strict,
     isinstance_str,
+    str_adjust_index_op,
     str_encode_ascii_strict,
     str_encode_latin1_strict,
     str_encode_utf8_strict,
+    str_get_item_unsafe_as_int_op,
+    str_range_check_op,
 )
 from mypyc.primitives.tuple_ops import isinstance_tuple, new_tuple_set_item_op
 
@@ -1142,8 +1145,17 @@ def translate_ord(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value
             index_type = builder.node_type(index_expr)
             if is_int_rprimitive(index_type) or is_fixed_width_rtype(index_type):
                 # This is ord(s[i]) where s is str and i is a native integer
-                # TODO: Generate specialized code here
-                assert False, "ord(s[i]) specialization not yet implemented"
+                # Generate specialized inline code using the helper
+                result = translate_getitem_with_bounds_check(
+                    builder,
+                    arg_expr.base,
+                    [arg_expr.index],
+                    expr,
+                    str_adjust_index_op,
+                    str_range_check_op,
+                    str_get_item_unsafe_as_int_op,
+                )
+                return result
 
     return None
 

diff --git a/mypyc/lib-rt/str_extra_ops.c b/mypyc/lib-rt/str_extra_ops.c
@@ -0,0 +1,4 @@
+#include "str_extra_ops.h"
+
+// All str extra ops are inline functions in str_extra_ops.h
+// This file exists to satisfy the SourceDep requirements
diff --git a/mypyc/lib-rt/str_extra_ops.h b/mypyc/lib-rt/str_extra_ops.h
@@ -0,0 +1,29 @@
+#ifndef MYPYC_STR_EXTRA_OPS_H
+#define MYPYC_STR_EXTRA_OPS_H
+
+#include <Python.h>
+#include <stdint.h>
+#include "CPy.h"
+
+// Optimized str indexing for ord(s[i])
+
+// If index is negative, convert to non-negative index (no range checking)
+static inline int64_t CPyStr_AdjustIndex(PyObject *obj, int64_t index) {
+    if (index < 0) {
+        return index + PyUnicode_GET_LENGTH(obj);
+    }
+    return index;
+}
+
+// Check if index is in valid range [0, len)
+static inline bool CPyStr_RangeCheck(PyObject *obj, int64_t index) {
+    return index >= 0 && index < PyUnicode_GET_LENGTH(obj);
+}
+
+// Get character at index as int (ord value) - no bounds checking, returns as CPyTagged
+static inline CPyTagged CPyStr_GetItemUnsafeAsInt(PyObject *obj, int64_t index) {
+    int kind = PyUnicode_KIND(obj);
+    return PyUnicode_READ(kind, PyUnicode_DATA(obj), index) << 1;
+}
+
+#endif
diff --git a/mypyc/primitives/str_ops.py b/mypyc/primitives/str_ops.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from mypyc.ir.deps import STR_EXTRA_OPS
 from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
 from mypyc.ir.rtypes import (
     RType,
@@ -10,6 +11,7 @@
     bytes_rprimitive,
     c_int_rprimitive,
     c_pyssize_t_rprimitive,
+    int64_rprimitive,
     int_rprimitive,
     list_rprimitive,
     object_rprimitive,
@@ -507,3 +509,38 @@
     c_function_name="CPyStr_Ord",
     error_kind=ERR_MAGIC,
 )
+
+# Optimized str indexing for ord(s[i])
+
+# str index adjustment - convert negative index to positive
+str_adjust_index_op = custom_primitive_op(
+    name="str_adjust_index",
+    arg_types=[str_rprimitive, int64_rprimitive],
+    return_type=int64_rprimitive,
+    c_function_name="CPyStr_AdjustIndex",
+    error_kind=ERR_NEVER,
+    experimental=True,
+    dependencies=[STR_EXTRA_OPS],
+)
+
+# str range check - check if index is in valid range
+str_range_check_op = custom_primitive_op(
+    name="str_range_check",
+    arg_types=[str_rprimitive, int64_rprimitive],
+    return_type=bool_rprimitive,
+    c_function_name="CPyStr_RangeCheck",
+    error_kind=ERR_NEVER,
+    experimental=True,
+    dependencies=[STR_EXTRA_OPS],
+)
+
+# str.__getitem__() as int - get character at index as int (ord value) - no bounds checking
+str_get_item_unsafe_as_int_op = custom_primitive_op(
+    name="str_get_item_unsafe_as_int",
+    arg_types=[str_rprimitive, int64_rprimitive],
+    return_type=int_rprimitive,
+    c_function_name="CPyStr_GetItemUnsafeAsInt",
+    error_kind=ERR_NEVER,
+    experimental=True,
+    dependencies=[STR_EXTRA_OPS],
+)
diff --git a/mypyc/test-data/irbuild-str.test b/mypyc/test-data/irbuild-str.test
@@ -515,6 +515,56 @@ def ord_str_index_i64(s: str, i: i64) -> int:
     return ord(s[i])
 [typing fixtures/typing-full.pyi]
 [out]
+def ord_str_index(s, i):
+    s :: str
+    i :: int
+    r0 :: native_int
+    r1 :: bit
+    r2, r3 :: i64
+    r4 :: ptr
+    r5 :: c_ptr
+    r6, r7 :: i64
+    r8, r9 :: bool
+    r10 :: int
+L0:
+    r0 = i & 1
+    r1 = r0 == 0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    r2 = i >> 1
+    r3 = r2
+    goto L3
+L2:
+    r4 = i ^ 1
+    r5 = r4
+    r6 = CPyLong_AsInt64(r5)
+    r3 = r6
+    keep_alive i
+L3:
+    r7 = CPyStr_AdjustIndex(s, r3)
+    r8 = CPyStr_RangeCheck(s, r7)
+    if r8 goto L5 else goto L4 :: bool
+L4:
+    r9 = raise IndexError('index out of range')
+    unreachable
+L5:
+    r10 = CPyStr_GetItemUnsafeAsInt(s, r7)
+    return r10
+def ord_str_index_i64(s, i):
+    s :: str
+    i, r0 :: i64
+    r1, r2 :: bool
+    r3 :: int
+L0:
+    r0 = CPyStr_AdjustIndex(s, i)
+    r1 = CPyStr_RangeCheck(s, r0)
+    if r1 goto L2 else goto L1 :: bool
+L1:
+    r2 = raise IndexError('index out of range')
+    unreachable
+L2:
+    r3 = CPyStr_GetItemUnsafeAsInt(s, r0)
+    return r3
 
 [case testStrip]
 from typing import NewType, Union