Change PyObject sort to be pickled form instead of reference

Author: saulshanabrook

.gitignore

@@ -85,3 +85,4 @@ inlined
 visualizer.tgz
 package
 .mypy_cache/
+*.json

Cargo.lock

@@ -29,6 +29,7 @@ lalrpop-util = { version = "0.22", features = ["lexer"] }
 ordered-float = "5"
 uuid = { version = "1.18", features = ["v4"] }
 rayon = "1.11"
+base64 = "0.22.1"
 
 # Use patched version of egglog in experimental
 [patch.'https://github.com/egraphs-good/egglog']

Cargo.toml

@@ -96,43 +96,55 @@ We define a custom "primitive sort" (i.e. a builtin type) for `PyObject`s. This
 
 ### Saving Python Objects
 
-To create an expression of type `PyObject`, we call the call the constructor with any Python object. It will
-save a reference to the object:
+To create an expression of type `PyObject`, call the constructor with any Python object. The value is immediately
+serialized with `cloudpickle.dumps`, and the serialized bytes (base64 encoded when printed) are what get stored
+inside the e-graph. This means the e-graph keeps a snapshot of the object rather than a live reference.
 
 ```{code-cell} python
-PyObject(1)
-```
-
-We see that this as saved internally as a pointer to the Python object. For hashable objects like `int` we store two integers, a hash of the type and a has of the value.
+from dataclasses import dataclass
 
-We can also store unhashable objects in the e-graph like lists.
+@dataclass
+class MyObject:
+    a: int = 10
 
-```{code-cell} python
-lst = PyObject([1, 2, 3])
-lst
+PyObject(MyObject())
 ```
 
-We see that this is stored with one number, simply the `id` of the object.
+The new serialization approach works for both hashable and unhashable Python values, and no longer depends on
+their `id()`. Subsequent inserts of equal values round-trip through `cloudpickle` so the e-graph can identify and
+merge them by value.
 
-```{admonition} Mutable Objects
-:class: warning
+```{admonition} Serialization requirements
+:class: note
 
-While it is possible to store unhashable objects in the e-graph, you have to be careful defining any rules which create new unhashable objects. If each time a rule is run, it creates a new object, then the e-graph will never saturate.
-
-Creating hashable objects is safer, since while the rule might create new Python objects each time it executes, they should have the same hash, i.e. be equal, so that the e-graph can saturate.
+`PyObject` relies on `cloudpickle`. Any object you store must be serializable by `cloudpickle.dumps`; objects such
+as open file handles, generators, or extension types that `cloudpickle` cannot handle will raise an error when you
+try to construct a `PyObject`.
 ```
 
 ### Retrieving Python Objects
 
-Like other primitives, we can retrieve the Python object from the e-graph by using the `.value` property:
+Like other primitives, we can retrieve a Python object by using the `.value` property. Deserialization happens on
+every access, so you receive a fresh copy each time rather than the original object.
 
 ```{code-cell} python
