Jan 10, 2023 · Aug 12, 2021 · Aug 24, 2022 · Jan 6, 2023 · Jan 6, 2023 · Jan 6, 2023
diff --git a/mypyc/analysis/ircheck.py b/mypyc/analysis/ircheck.py
            if isinstance(x, tuple):
                self.check_tuple_items_valid_literals(op, x)

    def check_frozenset_items_valid_literals(
        self, op: LoadLiteral, s: frozenset[object, ...]
    ) -> None:
        for x in s:
            if x is not None and not isinstance(x, (str, bytes, bool, int, float, complex, tuple)):
                self.fail(op, f"Invalid type for item of frozenset literal: {type(x)})")
            if isinstance(x, tuple):
                self.check_tuple_items_valid_literals(op, x)

    def visit_load_literal(self, op: LoadLiteral) -> None:
        expected_type = None
        if op.value is None:
        elif isinstance(op.value, tuple):
            expected_type = "builtins.tuple"
            self.check_tuple_items_valid_literals(op, op.value)
        elif isinstance(op.value, frozenset):
            # There's no frozenset_rprimitive type since it'd be pretty useless so we just pretend
            # it's a set (when it's really a frozenset).
            expected_type = "builtins.set"
            self.check_frozenset_items_valid_literals(op, op.value)

        assert expected_type is not None, "Missed a case for LoadLiteral check"

diff --git a/mypyc/codegen/emitmodule.py b/mypyc/codegen/emitmodule.py
        self.declare_global("const int []", "CPyLit_Tuple", initializer=init_tuple)
        # Descriptions of frozenset literals
        init_frozenset = c_array_initializer(literals.encoded_frozenset_values())
        self.declare_global('const int []', 'CPyLit_FrozenSet', initializer=init_frozenset)
        self.declare_global("const int []", "CPyLit_FrozenSet", initializer=init_frozenset)

    def generate_export_table(self, decl_emitter: Emitter, code_emitter: Emitter) -> None:
        """Generate the declaration and definition of the group's export struct.
diff --git a/mypyc/codegen/literals.py b/mypyc/codegen/literals.py

 # Supported Python literal types. All tuple / frozenset items must have supported
 # literal types as well, but we can't represent the type precisely.
 LiteralValue = Union[str, bytes, int, bool, float, complex, Tuple[object, ...], FrozenSet[object], None]
 LiteralValue = Union[
    str, bytes, int, bool, float, complex, Tuple[object, ...], FrozenSet[object], None
 ]

 # Some literals are singletons and handled specially (None, False and True)
 NUM_SINGLETONS: Final = 3
        return _encode_complex_values(self.complex_literals)

    def encoded_tuple_values(self) -> list[str]:
        """Encode tuple values into a C array.
        return self._encode_collection_values(self.tuple_literals)

    def encoded_frozenset_values(self) -> List[str]:
        return self._encode_collection_values(self.frozenset_literals)

    def _encode_collection_values(
        self, values: dict[tuple[object, ...], int] | dict[frozenset[object], int]
    ) -> list[str]:
        """Encode tuple/frozen values into a C array.

        The format of the result is like this:

           <number oftuples>
           <length of the firsttuple>
           <number ofcollections>
           <length of the firstcollection>
           <literal index of first item>
           ...
           <literal index of last item>
           <length of the secondtuple>
           <length of the secondcollection>
           ...
        """
        values = self.tuple_literals
        value_by_index = {index: value for value, index in values.items()}
        result = []
        num = len(values)
        result.append(str(num))
        for i in range(num):
            value = value_by_index[i]
            result.append(str(len(value)))
            for item in value:
                index = self.literal_index(cast(Any, item))
                result.append(str(index))
        return result

    def encoded_frozenset_values(self) -> List[str]:
        """Encode frozenset values into a C array.

        The format used here is identical to one used by tuples.
        """
        values = self.frozenset_literals
        value_by_index = {index: value for value, index in values.items()}
        set_count = len(values)
        result = []
        result.append(str(set_count))
        for i in range(set_count):
            value = value_by_index[i]
            result.append(str(len(value)))
            for item in value:
                index = self.literal_index(cast(Any, item))
                result.append(str(index))

        count = len(values)
        result.append(str(count))
        for i in range(count):
           value = value_by_index[i]
           result.append(str(len(value)))
           for item in value:
               index = self.literal_index(cast(Any, item))
               result.append(str(index))
        return result


diff --git a/mypyc/ir/ops.py b/mypyc/ir/ops.py

    def __init__(
        self,
        value: None | str | bytes | bool | int | float | complex | tuple[object, ...] | frozenset[object],
        value: None
        | str
        | bytes
        | bool
        | int
        | float
        | complex
        | tuple[object, ...]
        | frozenset[object],
        rtype: RType,
    ) -> None:
        self.value = value
diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py

 from __future__ import annotations

 from typing import Callable, cast
 from typing import Callable,Sequence, Tuple, Union,cast

 from mypy.nodes import (
    ARG_POS,
    is_list_rprimitive,
    is_none_rprimitive,
    object_rprimitive,
    set_rprimitive,
 )
 from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional
 from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op
 from mypyc.primitives.list_ops import list_append_op, list_extend_op, list_slice_op
 from mypyc.primitives.misc_ops import ellipsis_op, get_module_dict_op, new_slice_op, type_op
 from mypyc.primitives.registry import CFunctionDescription, builtin_names
 from mypyc.primitives.set_ops import set_add_op, set_update_op
 from mypyc.primitives.set_ops import set_add_op,set_in_op,set_update_op
 from mypyc.primitives.str_ops import str_slice_op
 from mypyc.primitives.tuple_ops import list_tuple_op, tuple_slice_op

    return target


 def precompute_set_literal(builder: IRBuilder, s: SetExpr) -> Value | None:
    """Try to pre-compute a frozenset literal during module initialization.

    Return None if it's not possible.

    Only references to "simple" final variables, tuple literals (with items that
    are themselves supported), and other non-container literals are supported.
    """
    SupportedValue = Union[str, bytes, bool, int, float, complex, Tuple[object, ...], None]

    def _set_literal_final_values(items: Sequence[Expression]) -> list[SupportedValue] | None:
        values: list[SupportedValue] = []
        for item in items:
            if isinstance(item, NameExpr) and item.name in ("None", "True", "False"):
                if item.name == "None":
                    values.append(None)
                elif item.name == "True":
                    values.append(True)
                elif item.name == "False":
                    values.append(False)
            elif isinstance(item, (NameExpr, MemberExpr)) and isinstance(item.node, Var):
                if item.node.is_final:
                    v = item.node.final_value
                    if isinstance(v, (str, int, float, bool)):
                        values.append(v)
            elif isinstance(item, (StrExpr, BytesExpr, IntExpr, FloatExpr, ComplexExpr)):
                values.append(item.value)
            elif isinstance(item, TupleExpr):
                t = _set_literal_final_values(item.items)
                if t is not None:
                    values.append(tuple(t))

        if len(values) != len(items):
            return None
        return values

    values = _set_literal_final_values(s.items)
    if values is not None:
        return builder.add(LoadLiteral(frozenset(values), set_rprimitive))
    return None


 def transform_comparison_expr(builder: IRBuilder, e: ComparisonExpr) -> Value:
    # x in (...)/[...]
    # x not in (...)/[...]
            else:
                return builder.true()

    # x in {...}
    # x not in {...}
    if (
        first_op in ("in", "not in")
        and len(e.operators) == 1
        and isinstance(e.operands[1], SetExpr)
    ):
        literal = precompute_set_literal(builder, e.operands[1])
        if literal is not None:
            lhs = e.operands[0]
            v = builder.builder.call_c(
                set_in_op, [builder.accept(lhs), literal], e.line, bool_rprimitive
            )
            if first_op == "not in":
                return builder.unary_op(v, "not", e.line)
            return v

    if len(e.operators) == 1:
        # Special some common simple cases
        if first_op in ("is", "is not"):
                    right = builder.accept(right_expr, can_borrow=True)
                    return builder.compare_tagged(left, right, first_op, e.line)

    # x in {...}
    # x not in {...}
    if (e.operators[0] in ['in', 'not in']
            and len(e.operators) == 1
            and isinstance(e.operands[1], SetExpr)):
        negated = e.operators[0] == 'not in'
        items = e.operands[1].items
        if items:
            # Precalculate a frozenset at module top level and use that instead
            # of always rebuilding the set literal every evaluation.
            # NOTE: this only supports set literals which only contain items
            # which can be literals and aren't containers or sequences.
            literal_safe = True
            literal_values: List[Union[str, complex, bytes, int, float, None]] = []
            for item in items:
                if isinstance(item, NameExpr) and item.name == 'None':
                    literal_safe = literal_safe and True
                    literal_values.append(None)
                elif isinstance(item, (BytesExpr, StrExpr, IntExpr, FloatExpr, ComplexExpr)):
                    literal_safe = literal_safe and True
                    literal_values.append(item.value)
                else:
                    literal_safe = False
            if literal_safe:
                left_operand = builder.accept(e.operands[0])
                set_literal = builder.add(LoadLiteral(frozenset(literal_values), set_rprimitive))
                value = builder.builder.call_c(set_in_op,
                                               [left_operand, set_literal], e.line, bool_rprimitive)
                if negated:
                    value = builder.unary_op(value, 'not', e.line)
                return value

    # TODO: Don't produce an expression when used in conditional context
    # All of the trickiness here is due to support for chained conditionals
    # (`e1 < e2 > e3`, etc). `e1 < e2 > e3` is approximately equivalent to
diff --git a/mypyc/test-data/irbuild-set.test b/mypyc/test-data/irbuild-set.test
    r12 = PySet_Add(r0, r11)
    r13 = r12 >= 0 :: signed
    return r0

 [case testOperatorInLiteral]
 def f(i: object) -> bool:
    return i in {1, 2.7, "daylily", b"daylily", (None, (None,))}
 [out]
 def f(i):
    i :: object
    r0 :: set
    r1 :: int32
    r2 :: bit
    r3 :: bool
 L0:
    r0 = frozenset({1, 2.7, (None, (None,)), 'daylily'})
    r1 = PySet_Contains(r0, i)
    r2 = r1 >= 0 :: signed
    r3 = truncate r1: int32 to builtins.bool
    return r3
Original file line number	Diff line number	Diff line change
Expand Up		@@ -248,6 +248,15 @@ def check_tuple_items_valid_literals(self, op: LoadLiteral, t: tuple[object, ...
		if isinstance(x, tuple):
		self.check_tuple_items_valid_literals(op, x)

		def check_frozenset_items_valid_literals(
		self, op: LoadLiteral, s: frozenset[object, ...]
		) -> None:
		for x in s:
		if x is not None and not isinstance(x, (str, bytes, bool, int, float, complex, tuple)):
		self.fail(op, f"Invalid type for item of frozenset literal: {type(x)})")
		if isinstance(x, tuple):
		self.check_tuple_items_valid_literals(op, x)

		def visit_load_literal(self, op: LoadLiteral) -> None:
		expected_type = None
		if op.value is None:
Expand All		@@ -267,6 +276,11 @@ def visit_load_literal(self, op: LoadLiteral) -> None:
		elif isinstance(op.value, tuple):
		expected_type = "builtins.tuple"
		self.check_tuple_items_valid_literals(op, op.value)
		elif isinstance(op.value, frozenset):
		# There's no frozenset_rprimitive type since it'd be pretty useless so we just pretend
		# it's a set (when it's really a frozenset).
		expected_type = "builtins.set"
		self.check_frozenset_items_valid_literals(op, op.value)

		assert expected_type is not None, "Missed a case for LoadLiteral check"

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -670,7 +670,7 @@ def generate_literal_tables(self) -> None:
		self.declare_global("const int []", "CPyLit_Tuple", initializer=init_tuple)
		# Descriptions of frozenset literals
		init_frozenset = c_array_initializer(literals.encoded_frozenset_values())
		self.declare_global('const int []', 'CPyLit_FrozenSet', initializer=init_frozenset)
		self.declare_global("const int []", "CPyLit_FrozenSet", initializer=init_frozenset)

		def generate_export_table(self, decl_emitter: Emitter, code_emitter: Emitter) -> None:
		"""Generate the declaration and definition of the group's export struct.
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,7 +6,9 @@

		# Supported Python literal types. All tuple / frozenset items must have supported
		# literal types as well, but we can't represent the type precisely.
		LiteralValue = Union[str, bytes, int, bool, float, complex, Tuple[object, ...], FrozenSet[object], None]
		LiteralValue = Union[
		str, bytes, int, bool, float, complex, Tuple[object, ...], FrozenSet[object], None
		]

		# Some literals are singletons and handled specially (None, False and True)
		NUM_SINGLETONS: Final = 3
Expand DownExpand Up		@@ -130,48 +132,36 @@ def encoded_complex_values(self) -> list[str]:
		return _encode_complex_values(self.complex_literals)

		def encoded_tuple_values(self) -> list[str]:
		"""Encode tuple values into a C array.
		return self._encode_collection_values(self.tuple_literals)

		def encoded_frozenset_values(self) -> List[str]:
		return self._encode_collection_values(self.frozenset_literals)

		def _encode_collection_values(
		self, values: dict[tuple[object, ...], int] \| dict[frozenset[object], int]
		) -> list[str]:
		"""Encode tuple/frozen values into a C array.

		The format of the result is like this:

		<number oftuples>
		<length of the firsttuple>
		<number ofcollections>
		<length of the firstcollection>
		<literal index of first item>
		...
		<literal index of last item>
		<length of the secondtuple>
		<length of the secondcollection>
		...
		"""
		values = self.tuple_literals
		value_by_index = {index: value for value, index in values.items()}
		result = []
		num = len(values)
		result.append(str(num))
		for i in range(num):
		value = value_by_index[i]
		result.append(str(len(value)))
		for item in value:
		index = self.literal_index(cast(Any, item))
		result.append(str(index))
		return result

		def encoded_frozenset_values(self) -> List[str]:
		"""Encode frozenset values into a C array.

		The format used here is identical to one used by tuples.
		"""
		values = self.frozenset_literals
		value_by_index = {index: value for value, index in values.items()}
		set_count = len(values)
		result = []
		result.append(str(set_count))
		for i in range(set_count):
		value = value_by_index[i]
		result.append(str(len(value)))
		for item in value:
		index = self.literal_index(cast(Any, item))
		result.append(str(index))

		count = len(values)
		result.append(str(count))
		for i in range(count):
		value = value_by_index[i]
		result.append(str(len(value)))
		for item in value:
		index = self.literal_index(cast(Any, item))
		result.append(str(index))
		return result


Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -606,7 +606,15 @@ class LoadLiteral(RegisterOp):

		def __init__(
		self,
		value: None \| str \| bytes \| bool \| int \| float \| complex \| tuple[object, ...] \| frozenset[object],
		value: None
		\| str
		\| bytes
		\| bool
		\| int
		\| float
		\| complex
		\| tuple[object, ...]
		\| frozenset[object],
		rtype: RType,
		) -> None:
		self.value = value
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,7 +6,7 @@

		from __future__ import annotations

		from typing import Callable, cast
		from typing import Callable,Sequence, Tuple, Union,cast

		from mypy.nodes import (
		ARG_POS,
Expand DownExpand Up		@@ -68,6 +68,7 @@
		is_list_rprimitive,
		is_none_rprimitive,
		object_rprimitive,
		set_rprimitive,
		)
		from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional
		from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op
Expand All		@@ -92,7 +93,7 @@
		from mypyc.primitives.list_ops import list_append_op, list_extend_op, list_slice_op
		from mypyc.primitives.misc_ops import ellipsis_op, get_module_dict_op, new_slice_op, type_op
		from mypyc.primitives.registry import CFunctionDescription, builtin_names
		from mypyc.primitives.set_ops import set_add_op, set_update_op
		from mypyc.primitives.set_ops import set_add_op,set_in_op,set_update_op
		from mypyc.primitives.str_ops import str_slice_op
		from mypyc.primitives.tuple_ops import list_tuple_op, tuple_slice_op

Expand DownExpand Up		@@ -600,6 +601,48 @@ def transform_conditional_expr(builder: IRBuilder, expr: ConditionalExpr) -> Val
		return target


		def precompute_set_literal(builder: IRBuilder, s: SetExpr) -> Value \| None:
		"""Try to pre-compute a frozenset literal during module initialization.

		Return None if it's not possible.

		Only references to "simple" final variables, tuple literals (with items that
		are themselves supported), and other non-container literals are supported.
		"""
		SupportedValue = Union[str, bytes, bool, int, float, complex, Tuple[object, ...], None]

		def _set_literal_final_values(items: Sequence[Expression]) -> list[SupportedValue] \| None:
		values: list[SupportedValue] = []
		for item in items:
		if isinstance(item, NameExpr) and item.name in ("None", "True", "False"):
		if item.name == "None":
		values.append(None)
		elif item.name == "True":
		values.append(True)
		elif item.name == "False":
		values.append(False)
		elif isinstance(item, (NameExpr, MemberExpr)) and isinstance(item.node, Var):
		if item.node.is_final:
		v = item.node.final_value
		if isinstance(v, (str, int, float, bool)):
		values.append(v)
		elif isinstance(item, (StrExpr, BytesExpr, IntExpr, FloatExpr, ComplexExpr)):
		values.append(item.value)
		elif isinstance(item, TupleExpr):
		t = _set_literal_final_values(item.items)
		if t is not None:
		values.append(tuple(t))

		if len(values) != len(items):
		return None
		return values

		values = _set_literal_final_values(s.items)
		if values is not None:
		return builder.add(LoadLiteral(frozenset(values), set_rprimitive))
		return None


		def transform_comparison_expr(builder: IRBuilder, e: ComparisonExpr) -> Value:
		# x in (...)/[...]
		# x not in (...)/[...]
Expand DownExpand Up		@@ -653,6 +696,23 @@ def transform_comparison_expr(builder: IRBuilder, e: ComparisonExpr) -> Value:
		else:
		return builder.true()

		# x in {...}
		# x not in {...}
		if (
		first_op in ("in", "not in")
		and len(e.operators) == 1
		and isinstance(e.operands[1], SetExpr)
		):
		literal = precompute_set_literal(builder, e.operands[1])
		if literal is not None:
		lhs = e.operands[0]
		v = builder.builder.call_c(
		set_in_op, [builder.accept(lhs), literal], e.line, bool_rprimitive
		)
		if first_op == "not in":
		return builder.unary_op(v, "not", e.line)
		return v

		if len(e.operators) == 1:
		# Special some common simple cases
		if first_op in ("is", "is not"):
Expand All		@@ -670,38 +730,6 @@ def transform_comparison_expr(builder: IRBuilder, e: ComparisonExpr) -> Value:
		right = builder.accept(right_expr, can_borrow=True)
		return builder.compare_tagged(left, right, first_op, e.line)

		# x in {...}
		# x not in {...}
		if (e.operators[0] in ['in', 'not in']
		and len(e.operators) == 1
		and isinstance(e.operands[1], SetExpr)):
		negated = e.operators[0] == 'not in'
		items = e.operands[1].items
		if items:
		# Precalculate a frozenset at module top level and use that instead
		# of always rebuilding the set literal every evaluation.
		# NOTE: this only supports set literals which only contain items
		# which can be literals and aren't containers or sequences.
		literal_safe = True
		literal_values: List[Union[str, complex, bytes, int, float, None]] = []
		for item in items:
		if isinstance(item, NameExpr) and item.name == 'None':
		literal_safe = literal_safe and True
		literal_values.append(None)
		elif isinstance(item, (BytesExpr, StrExpr, IntExpr, FloatExpr, ComplexExpr)):
		literal_safe = literal_safe and True
		literal_values.append(item.value)
		else:
		literal_safe = False
		if literal_safe:
		left_operand = builder.accept(e.operands[0])
		set_literal = builder.add(LoadLiteral(frozenset(literal_values), set_rprimitive))
		value = builder.builder.call_c(set_in_op,
		[left_operand, set_literal], e.line, bool_rprimitive)
		if negated:
		value = builder.unary_op(value, 'not', e.line)
		return value

		# TODO: Don't produce an expression when used in conditional context
		# All of the trickiness here is due to support for chained conditionals
		# (`e1 < e2 > e3`, etc). `e1 < e2 > e3` is approximately equivalent to
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -655,3 +655,21 @@ L0:
		r12 = PySet_Add(r0, r11)
		r13 = r12 >= 0 :: signed
		return r0

		[case testOperatorInLiteral]
		def f(i: object) -> bool:
		return i in {1, 2.7, "daylily", b"daylily", (None, (None,))}
		[out]
		def f(i):
		i :: object
		r0 :: set
		r1 :: int32
		r2 :: bit
		r3 :: bool
		L0:
		r0 = frozenset({1, 2.7, (None, (None,)), 'daylily'})
		r1 = PySet_Contains(r0, i)
		r2 = r1 >= 0 :: signed
		r3 = truncate r1: int32 to builtins.bool
		return r3