[mypyc] Improve support for frozenset (#18571)

Add optimized methods for `len` and `contains` on `frozenset` objects.

Author: cdce8p

mypyc/codegen/emit.py

@@ -34,6 +34,7 @@
     is_dict_rprimitive,
     is_fixed_width_rtype,
     is_float_rprimitive,
+    is_frozenset_rprimitive,
     is_int16_rprimitive,
     is_int32_rprimitive,
     is_int64_rprimitive,
@@ -609,6 +610,7 @@ def emit_cast(
             is_list_rprimitive(typ)
             or is_dict_rprimitive(typ)
             or is_set_rprimitive(typ)
+            or is_frozenset_rprimitive(typ)
             or is_str_rprimitive(typ)
             or is_range_rprimitive(typ)
             or is_float_rprimitive(typ)
@@ -625,6 +627,8 @@ def emit_cast(
                 prefix = "PyDict"
             elif is_set_rprimitive(typ):
                 prefix = "PySet"
+            elif is_frozenset_rprimitive(typ):
+                prefix = "PyFrozenSet"
             elif is_str_rprimitive(typ):
                 prefix = "PyUnicode"
             elif is_range_rprimitive(typ):

mypyc/doc/frozenset_operations.rst

@@ -0,0 +1,29 @@
+.. _frozenset-ops:
+
+Native frozenset operations
+======================
+
+These ``frozenset`` operations have fast, optimized implementations. Other
+frozenset operations use generic implementations that are often slower.
+
+Construction
+------------
+
+Construct empty frozenset:
+
+* ``frozenset()``
+
+Construct frozenset from iterable:
+
+* ``frozenset(x: Iterable)``
+
+
+Operators
+---------
+
+* ``item in s``
+
+Functions
+---------
+
+* ``len(s: set)``

mypyc/doc/index.rst

@@ -41,6 +41,7 @@ generate fast code.
    dict_operations
    set_operations
    tuple_operations
+   frozenset_operations
 
 .. toctree::
    :maxdepth: 2

mypyc/ir/rtypes.py

@@ -461,6 +461,11 @@ def __hash__(self) -> int:
 # Python set object (or an instance of a subclass of set).
 set_rprimitive: Final = RPrimitive("builtins.set", is_unboxed=False, is_refcounted=True)
 
+# Python frozenset object (or an instance of a subclass of frozenset).
+frozenset_rprimitive: Final = RPrimitive(
+    "builtins.frozenset", is_unboxed=False, is_refcounted=True
+)
+
 # Python str object. At the C layer, str is referred to as unicode
 # (PyUnicode).
 str_rprimitive: Final = RPrimitive("builtins.str", is_unboxed=False, is_refcounted=True)
@@ -565,6 +570,10 @@ def is_set_rprimitive(rtype: RType) -> bool:
     return isinstance(rtype, RPrimitive) and rtype.name == "builtins.set"
 
 
+def is_frozenset_rprimitive(rtype: RType) -> bool:
+    return isinstance(rtype, RPrimitive) and rtype.name == "builtins.frozenset"
+
+
 def is_str_rprimitive(rtype: RType) -> bool:
     return isinstance(rtype, RPrimitive) and rtype.name == "builtins.str"
 

mypyc/irbuild/ll_builder.py

@@ -99,6 +99,7 @@
     is_dict_rprimitive,
     is_fixed_width_rtype,
     is_float_rprimitive,
+    is_frozenset_rprimitive,
     is_int16_rprimitive,
     is_int32_rprimitive,
     is_int64_rprimitive,
@@ -2219,7 +2220,7 @@ def builtin_len(self, val: Value, line: int, use_pyssize_t: bool = False) -> Val
         size_value = None
         if is_list_rprimitive(typ) or is_tuple_rprimitive(typ) or is_bytes_rprimitive(typ):
             size_value = self.primitive_op(var_object_size, [val], line)
-        elif is_set_rprimitive(typ):
+        elif is_set_rprimitive(typ) or is_frozenset_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]))

mypyc/irbuild/mapper.py

@@ -33,6 +33,7 @@
     bytes_rprimitive,
     dict_rprimitive,
     float_rprimitive,
+    frozenset_rprimitive,
     int16_rprimitive,
     int32_rprimitive,
     int64_rprimitive,
@@ -89,6 +90,8 @@ def type_to_rtype(self, typ: Type | None) -> RType:
                 return dict_rprimitive
             elif typ.type.fullname == "builtins.set":
                 return set_rprimitive
+            elif typ.type.fullname == "builtins.frozenset":
+                return frozenset_rprimitive
             elif typ.type.fullname == "builtins.tuple":
                 return tuple_rprimitive  # Varying-length tuple
             elif typ.type.fullname == "builtins.range":

mypyc/primitives/set_ops.py

@@ -1,4 +1,4 @@
-"""Primitive set (and frozenset) ops."""
+"""Primitive set and frozenset ops."""
 
 from __future__ import annotations
 
@@ -7,6 +7,7 @@
     bit_rprimitive,
     bool_rprimitive,
     c_int_rprimitive,