-assert lst.value == [1, 2, 3]
+original = {"count": 1}
+expr = PyObject(original)
+
+restored = expr.value
+assert restored == original
+assert restored is not original
+
+# Mutating the copy does not affect the stored value.
+restored["count"] = 2
+assert expr.value == {"count": 1}
 ```
 
 ### Builtin methods
 
-Currently, we only support a few methods on `PyObject`s, but we plan to add more in the future.
+Currently, we only support a few methods on `PyObject`s, but we plan to add more in the future. Each builtin
+deserializes its inputs, performs the operation in Python, and then serializes the result back into a new
+`PyObject`, so previously stored values remain unchanged.
 
 Conversion to/from a string:
 

docs/reference/python-integration.md

@@ -37,6 +37,7 @@ array = [
     "numba>=0.59.1",
     "llvmlite>=0.42.0",
     "numpy>2",
+    "cloudpickle>=3",
 ]
 dev = [
     "ruff",

pyproject.toml

@@ -54,7 +54,6 @@ __all__ = [
     "PrintOverallStatistics",
     "PrintSize",
     "Push",
-    "PyObjectSort",
     "Relation",
     "Repeat",
     "Rewrite",
@@ -105,17 +104,10 @@ class SerializedEGraph:
     def map_ops(self, map: dict[str, str]) -> None: ...
     def split_classes(self, egraph: EGraph, ops: set[str]) -> None: ...
 
-@final
-class PyObjectSort:
-    def __init__(self) -> None: ...
-    def store(self, __o: object, /) -> _Expr: ...
-    def load(self, __e: _Expr, /) -> object: ...
-
 @final
 class EGraph:
     def __init__(
         self,
-        py_object_sort: PyObjectSort | None = None,
         /,
         *,
         fact_directory: str | Path | None = None,
@@ -142,7 +134,7 @@ class EGraph:
     def value_to_rational(self, v: Value) -> Fraction: ...
     def value_to_bigint(self, v: Value) -> int: ...
     def value_to_bigrat(self, v: Value) -> Fraction: ...
-    def value_to_pyobject(self, py_object_sort: PyObjectSort, v: Value) -> object: ...
+    def value_to_pyobject(self, v: Value) -> object: ...
     def value_to_map(self, v: Value) -> dict[Value, Value]: ...
     def value_to_multiset(self, v: Value) -> list[Value]: ...
     def value_to_vec(self, v: Value) -> list[Value]: ...

python/egglog/bindings.pyi

@@ -13,6 +13,7 @@
 from types import FunctionType, MethodType
 from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, cast, overload
 
+import cloudpickle
 from typing_extensions import TypeVarTuple, Unpack, deprecated
 
 from .conversion import convert, converter, get_type_args, resolve_literal
@@ -984,12 +985,15 @@ def value(self) -> object:
         expr = cast("RuntimeExpr", self).__egg_typed_expr__.expr
         if not isinstance(expr, PyObjectDecl):
             raise ExprValueError(self, "PyObject(x)")
-        return expr.value
+        return cloudpickle.loads(expr.pickled)
 
     __match_args__ = ("value",)
 
     def __init__(self, value: object) -> None: ...
 
+    @method(egg_fn="py-call")
+    def __call__(self, *args: object) -> PyObject: ...
+
     @method(egg_fn="py-from-string")
     @classmethod
     def from_string(cls, s: StringLike) -> PyObject: ...
@@ -1023,24 +1027,15 @@ class PyObjectFunction(Protocol):
     def __call__(self, *__args: PyObject) -> PyObject: ...
 
 
+@deprecated("use PyObject(fn) directly")
 def py_eval_fn(fn: Callable) -> PyObjectFunction:
     """
     Takes a python callable and maps it to a callable which takes and returns PyObjects.
 
     It translates it to a call which uses `py_eval` to call the function, passing in the
     args as locals, and using the globals from function.
     """
-
-    def inner(*__args: PyObject, __fn: Callable = fn) -> PyObject:
-        new_kvs: list[object] = []
-        eval_str = "__fn("
-        for i, arg in enumerate(__args):
-            new_kvs.extend((f"__arg_{i}", arg))
-            eval_str += f"__arg_{i}, "
-        eval_str += ")"
-        return py_eval(eval_str, PyObject({"__fn": __fn}).dict_update(*new_kvs), getattr(__fn, "__globals__", {}))
-
-    return inner
+    return PyObject(fn)
 
 
 @function(builtin=True, egg_fn="py-exec")

python/egglog/builtins.py

@@ -617,23 +617,7 @@ class LetRefDecl:
 
 @dataclass(frozen=True)
 class PyObjectDecl:
-    value: object
-
-    def __hash__(self) -> int:
-        """Tries using the hash of the value, if unhashable use the ID."""
-        try:
-            return hash((type(self.value), self.value))
-        except TypeError:
-            return id(self.value)
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, PyObjectDecl):
-            return False
-        return self.parts == other.parts
-
-    @property
-    def parts(self) -> tuple[type, object]:
-        return (type(self.value), self.value)
+    pickled: bytes
 
 
 LitType: TypeAlias = int | str | float | bool | None

python/egglog/declarations.py

@@ -8,6 +8,7 @@
 from functools import partial
 from typing import TYPE_CHECKING, TypeVar, overload
 
+import cloudpickle
 from typing_extensions import TypeVarTuple, Unpack
 
 from .declarations import *
@@ -64,7 +65,7 @@ def get_literal_value(x: object) -> object:
         case LitDecl(v):
             return v
         case PyObjectDecl(obj):
-            return obj
+            return cloudpickle.loads(obj)
         case PartialCallDecl(call):
             fn, args = _deconstruct_call_decl(x.__egg_decls_thunk__, call)
             if not args:

python/egglog/deconstruct.py

@@ -839,7 +839,7 @@ class EGraph:
     _token_stack: list[EGraph] = field(default_factory=list, repr=False)
 
     def __post_init__(self, seminaive: bool, save_egglog_string: bool) -> None:
-        egraph = bindings.EGraph(GLOBAL_PY_OBJECT_SORT, seminaive=seminaive, record=save_egglog_string)
+        egraph = bindings.EGraph(seminaive=seminaive, record=save_egglog_string)
         self._state = EGraphState(egraph)
 
     def _add_decls(self, *decls: DeclerationsLike) -> None: