[mypyc] Improve support for frozenset

Author: cdce8p

mypyc/codegen/emit.py

@@ -35,6 +35,7 @@
     is_dict_rprimitive,
     is_fixed_width_rtype,
     is_float_rprimitive,
+    is_frozenset_rprimitive,
     is_int16_rprimitive,
     is_int32_rprimitive,
     is_int64_rprimitive,
@@ -613,6 +614,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)
@@ -629,6 +631,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 @@
+.. _set-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

@@ -101,6 +101,7 @@
     is_dict_rprimitive,
     is_fixed_width_rtype,
     is_float_rprimitive,
+    is_frozenset_rprimitive,
     is_int16_rprimitive,
     is_int32_rprimitive,
     is_int64_rprimitive,
@@ -2225,7 +2226,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",