+    frozenset_rprimitive,
     object_rprimitive,
     pointer_rprimitive,
     set_rprimitive,
@@ -44,11 +45,21 @@
     error_kind=ERR_MAGIC,
 )
 
+# Construct an empty frozenset
+function_op(
+    name="builtins.frozenset",
+    arg_types=[],
+    return_type=frozenset_rprimitive,
+    c_function_name="PyFrozenSet_New",
+    error_kind=ERR_MAGIC,
+    extra_int_constants=[(0, pointer_rprimitive)],
+)
+
 # frozenset(obj)
 function_op(
     name="builtins.frozenset",
     arg_types=[object_rprimitive],
-    return_type=object_rprimitive,
+    return_type=frozenset_rprimitive,
     c_function_name="PyFrozenSet_New",
     error_kind=ERR_MAGIC,
 )
@@ -64,6 +75,17 @@
     ordering=[1, 0],
 )
 
+# item in frozenset
+binary_op(
+    name="in",
+    arg_types=[object_rprimitive, frozenset_rprimitive],
+    return_type=c_int_rprimitive,
+    c_function_name="PySet_Contains",
+    error_kind=ERR_NEG_INT,
+    truncated_type=bool_rprimitive,
+    ordering=[1, 0],
+)
+
 # set.remove(obj)
 method_op(
     name="remove",

mypyc/test-data/irbuild-frozenset.test

@@ -0,0 +1,115 @@
+[case testNewFrozenSet]
+from typing import FrozenSet
+def f() -> FrozenSet[int]:
+    return frozenset({1, 2, 3})
+[out]
+def f():
+    r0 :: set
+    r1 :: object
+    r2 :: i32
+    r3 :: bit
+    r4 :: object
+    r5 :: i32
+    r6 :: bit
+    r7 :: object
+    r8 :: i32
+    r9 :: bit
+    r10 :: frozenset
+L0:
+    r0 = PySet_New(0)
+    r1 = object 1
+    r2 = PySet_Add(r0, r1)
+    r3 = r2 >= 0 :: signed
+    r4 = object 2
+    r5 = PySet_Add(r0, r4)
+    r6 = r5 >= 0 :: signed
+    r7 = object 3
+    r8 = PySet_Add(r0, r7)
+    r9 = r8 >= 0 :: signed
+    r10 = PyFrozenSet_New(r0)
+    return r10
+
+[case testNewEmptyFrozenSet]
+from typing import FrozenSet
+def f1() -> FrozenSet[int]:
+    return frozenset()
+
+def f2() -> FrozenSet[int]:
+    return frozenset(())
+[out]
+def f1():
+    r0 :: frozenset
+L0:
+    r0 = PyFrozenSet_New(0)
+    return r0
+def f2():
+    r0 :: tuple[]
+    r1 :: object
+    r2 :: frozenset
+L0:
+    r0 = ()
+    r1 = box(tuple[], r0)
+    r2 = PyFrozenSet_New(r1)
+    return r2
+
+[case testNewFrozenSetFromIterable]
+from typing import FrozenSet, List, TypeVar
+
+T = TypeVar("T")
+
+def f(l: List[T]) -> FrozenSet[T]:
+    return frozenset(l)
+[out]
+def f(l):
+    l :: list
+    r0 :: frozenset
+L0:
+    r0 = PyFrozenSet_New(l)
+    return r0
+
+[case testFrozenSetSize]
+from typing import FrozenSet
+def f() -> int:
+    return len(frozenset((1, 2, 3)))
+[out]
+def f():
+    r0 :: tuple[int, int, int]
+    r1 :: object
+    r2 :: frozenset
+    r3 :: ptr
+    r4 :: native_int
+    r5 :: short_int
+L0:
+    r0 = (2, 4, 6)
+    r1 = box(tuple[int, int, int], r0)
+    r2 = PyFrozenSet_New(r1)
+    r3 = get_element_ptr r2 used :: PySetObject
+    r4 = load_mem r3 :: native_int*
+    keep_alive r2
+    r5 = r4 << 1
+    return r5
+
+[case testFrozenSetContains]
+from typing import FrozenSet
+def f() -> bool:
+    x = frozenset((3, 4))
+    return (5 in x)
+[out]
+def f():
+    r0 :: tuple[int, int]
+    r1 :: object
+    r2, x :: frozenset
+    r3 :: object
+    r4 :: i32
+    r5 :: bit
+    r6 :: bool
+L0:
+    r0 = (6, 8)
+    r1 = box(tuple[int, int], r0)
+    r2 = PyFrozenSet_New(r1)
+    x = r2
+    r3 = object 5
+    r4 = PySet_Contains(x, r3)
+    r5 = r4 >= 0 :: signed
+    r6 = truncate r4: i32 to builtins.bool
+    return r6

mypyc/test-data/irbuild-set.test

@@ -686,7 +686,7 @@ def not_precomputed_nested_set(i):
     r1 :: object
     r2 :: i32
     r3 :: bit
-    r4 :: object
+    r4 :: frozenset
     r5 :: set
     r6 :: i32
     r7 :: bit