diff --git a/src/kirin/ir/attrs/types.py b/src/kirin/ir/attrs/types.py index ece9555de..e38d719c6 100644 --- a/src/kirin/ir/attrs/types.py +++ b/src/kirin/ir/attrs/types.py @@ -368,7 +368,11 @@ def deserialize( ) -> "Literal": d = deserializer.deserialize(serUnit.data["value"]) type_attr = deserializer.deserialize(serUnit.data["type"]) - return Literal(d, type_attr) + # Reconstruct the serialized graph without reusing a process-local + # LiteralMeta cache entry that may reference a different type instance. + literal = object.__new__(cls) + literal.__init__(d, type_attr) + return literal @typing.final diff --git a/src/kirin/serialization/base/context.py b/src/kirin/serialization/base/context.py index 51337ac86..dbdef4437 100644 --- a/src/kirin/serialization/base/context.py +++ b/src/kirin/serialization/base/context.py @@ -34,6 +34,7 @@ class SerializationContext: Dialect_Lookup: dict[str, ir.Dialect] = field(default_factory=dict) SSA_Lookup: dict[str, ir.SSAValue] = field(default_factory=dict) + TypeAttribute_Lookup: dict[str, types.TypeAttribute] = field(default_factory=dict) Statement_Lookup: dict[str, ir.Statement] = field(default_factory=dict) Block_Lookup: dict[str, ir.Block] = field(default_factory=dict) Region_Lookup: dict[str, ir.Region] = field(default_factory=dict) @@ -56,6 +57,7 @@ class SerializationContext: def clear(self) -> None: self.SSA_Lookup.clear() + self.TypeAttribute_Lookup.clear() self.Block_Lookup.clear() self.Region_Lookup.clear() self.Statement_Lookup.clear() diff --git a/src/kirin/serialization/base/deserializer.py b/src/kirin/serialization/base/deserializer.py index 71638035a..80cf2f812 100644 --- a/src/kirin/serialization/base/deserializer.py +++ b/src/kirin/serialization/base/deserializer.py @@ -78,6 +78,10 @@ def deserialize_attr_ref(self, serUnit: SerializationUnit) -> types.TypeAttribut if not isinstance(attr_id, str): raise ValueError(f"attr_ref id must be a string, got {attr_id!r}") + existing = self._ctx.TypeAttribute_Lookup.get(attr_id) + if existing is not None: + return existing + definition = self._ctx.TypeAttribute_Definitions.get(attr_id) if definition is None: self._index_type_attribute_definitions() @@ -417,11 +421,16 @@ def deserialize_attribute(self, serUnit: SerializationUnit) -> ir.Attribute: def _deserialize_type_attribute_definition( self, serUnit: SerializationUnit, attr_id: str ) -> types.TypeAttribute: + existing = self._ctx.TypeAttribute_Lookup.get(attr_id) + if existing is not None: + return existing + attr = self._deserialize_attribute_full(serUnit) if not isinstance(attr, types.TypeAttribute): raise ValueError( f"TypeAttribute definition {attr_id!r} decoded as {type(attr).__name__}" ) + self._ctx.TypeAttribute_Lookup[attr_id] = attr return attr def _deserialize_attribute_full(self, serUnit: SerializationUnit) -> ir.Attribute: diff --git a/test/serialization/test_type_attribute_refs.py b/test/serialization/test_type_attribute_refs.py index 4f66315e3..78aacddcc 100644 --- a/test/serialization/test_type_attribute_refs.py +++ b/test/serialization/test_type_attribute_refs.py @@ -57,7 +57,7 @@ def _through_transport( @pytest.mark.parametrize("transport", TRANSPORTS) -def test_repeated_type_attribute_uses_one_definition_and_preserves_v1_isolation( +def test_repeated_type_attribute_preserves_shared_identity( transport: Transport, ) -> None: shared = types.TypeVar("Shared") @@ -79,13 +79,10 @@ def test_repeated_type_attribute_uses_one_definition_and_preserves_v1_isolation( decoded = method.dialects.decode(transported) first_decoded, second_decoded = decoded.fields - assert first_decoded is not second_decoded + assert first_decoded is not shared + assert first_decoded is second_decoded assert first_decoded == second_decoded assert isinstance(first_decoded, types.TypeVar) - - new_bound = types.TypeVar("NewBound") - first_decoded.bound = new_bound - assert second_decoded.bound is not new_bound assert decoded(7) == 7 @@ -115,7 +112,7 @@ def test_standalone_attribute_serialization_uses_a_reference() -> None: deserializer = Deserializer(passthrough.dialects) first_decoded = deserializer.deserialize(first) second_decoded = deserializer.deserialize(second) - assert first_decoded is not second_decoded + assert first_decoded is second_decoded assert first_decoded == second_decoded @@ -137,6 +134,8 @@ def test_distinct_but_equal_type_attributes_are_not_interned( assert "id" not in second_unit.data decoded = method.dialects.decode(_through_transport(module, transport)) + assert decoded.fields[0] is not first + assert decoded.fields[1] is not second assert decoded.fields[0] is not decoded.fields[1] assert decoded.fields[0] == decoded.fields[1] @@ -156,7 +155,7 @@ def test_general_pyattr_is_not_memoized(transport: Transport) -> None: decoded = method.dialects.decode(_through_transport(module, transport)) assert decoded.fields[0] is not decoded.fields[1] - assert decoded.fields[0].type is not decoded.fields[1].type + assert decoded.fields[0].type is decoded.fields[1].type assert decoded.fields[0].type == decoded.fields[1].type @@ -178,36 +177,98 @@ def test_forward_ref_from_reversed_pyattr_field_order(transport: Transport) -> N decoded = method.dialects.decode(_through_transport(module, transport)) decoded_wrapper = decoded.fields[0] - assert decoded_wrapper.type is not decoded_wrapper.data + assert decoded_wrapper.type is decoded_wrapper.data assert decoded_wrapper.type == decoded_wrapper.data @pytest.mark.parametrize("transport", TRANSPORTS) -def test_nested_type_graph_preserves_v1_value_semantics( +def test_nested_type_graph_preserves_identity_relationships( transport: Transport, ) -> None: - leaf = types.TypeVar("Leaf") - literal = types.Literal(1, leaf) - union = types.Union(literal, types.String) - generic = types.Generic(tuple, leaf, types.Vararg(leaf)) - method = _method_with_fields(union, union, generic, generic) + leaf = types.TypeVar(f"Leaf-{transport}") + first = types.Generic(tuple, leaf) + second = types.Generic(list, leaf) + assert first.vars[0] is second.vars[0] + method = _method_with_fields(first, first, second, second) module = method.dialects.encode(method) decoded = method.dialects.decode(_through_transport(module, transport)) + decoded_first = decoded.fields[0] + decoded_second = decoded.fields[2] + + assert decoded_first is not first + assert decoded_second is not second + assert decoded_first is decoded.fields[1] + assert decoded_second is decoded.fields[3] + assert decoded_first is not decoded_second + assert decoded_first.vars[0] is decoded_second.vars[0] + + +@pytest.mark.parametrize("transport", TRANSPORTS) +def test_union_and_literal_roundtrip_by_value(transport: Transport) -> None: + literal = types.Literal(1) + union = types.Union(literal, types.String) + assert isinstance(union, types.Union) + method = _method_with_fields(union) + + decoded = method.dialects.decode( + _through_transport(method.dialects.encode(method), transport) + ) decoded_union = decoded.fields[0] - decoded_generic = decoded.fields[2] + assert isinstance(decoded_union, types.Union) + assert decoded_union.is_structurally_equal(union) + assert any( + isinstance(attr, types.Literal) and attr.is_structurally_equal(literal) + for attr in decoded_union.types + ) + - assert decoded_union is not decoded.fields[1] - assert decoded_union == decoded.fields[1] +@pytest.mark.parametrize("transport", TRANSPORTS) +def test_literal_preserves_shared_type_identity(transport: Transport) -> None: + leaf = types.TypeVar(f"LiteralLeaf-{transport}") + literal = types.Literal(transport, leaf) + union = types.Union(literal, types.String) + generic = types.Generic(tuple, leaf) + assert isinstance(union, types.Union) + assert literal.type is generic.vars[0] + method = _method_with_fields(union, generic) + + decoded = method.dialects.decode( + _through_transport(method.dialects.encode(method), transport) + ) + decoded_union = decoded.fields[0] + decoded_generic = decoded.fields[1] + assert isinstance(decoded_union, types.Union) decoded_literal = next( attr for attr in decoded_union.types if isinstance(attr, types.Literal) ) + assert decoded_literal is not literal + assert decoded_literal.type is decoded_generic.vars[0] + + +@pytest.mark.parametrize("transport", TRANSPORTS) +def test_nested_repeated_type_attributes_preserve_dag_topology( + transport: Transport, +) -> None: + depth = 8 + typ: types.TypeAttribute = types.TypeVar("Leaf") + for _ in range(depth): + typ = types.Generic(tuple, typ, types.Vararg(typ)) + + method = _method_with_fields(typ) + decoded = method.dialects.decode( + _through_transport(method.dialects.encode(method), transport) + ) + decoded_type = decoded.fields[0] + + for _ in range(depth): + assert isinstance(decoded_type, types.Generic) + assert decoded_type.vararg is not None + nested_type = decoded_type.vars[0] + assert nested_type is decoded_type.vararg.typ + decoded_type = nested_type - assert decoded_generic is not decoded.fields[3] - assert decoded_generic == decoded.fields[3] - assert decoded_literal.is_structurally_equal(literal) - assert decoded_generic.vars[0] is not decoded_generic.vararg.typ - assert decoded_generic.vars[0] == decoded_generic.vararg.typ + assert isinstance(decoded_type, types.TypeVar) @pytest.mark.parametrize("transport_name", ["json", "cbson"])