Serializer framework for Kirin (#492)

Serializer framework. With base Serializer, JSONSerializer, and
BinarySerializer.

---------

Co-authored-by: kaihsin <[email protected]>

Author: zhenrongliew

src/kirin/analysis/const/lattice.py

@@ -51,6 +51,11 @@ def is_subseteq(self, other: Result) -> bool:
     def __hash__(self) -> int:
         return id(self)
 
+    def is_structurally_equal(
+        self, other: ir.Attribute, context: dict | None = None
+    ) -> bool:
+        return isinstance(other, Unknown)
+
 
 @final
 @dataclass
@@ -63,6 +68,11 @@ def is_subseteq(self, other: Result) -> bool:
     def __hash__(self) -> int:
         return id(self)
 
+    def is_structurally_equal(
+        self, other: ir.Attribute, context: dict | None = None
+    ) -> bool:
+        return isinstance(other, Bottom)
+
 
 @final
 @dataclass
@@ -85,6 +95,15 @@ def __hash__(self) -> int:
         # the data is guaranteed to be unique.
         return id(self)
 
+    def is_structurally_equal(
+        self, other: ir.Attribute, context: dict | None = None
+    ) -> bool:
+        if not isinstance(other, Value):
+            return False
+        if isinstance(self.data, ir.Attribute) and isinstance(other.data, ir.Attribute):
+            return self.data.is_structurally_equal(other.data, context=context)
+        return self.data.is_structurally_equal(other.data, context=context)
+
 
 @dataclass
 class PartialConst(Result):
@@ -158,6 +177,18 @@ def is_subseteq_Value(self, other: Value) -> bool:
     def __hash__(self) -> int:
         return hash(self.data)
 
+    def is_structurally_equal(
+        self, other: ir.Attribute, context: dict | None = None
+    ) -> bool:
+        if not isinstance(other, PartialTuple):
+            return False
+        if len(self.data) != len(other.data):
+            return False
+        return all(
+            x.is_structurally_equal(y, context=context)
+            for x, y in zip(self.data, other.data)
+        )
+
 
 @final
 @dataclass
@@ -230,3 +261,17 @@ def meet(self, other: Result) -> Result:
             tuple(x.meet(y) for x, y in zip(self.captured, other.captured)),
             self.argnames,
         )
+
+    def is_structurally_equal(
+        self, other: ir.Attribute, context: dict | None = None
+    ) -> bool:
+        return (
+            isinstance(other, PartialLambda)
+            and self.code.is_structurally_equal(other.code, context=context)
+            and self.argnames == other.argnames
+            and len(self.captured) == len(other.captured)
+            and all(
+                x.is_structurally_equal(y, context=context)
+                for x, y in zip(self.captured, other.captured)
+            )
+        )

src/kirin/dialects/func/attrs.py

@@ -1,9 +1,14 @@
-from typing import Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar
 from dataclasses import dataclass
 
 from kirin import types
 from kirin.ir import Method, Attribute
 from kirin.print.printer import Printer
+from kirin.serialization.base.serializationunit import SerializationUnit
+
+if TYPE_CHECKING:
+    from kirin.serialization.base.serializer import Serializer
+    from kirin.serialization.base.deserializer import Deserializer
 
 from ._dialect import dialect
 
@@ -38,3 +43,24 @@ def print_impl(self, printer: Printer) -> None:
         printer.print_seq(self.inputs, delim=", ", prefix="(", suffix=")")
         printer.plain_print(" -> ")
         printer.print(self.output)
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, Signature):
+            return False
+        return self.inputs == value.inputs and self.output == value.output
+
+    def serialize(self, serializer: "Serializer") -> "SerializationUnit":
+        return serializer.serialize_signature(self)
+
+    def is_structurally_equal(
+        self,
+        other: Attribute,
+        context: dict | None = None,
+    ) -> bool:
+        return self == other
+
+    @classmethod
+    def deserialize(
+        cls, serUnit: "SerializationUnit", deserializer: "Deserializer"
+    ) -> "Signature":
+        return deserializer.deserialize_signature(serUnit)

src/kirin/dialects/ilist/runtime.py

@@ -1,10 +1,15 @@
 # TODO: replace with something faster
-from typing import Any, Generic, TypeVar, overload
+from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload
 from dataclasses import dataclass
 from collections.abc import Sequence
 
 from kirin import ir, types
 from kirin.print.printer import Printer
+from kirin.serialization.base.serializationunit import SerializationUnit
+
+if TYPE_CHECKING:
+    from kirin.serialization.base.serializer import Serializer
+    from kirin.serialization.base.deserializer import Deserializer
 
 T = TypeVar("T")
 L = TypeVar("L")
@@ -85,8 +90,25 @@ def unwrap(self) -> Sequence[T]:
         return self
 
     def print_impl(self, printer: Printer) -> None:
-        printer.plain_print("IList(")
         printer.print_seq(
             self.data, delim=", ", prefix="[", suffix="]", emit=printer.plain_print
         )
         printer.plain_print(")")
+
+    def is_structurally_equal(
+        self, other: ir.Attribute, context: dict | None = None
+    ) -> bool:
+        return (
+            isinstance(other, IList)
+            and self.data == other.data
+            and self.elem.is_equal(other.elem)
+        )
+
+    def serialize(self, serializer: "Serializer") -> "SerializationUnit":
+        return serializer.serialize_ilist(self)
+
+    @classmethod
+    def deserialize(
+        cls, serUnit: "SerializationUnit", deserializer: "Deserializer"
+    ) -> "IList":
+        return deserializer.deserialize_ilist(serUnit)

src/kirin/dialects/py/slice.py

@@ -89,6 +89,16 @@ def __hash__(self):
     def print_impl(self, printer: Printer) -> None:
         return printer.plain_print(f"slice({self.start}, {self.stop}, {self.step})")
 
+    def is_structurally_equal(
+        self, other: ir.Attribute, context: dict | None = None
+    ) -> bool:
+        return (
+            isinstance(other, SliceAttribute)
+            and self.start == other.start
+            and self.stop == other.stop
+            and self.step == other.step
+        )
+
 
 @dialect.register
 class Concrete(interp.MethodTable):

src/kirin/idtable.py

@@ -55,3 +55,14 @@ def __getitem__(self, value: T) -> str:
             return self.table[value]
         else:
             return self.add(value)
+
+    def __len__(self) -> int:
+        return len(self.table)
+
+    def size(self) -> int:
+        return len(self)
+
+    def clear(self) -> None:
+        self.table.clear()
+        self.name_count.clear()
+        self.next_id = 0

src/kirin/ir/attrs/abc.py

@@ -2,6 +2,8 @@
 from typing import TYPE_CHECKING, TypeVar, ClassVar, Optional
 from dataclasses import field, dataclass
 
+from typing_extensions import Self
+
 from kirin.print import Printable
 from kirin.ir.traits import Trait
 from kirin.lattice.abc import LatticeMeta, SingletonMeta
@@ -53,6 +55,10 @@ def __hash__(self) -> int: ...
     @abstractmethod
     def __eq__(self, value: object) -> bool: ...
 
+    @abstractmethod
+    def is_structurally_equal(self, other: Self, context: dict | None = None) -> bool:
+        return self == other
+
     @classmethod
     def has_trait(cls, trait_type: type[Trait["Attribute"]]) -> bool:
         """Check if the Statement has a specific trait.

src/kirin/ir/attrs/py.py

@@ -1,7 +1,15 @@
-from typing import TypeVar
+from typing import TYPE_CHECKING, Any, Type, TypeVar
 from dataclasses import dataclass
 
+from typing_extensions import Protocol, runtime_checkable
+
 from kirin.print import Printer
+from kirin.ir.attrs.abc import Attribute
+from kirin.serialization.base.serializationunit import SerializationUnit
+
+if TYPE_CHECKING:
+    from kirin.serialization.base.serializer import Serializer
+    from kirin.serialization.base.deserializer import Deserializer
 
 from .data import Data
 from .types import PyClass, TypeAttribute
@@ -51,3 +59,32 @@ def print_impl(self, printer: Printer) -> None:
 
     def unwrap(self) -> T:
         return self.data
+
+    def is_structurally_equal(
+        self, other: Attribute, context: dict | None = None
+    ) -> bool:
+        if not isinstance(other, PyAttr):
+            return False
+        if self.type != other.type:
+            return False
+        if isinstance(self.data, StructurallyEqual) and isinstance(
+            other.data, StructurallyEqual
+        ):
+            return self.data.is_structurally_equal(other.data, context=context)
+        return self.data == other.data
+
+    def serialize(self, serializer: "Serializer") -> "SerializationUnit":
+        return serializer.serialize_pyattr(self)
+
+    @classmethod
+    def deserialize(
+        cls: Type["PyAttr"], serUnit: "SerializationUnit", deserializer: "Deserializer"
+    ) -> "PyAttr":
+        return deserializer.deserialize_pyattr(serUnit)
+
+
+@runtime_checkable
+class StructurallyEqual(Protocol):
+    def is_structurally_equal(
+        self, other: Any, context: dict | None = None
+    ) -> bool: ...