Add explicit serialization codec versioning

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,
+)

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:

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(

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)

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)