[mypyc] Implement new-style builtins.len for all supported types (#9284)

This PR completes the support of the new style builtins.len for dict, 
set, tuple and generic cases.

Author: TH3CHARLie

mypyc/ir/rtypes.py

@@ -681,3 +681,21 @@ def is_optional_type(rtype: RType) -> bool:
     name='PyVarObject',
     names=['ob_base', 'ob_size'],
     types=[PyObject, c_pyssize_t_rprimitive])
+
+setentry = RStruct(
+    name='setentry',
+    names=['key', 'hash'],
+    types=[pointer_rprimitive, c_pyssize_t_rprimitive])
+
+smalltable = RStruct(
+    name='smalltable',
+    names=[],
+    types=[setentry] * 8)
+
+PySetObject = RStruct(
+    name='PySetObject',
+    names=['ob_base', 'fill', 'used', 'mask', 'table', 'hash', 'finger',
+           'smalltable', 'weakreflist'],
+    types=[PyObject, c_pyssize_t_rprimitive, c_pyssize_t_rprimitive, c_pyssize_t_rprimitive,
+           pointer_rprimitive, c_pyssize_t_rprimitive, c_pyssize_t_rprimitive, smalltable,
+           pointer_rprimitive])

mypyc/irbuild/builder.py

@@ -238,8 +238,8 @@ def binary_int_op(self, type: RType, lhs: Value, rhs: Value, op: int, line: int)
     def compare_tagged(self, lhs: Value, rhs: Value, op: str, line: int) -> Value:
         return self.builder.compare_tagged(lhs, rhs, op, line)
 
-    def list_len(self, val: Value, line: int) -> Value:
-        return self.builder.list_len(val, line)
+    def builtin_len(self, val: Value, line: int) -> Value:
+        return self.builder.builtin_len(val, line)
 
     @property
     def environment(self) -> Environment:
@@ -511,7 +511,7 @@ def process_iterator_tuple_assignment(self,
         if target.star_idx is not None:
             post_star_vals = target.items[split_idx + 1:]
             iter_list = self.call_c(to_list, [iterator], line)
-            iter_list_len = self.list_len(iter_list, line)
+            iter_list_len = self.builtin_len(iter_list, line)
             post_star_len = self.add(LoadInt(len(post_star_vals)))
             condition = self.binary_op(post_star_len, iter_list_len, '<=', line)
 

mypyc/irbuild/for_helpers.py

@@ -332,14 +332,7 @@ def gen_cleanup(self) -> None:
 
     def load_len(self, expr: Union[Value, AssignmentTarget]) -> Value:
         """A helper to get collection length, used by several subclasses."""
-        val = self.builder.read(expr, self.line)
-        if is_list_rprimitive(val.type):
-            return self.builder.builder.list_len(self.builder.read(expr, self.line), self.line)
-        return self.builder.builder.builtin_call(
-            [self.builder.read(expr, self.line)],
-            'builtins.len',
-            self.line,
-        )
+        return self.builder.builder.builtin_len(self.builder.read(expr, self.line), self.line)
 
 
 class ForIterable(ForGenerator):

mypyc/irbuild/ll_builder.py

@@ -27,7 +27,8 @@
 from mypyc.ir.rtypes import (
     RType, RUnion, RInstance, optional_value_type, int_rprimitive, float_rprimitive,
     bool_rprimitive, list_rprimitive, str_rprimitive, is_none_rprimitive, object_rprimitive,
-    c_pyssize_t_rprimitive, is_short_int_rprimitive, is_tagged, PyVarObject, short_int_rprimitive
+    c_pyssize_t_rprimitive, is_short_int_rprimitive, is_tagged, PyVarObject, short_int_rprimitive,
+    is_list_rprimitive, is_tuple_rprimitive, is_dict_rprimitive, is_set_rprimitive, PySetObject
 )
 from mypyc.ir.func_ir import FuncDecl, FuncSignature
 from mypyc.ir.class_ir import ClassIR, all_concrete_classes
@@ -45,10 +46,10 @@
 )
 from mypyc.primitives.tuple_ops import list_tuple_op, new_tuple_op
 from mypyc.primitives.dict_ops import (
-    dict_update_in_display_op, dict_new_op, dict_build_op
+    dict_update_in_display_op, dict_new_op, dict_build_op, dict_size_op
 )
 from mypyc.primitives.generic_ops import (
-    py_getattr_op, py_call_op, py_call_with_kwargs_op, py_method_call_op
+    py_getattr_op, py_call_op, py_call_with_kwargs_op, py_method_call_op, generic_len_op
 )
 from mypyc.primitives.misc_ops import (
     none_op, none_object_op, false_op, fast_isinstance_op, bool_op, type_is_op
@@ -704,7 +705,7 @@ def add_bool_branch(self, value: Value, true: BasicBlock, false: BasicBlock) ->
             zero = self.add(LoadInt(0))
             value = self.binary_op(value, zero, '!=', value.line)
         elif is_same_type(value.type, list_rprimitive):
-            length = self.list_len(value, value.line)
+            length = self.builtin_len(value, value.line)
             zero = self.add(LoadInt(0))
             value = self.binary_op(length, zero, '!=', value.line)
         elif (isinstance(value.type, RInstance) and value.type.class_ir.is_ext_class
@@ -811,12 +812,29 @@ def matching_call_c(self,
     def binary_int_op(self, type: RType, lhs: Value, rhs: Value, op: int, line: int) -> Value:
         return self.add(BinaryIntOp(type, lhs, rhs, op, line))
 
-    def list_len(self, val: Value, line: int) -> Value:
-        elem_address = self.add(GetElementPtr(val, PyVarObject, 'ob_size'))
-        size_value = self.add(LoadMem(c_pyssize_t_rprimitive, elem_address))
-        offset = self.add(LoadInt(1, -1, rtype=c_pyssize_t_rprimitive))
-        return self.binary_int_op(short_int_rprimitive, size_value, offset,
-                                  BinaryIntOp.LEFT_SHIFT, -1)
+    def builtin_len(self, val: Value, line: int) -> Value:
+        typ = val.type
+        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))
+            offset = self.add(LoadInt(1, line, rtype=c_pyssize_t_rprimitive))
+            return self.binary_int_op(short_int_rprimitive, size_value, offset,
+                                      BinaryIntOp.LEFT_SHIFT, line)
+        elif is_dict_rprimitive(typ):
+            size_value = self.call_c(dict_size_op, [val], line)
+            offset = self.add(LoadInt(1, line, rtype=c_pyssize_t_rprimitive))
+            return self.binary_int_op(short_int_rprimitive, size_value, offset,
+                                      BinaryIntOp.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))
+            offset = self.add(LoadInt(1, line, rtype=c_pyssize_t_rprimitive))
+            return self.binary_int_op(short_int_rprimitive, size_value, offset,
+                                      BinaryIntOp.LEFT_SHIFT, line)
+        # generic case
+        else:
+            return self.call_c(generic_len_op, [val], line)
+
     # Internal helpers
 
     def decompose_union_helper(self,

mypyc/irbuild/specialize.py

@@ -22,7 +22,7 @@
 )
 from mypyc.ir.rtypes import (
     RType, RTuple, str_rprimitive, list_rprimitive, dict_rprimitive, set_rprimitive,
-    bool_rprimitive, is_dict_rprimitive, is_list_rprimitive,
+    bool_rprimitive, is_dict_rprimitive
 )
 from mypyc.primitives.dict_ops import dict_keys_op, dict_values_op, dict_items_op
 from mypyc.primitives.misc_ops import true_op, false_op
