diff --git a/src/kirin/serialization/core/__init__.py b/src/kirin/serialization/core/__init__.py index 6852b0c98..a437eb9e1 100644 --- a/src/kirin/serialization/core/__init__.py +++ b/src/kirin/serialization/core/__init__.py @@ -1,4 +1,9 @@ from .serializable import Serializable as Serializable from .deserializable import Deserializable as Deserializable from .serializationunit import SerializationUnit as SerializationUnit -from .serializationmodule import SerializationModule as SerializationModule +from .serializationmodule import ( + CURRENT_CODEC_VERSION as CURRENT_CODEC_VERSION, + SUPPORTED_CODEC_VERSIONS as SUPPORTED_CODEC_VERSIONS, + CodecVersion as CodecVersion, + SerializationModule as SerializationModule, +) diff --git a/src/kirin/serialization/core/serializationmodule.py b/src/kirin/serialization/core/serializationmodule.py index 9657451c1..19e9360b3 100644 --- a/src/kirin/serialization/core/serializationmodule.py +++ b/src/kirin/serialization/core/serializationmodule.py @@ -1,11 +1,25 @@ +from enum import IntEnum, unique from typing import TYPE_CHECKING if TYPE_CHECKING: from kirin.serialization.core.serializationunit import SerializationUnit +@unique +class CodecVersion(IntEnum): + """Kirin-owned serialization codec versions.""" + + LEGACY_UNVERSIONED = 1 + V2 = 2 + + +CURRENT_CODEC_VERSION = CodecVersion.V2 +SUPPORTED_CODEC_VERSIONS: frozenset[CodecVersion] = frozenset({CURRENT_CODEC_VERSION}) + + class SerializationModule: body: "SerializationUnit" + codec_version: CodecVersion version: str def __init__( @@ -14,6 +28,7 @@ def __init__( version: str = "", ): self.body = body + self.codec_version = CURRENT_CODEC_VERSION self.version = version def check_version(self, expect_version: str) -> bool: diff --git a/src/kirin/serialization/jsonserializer.py b/src/kirin/serialization/jsonserializer.py index 56b95c28d..2ada22dc3 100644 --- a/src/kirin/serialization/jsonserializer.py +++ b/src/kirin/serialization/jsonserializer.py @@ -2,7 +2,12 @@ from typing import Any, Optional from kirin.serialization.core.serializationunit import SerializationUnit -from kirin.serialization.core.serializationmodule import SerializationModule +from kirin.serialization.core.serializationmodule import ( + CURRENT_CODEC_VERSION, + SUPPORTED_CODEC_VERSIONS, + CodecVersion, + SerializationModule, +) COMPACT_UNIT_TAG = "$u" # Compact SerializationUnit wrapper. ESCAPED_MAP_TAG = "$m" # Escaped singleton user mapping. @@ -50,6 +55,7 @@ def _to_jsonifiable(self, obj: Any) -> Any: if isinstance(obj, SerializationModule): return { "__serialization_module__": True, + "codec_version": int(CURRENT_CODEC_VERSION), "version": obj.version, "body": self._to_jsonifiable(obj.body), } @@ -75,6 +81,7 @@ def _to_jsonifiable(self, obj: Any) -> Any: def _from_jsonifiable(self, obj: Any, *, allow_compact_tags: bool = True) -> Any: if isinstance(obj, dict): if obj.get("__serialization_module__"): + codec_version = self._decode_codec_version(obj) raw_body = obj.get("body") is_verbose = isinstance(raw_body, dict) and bool( raw_body.get("__serialization_unit__") @@ -84,7 +91,9 @@ def _from_jsonifiable(self, obj: Any, *, allow_compact_tags: bool = True) -> Any raw_body, allow_compact_tags=child_allow_compact_tags ) version = obj.get("version", "") - return SerializationModule(body=body, version=version) + module = SerializationModule(body=body, version=version) + module.codec_version = codec_version + return module if obj.get("__serialization_unit__"): data = self._from_jsonifiable( obj.get("data", {}), allow_compact_tags=False @@ -112,6 +121,33 @@ def _from_jsonifiable(self, obj: Any, *, allow_compact_tags: bool = True) -> Any ] return obj + def _decode_codec_version(self, envelope: dict[str, Any]) -> CodecVersion: + if "codec_version" not in envelope: + raise ValueError( + "legacy unversioned serialization codec (v1) is unsupported; " + f"current codec version is v{int(CURRENT_CODEC_VERSION)}" + ) + + raw_version = envelope["codec_version"] + if isinstance(raw_version, bool) or not isinstance(raw_version, int): + raise ValueError(f"codec_version must be an integer, got {raw_version!r}") + + try: + codec_version = CodecVersion(raw_version) + except ValueError: + qualifier = "future " if raw_version > int(CURRENT_CODEC_VERSION) else "" + raise ValueError( + f"unsupported {qualifier}codec version {raw_version}; " + f"current codec version is {int(CURRENT_CODEC_VERSION)}" + ) from None + + if codec_version not in SUPPORTED_CODEC_VERSIONS: + raise ValueError( + f"unsupported codec version {int(codec_version)}; " + f"current codec version is {int(CURRENT_CODEC_VERSION)}" + ) + return codec_version + def _decode_compact_unit(self, payload: Any) -> SerializationUnit: if not isinstance(payload, list) or len(payload) not in (2, 4): raise ValueError( diff --git a/test/serialization/test_compact_codec.py b/test/serialization/test_compact_codec.py index f24fc7867..d5824e79c 100644 --- a/test/serialization/test_compact_codec.py +++ b/test/serialization/test_compact_codec.py @@ -10,7 +10,10 @@ from kirin.serialization.bsonserializer import CompressedBSONSerializer from kirin.serialization.jsonserializer import JSONtifiable, JSONSerializer from kirin.serialization.core.serializationunit import SerializationUnit -from kirin.serialization.core.serializationmodule import SerializationModule +from kirin.serialization.core.serializationmodule import ( + CURRENT_CODEC_VERSION, + SerializationModule, +) Transport = tuple[ Callable[[SerializationModule], str | bytes], @@ -331,83 +334,34 @@ def test_public_transports_preserve_caller_version( payload = encode(module) + assert load(payload)["codec_version"] == int(CURRENT_CODEC_VERSION) assert load(payload)["version"] == version + assert decode(payload).codec_version is CURRENT_CODEC_VERSION assert decode(payload).version == version -VERBOSE_INT = { - "__serialization_unit__": True, - "kind": "int", - "module_name": "builtins", - "class_name": "int", - "data": {"value": str(2**256 + 1)}, -} -VERBOSE_MODULE = { - "__serialization_module__": True, - "version": "caller-version", - "body": { - "__serialization_unit__": True, - "kind": "list", - "module_name": "builtins", - "class_name": "list", - "data": {"value": [VERBOSE_INT]}, - }, -} - -VERBOSE_CONTROL_TAG_DATA = { - "unit-shaped": {"$u": ["int", {"value": "5"}]}, - "map-shaped": {"$m": [["user", "value"]]}, -} -VERBOSE_CONTROL_TAG_MODULE = { +LEGACY_UNVERSIONED_MODULE = { "__serialization_module__": True, "version": "caller-version", - "body": { - "__serialization_unit__": True, - "kind": "custom", - "module_name": "test.extension", - "class_name": "CustomUnit", - "data": VERBOSE_CONTROL_TAG_DATA, - }, + "body": {"$u": None}, } @pytest.mark.parametrize("transport_name", ["json", "bson"]) -def test_public_readers_accept_verbose_v1_and_reencode_compact( - transport_name: str, -) -> None: - if transport_name == "json": - serializer = JSONSerializer() - decoded = serializer.decode(json.dumps(VERBOSE_MODULE)) - compact_wire = json.loads(serializer.encode(decoded)) - else: - serializer = CompressedBSONSerializer() - payload = gzip.compress(bson.encode(VERBOSE_MODULE), mtime=0) - decoded = serializer.decode(payload) - compact_wire = bson.decode(gzip.decompress(serializer.encode(decoded))) - - assert decoded.version == "caller-version" - assert decoded.body.kind == "list" - assert decoded.body.data["value"][0].data["value"] == str(2**256 + 1) - assert set(compact_wire["body"]) == {"$u"} - - -@pytest.mark.parametrize("transport_name", ["json", "bson"]) -def test_verbose_v1_control_tag_shaped_mappings_remain_literal( +def test_public_readers_reject_legacy_unversioned_payloads( transport_name: str, ) -> None: if transport_name == "json": - module = JSONSerializer().decode(json.dumps(VERBOSE_CONTROL_TAG_MODULE)) + payload = json.dumps(LEGACY_UNVERSIONED_MODULE) + with pytest.raises( + ValueError, + match=r"legacy unversioned serialization codec \(v1\) is unsupported", + ): + JSONSerializer().decode(payload) else: - payload = gzip.compress(bson.encode(VERBOSE_CONTROL_TAG_MODULE), mtime=0) - module = CompressedBSONSerializer().decode(payload) - - assert module.body.data == VERBOSE_CONTROL_TAG_DATA - - -def test_missing_verbose_version_keeps_the_existing_empty_default() -> None: - payload = dict(VERBOSE_MODULE) - payload.pop("version") - - decoded = JSONSerializer().decode(json.dumps(payload)) - - assert decoded.version == "" + payload = gzip.compress(bson.encode(LEGACY_UNVERSIONED_MODULE), mtime=0) + with pytest.raises( + ValueError, + match=r"legacy unversioned serialization codec \(v1\) is unsupported", + ): + CompressedBSONSerializer().decode(payload) diff --git a/test/serialization/test_ssa_refs.py b/test/serialization/test_ssa_refs.py index be69be51e..c5d974f2d 100644 --- a/test/serialization/test_ssa_refs.py +++ b/test/serialization/test_ssa_refs.py @@ -2,7 +2,6 @@ from __future__ import annotations -import copy import gzip import json from typing import Any, Literal @@ -342,70 +341,3 @@ def test_malformed_ssa_refs_fail_with_context( with pytest.raises(ValueError, match=error): straight_line.dialects.decode(module) - - -def _with_verbose_v1_operands(module: SerializationModule) -> SerializationModule: - module = copy.deepcopy(module) - definitions = { - unit.data["id"]: unit - for unit in _walk_units(module.body) - if unit.kind in ("block-arg", "result-value") - } - - for statement in ( - unit for unit in _walk_units(module.body) if unit.kind == "statement" - ): - operands = statement.data["_args"].data["value"] - statement.data["_args"].data["value"] = [ - copy.deepcopy(definitions[operand.data["id"]]) for operand in operands - ] - return module - - -def _to_verbose_v1(value: Any) -> Any: - if isinstance(value, SerializationModule): - return { - "__serialization_module__": True, - "version": value.version, - "body": _to_verbose_v1(value.body), - } - if isinstance(value, SerializationUnit): - return { - "__serialization_unit__": True, - "kind": value.kind, - "module_name": value.module_name, - "class_name": value.class_name, - "data": _to_verbose_v1(value.data), - } - if isinstance(value, dict): - return {key: _to_verbose_v1(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_to_verbose_v1(item) for item in value] - return value - - -@pytest.mark.parametrize("transport", ["json", "cbson"]) -def test_verbose_v1_full_ssa_operands_remain_readable(transport: str) -> None: - old_module = _with_verbose_v1_operands(branching.dialects.encode(branching)) - verbose = _to_verbose_v1(old_module) - - if transport == "json": - module = JSONSerializer().decode(json.dumps(verbose)) - else: - payload = gzip.compress(bson.encode(verbose), mtime=0) - module = CompressedBSONSerializer().decode(payload) - - operand_kinds = { - operand.kind - for statement in _walk_units(module.body) - if statement.kind == "statement" - for operand in statement.data["_args"].data["value"] - } - assert operand_kinds <= {"block-arg", "result-value"} - assert operand_kinds == {"block-arg", "result-value"} - - decoded = branching.dialects.decode(module) - _assert_exact_reverse_uses(decoded) - decoded.verify() - assert decoded(3, True) == branching(3, True) - assert decoded(3, False) == branching(3, False) diff --git a/test/serialization/test_version.py b/test/serialization/test_version.py index 9a325eb83..08d766eef 100644 --- a/test/serialization/test_version.py +++ b/test/serialization/test_version.py @@ -1,9 +1,16 @@ +import json + import pytest from kirin.prelude import basic from kirin.serialization.jsonserializer import JSONSerializer from kirin.serialization.base.serializer import Serializer -from kirin.serialization.core.serializationmodule import SerializationModule +from kirin.serialization.core.serializationmodule import ( + CURRENT_CODEC_VERSION, + SUPPORTED_CODEC_VERSIONS, + CodecVersion, + SerializationModule, +) @basic @@ -21,6 +28,14 @@ def test_serialization_module_default_version_is_empty(): assert mod.version == "" +def test_serialization_module_uses_current_codec_version(): + mod = SerializationModule(body=basic.encode(simple_kernel).body) + + assert mod.codec_version is CURRENT_CODEC_VERSION + assert CURRENT_CODEC_VERSION is CodecVersion.V2 + assert SUPPORTED_CODEC_VERSIONS == {CodecVersion.V2} + + def test_serialization_module_stores_version(): mod = _empty_module(version="1.2.3") assert mod.version == "1.2.3" @@ -66,10 +81,61 @@ def test_dialect_group_encode_with_version(): def test_encode_json_round_trips_version(): json_str = basic.encode_json(simple_kernel, version="0.9.0") + payload = json.loads(json_str) decoded_module = JSONSerializer().decode(json_str) + + assert payload["codec_version"] == int(CURRENT_CODEC_VERSION) + assert decoded_module.codec_version is CURRENT_CODEC_VERSION assert decoded_module.version == "0.9.0" +@pytest.mark.parametrize( + ("payload", "error"), + [ + pytest.param( + { + "__serialization_module__": True, + "body": {"$u": None}, + }, + r"legacy unversioned serialization codec \(v1\) is unsupported", + id="missing-is-legacy-v1", + ), + pytest.param( + { + "__serialization_module__": True, + "codec_version": 1, + "body": {"$u": None}, + }, + r"unsupported codec version 1", + id="explicit-v1", + ), + pytest.param( + { + "__serialization_module__": True, + "codec_version": 3, + "body": {"$u": None}, + }, + r"unsupported future codec version 3", + id="future-before-body", + ), + ], +) +def test_unsupported_codec_versions_are_rejected_before_body( + payload: dict[str, object], error: str +) -> None: + with pytest.raises(ValueError, match=error): + JSONSerializer().decode(json.dumps(payload)) + + +def test_missing_caller_version_keeps_empty_default() -> None: + payload = json.loads(basic.encode_json(simple_kernel)) + payload.pop("version") + + decoded = JSONSerializer().decode(json.dumps(payload)) + + assert decoded.version == "" + + def test_decode_json_no_expected_version_succeeds(): json_str = basic.encode_json(simple_kernel, version="1.0.0") method = basic.decode_json(json_str)