Repository: python/mypy
Pull request: [mypyc] Add a special case for str_rprimitive in builder.builtin_len
### Description

Closes https://github.com/mypyc/mypyc/issues/835

* Add a branch for `str_rpimitive` in `builtin_len`
* Reduce redundant code
* Faster list/tuple built from str

## Test Plan

* Add run tests for faster list/tuple built from str
* Modify one irbuild test. The original one is redundant.



Problem statement:
Add specialization for strings in builtin_len
There is no special shortcut for `str_rprimitive` in `builtin_len` currently. 

We may:
* Add a branch for `str_rpimitive`
* Reduce redundant codes

We may directly get `length` from `PyUnicodeObject` using `LoadMem` method.

```c
#define PyUnicode_GET_LENGTH(op)                \
    (assert(PyUnicode_Check(op)),               \
     assert(PyUnicode_IS_READY(op)),            \
     ((PyASCIIObject *)(op))->length)

#define PyUnicode_IS_READY(op) (((PyASCIIObject*)op)->state.ready)
```

According to CPython's comments, the `length` member would be 0 when `PyUnicode` is not ready. Maybe it's ok to omit the `IS_READY` check?

Hints from issue discussion:
> Maybe it's ok to omit the IS_READY check?

It doesn't seem possible until Python 3.12, according to docs: https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_READY

Because of the readiness checks, it may be impractical to inline any/most `str` operations, unfortunately.

Also it would be good to have tests that use not-ready string objects as well.

Commit message: Modify irbuild test and add run test for tuple creation

Patch to review:
diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py
index 2e44ec3afaf5..98878cea64c6 100644
--- a/mypyc/irbuild/for_helpers.py
+++ b/mypyc/irbuild/for_helpers.py
@@ -16,7 +16,7 @@
 )
 from mypyc.ir.rtypes import (
     RType, is_short_int_rprimitive, is_list_rprimitive, is_sequence_rprimitive,
-    is_tuple_rprimitive, is_dict_rprimitive,
+    is_tuple_rprimitive, is_dict_rprimitive, is_str_rprimitive,
     RTuple, short_int_rprimitive, int_rprimitive
 )
 from mypyc.primitives.registry import CFunctionDescription
@@ -164,7 +164,8 @@ def sequence_from_generator_preallocate_helper(
     """
     if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0:
         rtype = builder.node_type(gen.sequences[0])
-        if is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype):
+        if (is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype)
+                or is_str_rprimitive(rtype)):
             sequence = builder.accept(gen.sequences[0])
             length = builder.builder.builtin_len(sequence, gen.line, use_pyssize_t=True)
             target_op = empty_op_llbuilder(length, gen.line)
diff --git a/mypyc/irbuild/ll_builder.py b/mypyc/irbuild/ll_builder.py
index ef2e482e45c2..3ad1eed91303 100644
--- a/mypyc/irbuild/ll_builder.py
+++ b/mypyc/irbuild/ll_builder.py
@@ -51,7 +51,7 @@
     list_tuple_op, new_tuple_op, new_tuple_with_length_op
 )
 from mypyc.primitives.dict_ops import (
-    dict_update_in_display_op, dict_new_op, dict_build_op, dict_size_op
+    dict_update_in_display_op, dict_new_op, dict_build_op, dict_ssize_t_size_op
 )
 from mypyc.primitives.generic_ops import (
     py_getattr_op, py_call_op, py_call_with_kwargs_op, py_method_call_op,
@@ -63,7 +63,9 @@
 )
 from mypyc.primitives.int_ops import int_comparison_op_mapping
 from mypyc.primitives.exc_ops import err_occurred_op, keep_propagating_op
-from mypyc.primitives.str_ops import unicode_compare, str_check_if_true
+from mypyc.primitives.str_ops import (
+    unicode_compare, str_check_if_true, str_ssize_t_size_op
+)
 from mypyc.primitives.set_ops import new_set_op
 from mypyc.rt_subtype import is_runtime_subtype
 from mypyc.subtype import is_subtype
@@ -1116,32 +1118,28 @@ def builtin_len(self, val: Value, line: int, use_pyssize_t: bool = False) -> Val
         Return c_pyssize_t if use_pyssize_t is true (unshifted).
         """
         typ = val.type
+        size_value = None
         if is_list_rprimitive(typ) or is_tuple_rprimitive(typ):
             elem_address = self.add(GetElementPtr(val, PyVarObject, 'ob_size'))
             size_value = self.add(LoadMem(c_pyssize_t_rprimitive, elem_address))
             self.add(KeepAlive([val]))
-            if use_pyssize_t:
-                return size_value
-            offset = Integer(1, c_pyssize_t_rprimitive, line)
-            return self.int_op(short_int_rprimitive, size_value, offset,
-                               IntOp.LEFT_SHIFT, line)
-        elif is_dict_rprimitive(typ):
-            size_value = self.call_c(dict_size_op, [val], line)
-            if use_pyssize_t:
-                return size_value
-            offset = Integer(1, c_pyssize_t_rprimitive, line)
-            return self.int_op(short_int_rprimitive, size_value, offset,
-                               IntOp.LEFT_SHIFT, line)
         elif is_set_rprimitive(typ):
             elem_address = self.add(GetElementPtr(val, PySetObject, 'used'))
             size_value = self.add(LoadMem(c_pyssize_t_rprimitive, elem_address))
             self.add(KeepAlive([val]))
+        elif is_dict_rprimitive(typ):
+            size_value = self.call_c(dict_ssize_t_size_op, [val], line)
+        elif is_str_rprimitive(typ):
+            size_value = self.call_c(str_ssize_t_size_op, [val], line)
+
+        if size_value is not None:
             if use_pyssize_t:
                 return size_value
             offset = Integer(1, c_pyssize_t_rprimitive, line)
             return self.int_op(short_int_rprimitive, size_value, offset,
                                IntOp.LEFT_SHIFT, line)
-        elif isinstance(typ, RInstance):
+
+        if isinstance(typ, RInstance):
             # TODO: Support use_pyssize_t
             assert not use_pyssize_t
             length = self.gen_method_call(val, '__len__', [], int_rprimitive, line)
@@ -1155,12 +1153,12 @@ def builtin_len(self, val: Value, line: int, use_pyssize_t: bool = False) -> Val
             self.add(Unreachable())
             self.activate_block(ok)
             return length
+
+        # generic case
+        if use_pyssize_t:
+            return self.call_c(generic_ssize_t_len_op, [val], line)
         else:
-            # generic case
-            if use_pyssize_t:
-                return self.call_c(generic_ssize_t_len_op, [val], line)
-            else:
-                return self.call_c(generic_len_op, [val], line)
+            return self.call_c(generic_len_op, [val], line)
 
     def new_tuple(self, items: List[Value], line: int) -> Value:
         size = Integer(len(items), c_pyssize_t_rprimitive)  # type: Value
diff --git a/mypyc/irbuild/specialize.py b/mypyc/irbuild/specialize.py
index ef2986cda0b1..367a2fc33268 100644
--- a/mypyc/irbuild/specialize.py
+++ b/mypyc/irbuild/specialize.py
@@ -121,7 +121,7 @@ def dict_methods_fast_path(
 def translate_list_from_generator_call(
         builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
     # Special case for simplest list comprehension, for example
-    #     list(f(x) for x in other_list/other_tuple)
+    #     list(f(x) for x in some_list/some_tuple/some_str)
     # translate_list_comprehension() would take care of other cases if this fails.
     if (len(expr.args) == 1
             and expr.arg_kinds[0] == ARG_POS
@@ -137,7 +137,7 @@ def translate_list_from_generator_call(
 def translate_tuple_from_generator_call(
         builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
     # Special case for simplest tuple creation from a generator, for example
-    #     tuple(f(x) for x in other_list/other_tuple)
+    #     tuple(f(x) for x in some_list/some_tuple/some_str)
     # translate_safe_generator_call() would take care of other cases if this fails.
     if (len(expr.args) == 1
             and expr.arg_kinds[0] == ARG_POS
diff --git a/mypyc/primitives/dict_ops.py b/mypyc/primitives/dict_ops.py
index e176ffcac316..5ae96a174904 100644
--- a/mypyc/primitives/dict_ops.py
+++ b/mypyc/primitives/dict_ops.py
@@ -240,7 +240,7 @@
     c_function_name='CPyDict_CheckSize',
     error_kind=ERR_FALSE)
 
-dict_size_op = custom_op(
+dict_ssize_t_size_op = custom_op(
     arg_types=[dict_rprimitive],
     return_type=c_pyssize_t_rprimitive,
     c_function_name='PyDict_Size',
diff --git a/mypyc/primitives/str_ops.py b/mypyc/primitives/str_ops.py
index 4e56d885528b..a94cd480ec8b 100644
--- a/mypyc/primitives/str_ops.py
+++ b/mypyc/primitives/str_ops.py
@@ -5,7 +5,8 @@
 from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
 from mypyc.ir.rtypes import (
     RType, object_rprimitive, str_rprimitive, int_rprimitive, list_rprimitive,
-    c_int_rprimitive, pointer_rprimitive, bool_rprimitive, bit_rprimitive
+    c_int_rprimitive, pointer_rprimitive, bool_rprimitive, bit_rprimitive,
+    c_pyssize_t_rprimitive
 )
 from mypyc.primitives.registry import (
     method_op, binary_op, function_op,
@@ -134,3 +135,10 @@
     c_function_name="CPyStr_IsTrue",
     error_kind=ERR_NEVER,
 )
+
+str_ssize_t_size_op = custom_op(
+    arg_types=[str_rprimitive],
+    return_type=c_pyssize_t_rprimitive,
+    c_function_name='PyUnicode_GET_LENGTH',
+    error_kind=ERR_NEVER)
+
diff --git a/mypyc/test-data/irbuild-str.test b/mypyc/test-data/irbuild-str.test
index 3de91be40ea2..936cd4b0b4f0 100644
--- a/mypyc/test-data/irbuild-str.test
+++ b/mypyc/test-data/irbuild-str.test
@@ -158,4 +158,3 @@ L2:
     return 0
 L3:
     unreachable
-
diff --git a/mypyc/test-data/irbuild-tuple.test b/mypyc/test-data/irbuild-tuple.test
index 20590645d5d3..00fb66aa6a02 100644
--- a/mypyc/test-data/irbuild-tuple.test
+++ b/mypyc/test-data/irbuild-tuple.test
@@ -350,14 +350,13 @@ L4:
     return 1
 
 
-[case testTupleBuiltFromList2]
+[case testTupleBuiltFromStr]
 def f2(val: str) -> str:
     return val + "f2"
 
 def test() -> None:
-    source = ["a", "b", "c"]
+    source = "abc"
     a = tuple(f2(x) for x in source)
-    print(a)
 [out]
 def f2(val):
     val, r0, r1 :: str
@@ -366,71 +365,41 @@ L0:
     r1 = PyUnicode_Concat(val, r0)
     return r1
 def test():
-    r0, r1, r2 :: str
-    r3 :: list
-    r4, r5, r6, r7 :: ptr
-    source :: list
-    r8 :: ptr
-    r9 :: native_int
-    r10 :: tuple
-    r11 :: short_int
-    r12 :: ptr
-    r13 :: native_int
-    r14 :: short_int
-    r15 :: bit
-    r16 :: object
-    r17, x, r18 :: str
-    r19 :: bit
-    r20 :: short_int
+    r0, source :: str
+    r1 :: native_int
+    r2 :: tuple
+    r3 :: short_int
+    r4 :: native_int
+    r5 :: short_int
+    r6 :: bit
+    r7, x, r8 :: str
+    r9 :: bit
+    r10 :: short_int
     a :: tuple
-    r21 :: object
-    r22 :: str
-    r23, r24 :: object
 L0:
-    r0 = 'a'
-    r1 = 'b'
-    r2 = 'c'
-    r3 = PyList_New(3)
-    r4 = get_element_ptr r3 ob_item :: PyListObject
-    r5 = load_mem r4 :: ptr*
-    set_mem r5, r0 :: builtins.object*
-    r6 = r5 + WORD_SIZE*1
-    set_mem r6, r1 :: builtins.object*
-    r7 = r5 + WORD_SIZE*2
-    set_mem r7, r2 :: builtins.object*
-    keep_alive r3
-    source = r3
-    r8 = get_element_ptr source ob_size :: PyVarObject
-    r9 = load_mem r8 :: native_int*
-    keep_alive source
-    r10 = PyTuple_New(r9)
-    r11 = 0
+    r0 = 'abc'
+    source = r0
+    r1 = PyUnicode_GET_LENGTH(source)
+    r2 = PyTuple_New(r1)
+    r3 = 0
 L1:
-    r12 = get_element_ptr source ob_size :: PyVarObject
-    r13 = load_mem r12 :: native_int*
-    keep_alive source
-    r14 = r13 << 1
-    r15 = r11 < r14 :: signed
-    if r15 goto L2 else goto L4 :: bool
+    r4 = PyUnicode_GET_LENGTH(source)
+    r5 = r4 << 1
+    r6 = r3 < r5 :: signed
+    if r6 goto L2 else goto L4 :: bool
 L2:
-    r16 = CPyList_GetItemUnsafe(source, r11)
-    r17 = cast(str, r16)
-    x = r17
-    r18 = f2(x)
-    r19 = CPySequenceTuple_SetItemUnsafe(r10, r11, r18)
+    r7 = CPyStr_GetItem(source, r3)
+    x = r7
+    r8 = f2(x)
+    r9 = CPySequenceTuple_SetItemUnsafe(r2, r3, r8)
 L3:
-    r20 = r11 + 2
-    r11 = r20
+    r10 = r3 + 2
+    r3 = r10
     goto L1
 L4:
-    a = r10
-    r21 = builtins :: module
-    r22 = 'print'
-    r23 = CPyObject_GetAttr(r21, r22)
-    r24 = PyObject_CallFunctionObjArgs(r23, a, 0)
+    a = r2
     return 1
 
-
 [case testTupleBuiltFromVariableLengthTuple]
 from typing import Tuple
 
diff --git a/mypyc/test-data/run-lists.test b/mypyc/test-data/run-lists.test
index a78b2d65de4b..9ef337c4e075 100644
--- a/mypyc/test-data/run-lists.test
+++ b/mypyc/test-data/run-lists.test
@@ -357,3 +357,7 @@ def test() -> None:
     source_e = [0, 1, 2]
     e = list((x ** 2) for x in (y + 2 for y in source_e))
     assert e == [4, 9, 16]
+    source_str = "abcd"
+    f = list("str:" + x for x in source_str)
+    assert f == ["str:a", "str:b", "str:c", "str:d"]
+
diff --git a/mypyc/test-data/run-tuples.test b/mypyc/test-data/run-tuples.test
index c4abfb1d3a36..5ef243b10aa0 100644
--- a/mypyc/test-data/run-tuples.test
+++ b/mypyc/test-data/run-tuples.test
@@ -196,6 +196,10 @@ def test_sequence_generator() -> None:
     a = tuple(f8(x) for x in source_fixed_length_tuple)
     assert a == (False, True, False, True)
 
+    source_str = "abbc"
+    a = tuple("s:" + x for x in source_str)
+    assert a == ("s:a", "s:b", "s:b", "s:c")
+
 TUPLE: Final[Tuple[str, ...]] = ('x', 'y')
 
 def test_final_boxed_tuple() -> None:


Code review comments:
diff_hunk: @@ -134,3 +135,10 @@
     c_function_name="CPyStr_IsTrue",
     error_kind=ERR_NEVER,
 )
+
+str_ssize_t_size_op = custom_op(
+    arg_types=[str_rprimitive],
+    return_type=c_pyssize_t_rprimitive,
+    c_function_name='PyUnicode_GET_LENGTH',
path: mypyc/primitives/str_ops.py
text: @user1:
not a full review, but I think using `PyUnicode_GET_LENGTH` here breaks the design since it's a macro not a function.

https://github.com/python/cpython/blob/9ad8f109ac037ac088e09dcd2da523322c5caf34/Include/cpython/unicodeobject.h#L412

@author:
we use PyUnicode_GET_LENGTH in many places already and it's a public API according to https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_GET_LENGTH

@user1:
The problem is not about whether it is a public API or not, the problem is about it being a macro. As you already saw when defining the custom_op, the field is called `c_function_name`, we want the called object to be a function instead of a macro which will be textually substituted after preprocessing. Macro also heavily limits the compiler's expressiveness when we are porting to backends with no macro support.

@author:
Thanks for pointing out. Maybe I should write a wrapper function for PyUnicode_READY() and PyUnicode_GET_LENGTH()?

@user1:
I took a look at their cpython implementation. It's a bit of tricky to implement them directly using mypyc IR. So I guess wrapper is the way to go, at least for now

@author:
I am going to keep calling `PyUnicode_GET_LENGTH` directly at least for this PR, since `PyUnicode_READY` would be deprecated in 3.12.