@@ -75,9 +75,9 @@ def translate_len(
             # though we still need to evaluate it.
             builder.accept(expr.args[0])
             return builder.add(LoadInt(len(expr_rtype.types)))
-        elif is_list_rprimitive(expr_rtype):
+        else:
             obj = builder.accept(expr.args[0])
-            return builder.list_len(obj, -1)
+            return builder.builtin_len(obj, -1)
     return None
 
 

mypyc/primitives/dict_ops.py

@@ -1,16 +1,14 @@
 """Primitive dict ops."""
 
-from typing import List
-
-from mypyc.ir.ops import EmitterInterface, ERR_FALSE, ERR_MAGIC, ERR_NEVER, ERR_NEG_INT
+from mypyc.ir.ops import ERR_FALSE, ERR_MAGIC, ERR_NEVER, ERR_NEG_INT
 from mypyc.ir.rtypes import (
     dict_rprimitive, object_rprimitive, bool_rprimitive, int_rprimitive,
     list_rprimitive, dict_next_rtuple_single, dict_next_rtuple_pair, c_pyssize_t_rprimitive,
     c_int_rprimitive
 )
 
 from mypyc.primitives.registry import (
-    name_ref_op, method_op, func_op,
+    name_ref_op, method_op,
     simple_emit, name_emit, c_custom_op, c_method_op, c_function_op, c_binary_op
 )
 
@@ -168,21 +166,6 @@
     c_function_name='CPyDict_Items',
     error_kind=ERR_MAGIC)
 
-
-def emit_len(emitter: EmitterInterface, args: List[str], dest: str) -> None:
-    temp = emitter.temp_name()
-    emitter.emit_declaration('Py_ssize_t %s;' % temp)
-    emitter.emit_line('%s = PyDict_Size(%s);' % (temp, args[0]))
-    emitter.emit_line('%s = CPyTagged_ShortFromSsize_t(%s);' % (dest, temp))
-
-
-# len(dict)
-func_op(name='builtins.len',
-        arg_types=[dict_rprimitive],
-        result_type=int_rprimitive,
-        error_kind=ERR_NEVER,
-        emit=emit_len)
-
 # PyDict_Next() fast iteration
 dict_key_iter_op = c_custom_op(
     arg_types=[dict_rprimitive],
@@ -226,3 +209,9 @@ def emit_len(emitter: EmitterInterface, args: List[str], dest: str) -> None:
     return_type=bool_rprimitive,
     c_function_name='CPyDict_CheckSize',
     error_kind=ERR_FALSE)
+
+dict_size_op = c_custom_op(
+    arg_types=[dict_rprimitive],
+    return_type=c_pyssize_t_rprimitive,
+    c_function_name='PyDict_Size',
+    error_kind=ERR_NEVER)

mypyc/primitives/generic_ops.py

@@ -12,7 +12,7 @@
 from mypyc.ir.ops import ERR_NEVER, ERR_MAGIC, ERR_NEG_INT
 from mypyc.ir.rtypes import object_rprimitive, int_rprimitive, bool_rprimitive, c_int_rprimitive
 from mypyc.primitives.registry import (
-    binary_op, unary_op, func_op, custom_op, call_emit, simple_emit,
+    binary_op, unary_op, custom_op, call_emit, simple_emit,
     call_negative_magic_emit, negative_int_emit,
     c_binary_op, c_unary_op, c_method_op, c_function_op, c_custom_op
 )
@@ -226,12 +226,11 @@
     emit=simple_emit('{dest} = CPyObject_CallMethodObjArgs({comma_args}, NULL);'))
 
 # len(obj)
-func_op(name='builtins.len',
-        arg_types=[object_rprimitive],
-        result_type=int_rprimitive,
-        error_kind=ERR_NEVER,
-        emit=call_emit('CPyObject_Size'),
-        priority=0)
+generic_len_op = c_custom_op(
+    arg_types=[object_rprimitive],
+    return_type=int_rprimitive,
+    c_function_name='CPyObject_Size',
+    error_kind=ERR_NEVER)
 
 # iter(obj)
 iter_op = c_function_op(name='builtins.iter',

mypyc/primitives/set_ops.py

@@ -3,11 +3,10 @@
 from mypyc.primitives.registry import (
     func_op, simple_emit, c_function_op, c_method_op, c_binary_op
 )
-from mypyc.ir.ops import ERR_MAGIC, ERR_FALSE, ERR_NEVER, ERR_NEG_INT, EmitterInterface
+from mypyc.ir.ops import ERR_MAGIC, ERR_FALSE, ERR_NEG_INT
 from mypyc.ir.rtypes import (
-    object_rprimitive, bool_rprimitive, set_rprimitive, int_rprimitive, c_int_rprimitive
+    object_rprimitive, bool_rprimitive, set_rprimitive, c_int_rprimitive
 )
-from typing import List
 
 
 # Construct an empty set.
@@ -35,23 +34,6 @@
     c_function_name='PyFrozenSet_New',
     error_kind=ERR_MAGIC)
 
-
-def emit_len(emitter: EmitterInterface, args: List[str], dest: str) -> None:
-    temp = emitter.temp_name()
-    emitter.emit_declaration('Py_ssize_t %s;' % temp)
-    emitter.emit_line('%s = PySet_GET_SIZE(%s);' % (temp, args[0]))
-    emitter.emit_line('%s = CPyTagged_ShortFromSsize_t(%s);' % (dest, temp))
-
-
-# len(set)
-func_op(
-    name='builtins.len',
-    arg_types=[set_rprimitive],
-    result_type=int_rprimitive,
-    error_kind=ERR_NEVER,
-    emit=emit_len,
-)
-
 # item in set
 c_binary_op(
     name='in',

mypyc/primitives/tuple_ops.py

@@ -4,14 +4,10 @@
 objects, i.e. tuple_rprimitive (RPrimitive), not RTuple.
 """
 
-from typing import List
-
-from mypyc.ir.ops import (
-    EmitterInterface, ERR_NEVER, ERR_MAGIC
-)
+from mypyc.ir.ops import ERR_MAGIC
 from mypyc.ir.rtypes import tuple_rprimitive, int_rprimitive, list_rprimitive, object_rprimitive
 from mypyc.primitives.registry import (
-    func_op, c_method_op, custom_op, simple_emit, c_function_op
+    c_method_op, custom_op, simple_emit, c_function_op
 )
 
 
@@ -33,22 +29,6 @@
     format_str='{dest} = ({comma_args}) :: tuple',
     emit=simple_emit('{dest} = PyTuple_Pack({num_args}{comma_if_args}{comma_args});'))
 
-
-def emit_len(emitter: EmitterInterface, args: List[str], dest: str) -> None:
-    temp = emitter.temp_name()
-    emitter.emit_declaration('Py_ssize_t %s;' % temp)
-    emitter.emit_line('%s = PyTuple_GET_SIZE(%s);' % (temp, args[0]))
-    emitter.emit_line('%s = CPyTagged_ShortFromSsize_t(%s);' % (dest, temp))
-
-
-# len(tuple)
-tuple_len_op = func_op(
-    name='builtins.len',
-    arg_types=[tuple_rprimitive],
-    result_type=int_rprimitive,
-    error_kind=ERR_NEVER,
-    emit=emit_len)
-
 # Construct tuple from a list.
 list_tuple_op = c_function_op(
     name='builtins.tuple',