Fix all tests

Author: saulshanabrook

Cargo.lock

@@ -1,5 +1,5 @@
 [package]
-name = "egglog"
+name = "egglog_python"
 version = "7.2.0"
 edition = "2021"
 
@@ -13,10 +13,17 @@ pyo3 = { version = "0.21", features = ["extension-module"] }
 
 # https://github.com/egraphs-good/egglog/compare/ceed816e9369570ffed9feeba157b19471dda70d...main
 # egglog = { git = "https://github.com/egraphs-good/egglog", rev = "fb4a9f114f9bb93154d6eff0dbab079b5cb4ebb6" }
-egglog = { path = "../egg-smol" }
+# egglog = { path = "../egg-smol" }
 # egglog = { git = "https://github.com/oflatt/egg-smol", branch = "oflatt-fast-terms" }
-# egglog = { git = "https://github.com/saulshanabrook/egg-smol", rev = "38b3014b34399cc78887ede09c845b2a5d6c7d19" }
-egraph-serialize = { git = "https://github.com/egraphs-good/egraph-serialize", rev = "5838c036623e91540831745b1574539e01c8cb23" }
+egglog = { git = "https://github.com/saulshanabrook/egg-smol", rev = "a555b2f5e82c684442775cc1a5da94b71930113c" }
+egraph-serialize = { git = "https://github.com/saulshanabrook/egraph-serialize", rev = "1c205fcc6d3426800b828e9264dbadbd4a5ef6e9", features = [
+    "serde",
+    "graphviz",
+] }
+# egraph-serialize = { path = "../egraph-serialize", features = [
+#     "serde",
+#     "graphviz",
+# ] }
 serde_json = "*"
 pyo3-log = "0.10.0"
 log = "0.4.21"
@@ -26,8 +33,8 @@ uuid = { version = "1.8.0", features = ["v4"] }
 num-rational = "*"
 
 # Use unreleased version of egraph-serialize in egglog as well
-[patch.crates-io]
-egraph-serialize = { git = "https://github.com/egraphs-good/egraph-serialize", rev = "5838c036623e91540831745b1574539e01c8cb23" }
+# [patch.crates-io]
+# egraph-serialize = { git = "https://github.com/egraphs-good/egraph-serialize", rev = "5838c036623e91540831745b1574539e01c8cb23" }
 # egraph-serialize = { path = "../egraph-serialize" }
 
 # enable debug symbols for easier profiling

Cargo.toml

@@ -32,10 +32,10 @@ conda activate egglog-python
 Then install the package in editable mode with the development dependencies:
 
 ```bash
-pip install -e .[dev]
+maturin develop -E .[dev]
 ```
 
-Anytime you change the rust code, you can run `pip install -e .` to recompile the rust code.
+Anytime you change the rust code, you can run `maturin develop -E` to recompile the rust code.
 
 ### Running Tests
 

docs/reference/contributing.md

@@ -0,0 +1,10 @@
+## Lambda Functions
+
+Anonymous functions are mapped to uniquely named global functions with a default rewrite of their body.
+
+Constraints:
+
+1. they should be uniquqe based on their types, arguments and body
+2. Their name is a version of that
+3. Their egg value is their name but simplified
+4. Their pretty/serialized version (for graphviz) is that version but without the types.

docs/reference/design.md

@@ -62,7 +62,7 @@ docs = [
     "ablog",
 ]
 
-[tool.ruff]
+[tool.ruff.lint]
 ignore = [
     # Allow uppercase vars
     "N806",
@@ -74,7 +74,7 @@ ignore = [
     # Allow exec
     "S102",
     "S307",
-    "PGH001",
+    "S307",
     # allow star imports
     "F405",
     "F403",
@@ -173,12 +173,16 @@ ignore = [
     "D401",
     # Allow private member refs
     "SLF001",
+    # allow blind exception to add context
+    "BLE001",
 ]
+select = ["ALL"]
+
+[tool.ruff]
 
 line-length = 120
 # Allow lines to be as long as 120.
 src = ["python"]
-select = ["ALL"]
 extend-exclude = ["python/tests/__snapshots__"]
 unsafe-fixes = true
 

pyproject.toml

@@ -14,6 +14,7 @@ class SerializedEGraph:
     def to_dot(self) -> str: ...
     def to_json(self) -> str: ...
     def map_ops(self, map: dict[str, str]) -> None: ...
+    def split_e_classes(self, egraph: EGraph, ops: set[str]) -> None: ...
 
 @final
 class PyObjectSort:
@@ -43,7 +44,6 @@ class EGraph:
         max_functions: int | None = None,
         max_calls_per_function: int | None = None,
         include_temporary_functions: bool = False,
-        split_primitive_outputs: bool = False,
     ) -> SerializedEGraph: ...
     def eval_py_object(self, __expr: _Expr) -> object: ...
     def eval_i64(self, __expr: _Expr) -> int: ...

python/egglog/bindings.pyi

@@ -408,6 +408,12 @@ def __getitem__(self, index: i64Like) -> T: ...
     @method(egg_fn="rebuild")
     def rebuild(self) -> Vec[T]: ...
 
+    @method(egg_fn="vec-remove")
+    def remove(self, index: i64Like) -> Vec[T]: ...
+
+    @method(egg_fn="vec-set")
+    def set(self, index: i64Like, value: T) -> Vec[T]: ...
+
 
 class PyObject(Expr, builtin=True):
     def __init__(self, value: object) -> None: ...

python/egglog/builtins.py

@@ -13,10 +13,11 @@
 from typing_extensions import Self, assert_never
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable
+    from collections.abc import Callable, Iterable, Mapping
 
 
 __all__ = [
+    "replace_typed_expr",
     "Declarations",
     "DeclerationsLike",
     "DelayedDeclerations",
@@ -29,6 +30,7 @@
     "MethodRef",
     "ClassMethodRef",
     "FunctionRef",
+    "UnnamedFunctionRef",
     "ConstantRef",
     "ClassVariableRef",
     "PropertyRef",
@@ -83,17 +85,14 @@ class DelayedDeclerations:
 
     @property
     def __egg_decls__(self) -> Declarations:
+        thunk = self.__egg_decls_thunk__
         try:
-            return self.__egg_decls_thunk__()
+            return thunk()
         # Catch attribute error, so that it isn't bubbled up as a missing attribute and fallbacks on `__getattr__`
         # instead raise explicitly
         except AttributeError as err:
             msg = f"Cannot resolve declerations for {self}"
             raise RuntimeError(msg) from err
-        # Might as well catch others too so we have more context when they raise
-        except Exception as err:  # noqa: BLE001
-            msg = f"Cannot resolve declerations for {self}"
-            raise RuntimeError(msg) from err
 
 
 @runtime_checkable
@@ -121,7 +120,8 @@ def upcast_declerations(declerations_like: Iterable[DeclerationsLike]) -> list[D
 
 @dataclass
 class Declarations:
-    _functions: dict[str, FunctionDecl | RelationDecl] = field(default_factory=dict)
+    # TODO: Replace with set of unnamed function decls
+    _functions: dict[str | UnnamedFunctionRef, FunctionDecl | RelationDecl] = field(default_factory=dict)
     _constants: dict[str, ConstantDecl] = field(default_factory=dict)
     _classes: dict[str, ClassDecl] = field(default_factory=dict)
     _rulesets: dict[str, RulesetDecl | CombinedRulesetDecl] = field(default_factory=lambda: {"": RulesetDecl([])})
@@ -323,9 +323,26 @@ def __str__(self) -> str:
 ##
 
 
+@dataclass(frozen=True)
+class UnnamedFunctionRef:
+    """
+    A reference to a function that doesn't have a name, but does have a body.
+    """
+
+    arg_types: tuple[JustTypeRef, ...]
+    arg_names: tuple[str, ...]
+    res: TypedExprDecl
+
+    @property
+    def args(self) -> tuple[TypedExprDecl, ...]:
+        return tuple(
+            TypedExprDecl(tp, VarDecl(name, False)) for tp, name in zip(self.arg_types, self.arg_names, strict=True)
+        )
+
+
 @dataclass(frozen=True)
 class FunctionRef:
-    name: str
+    name: str | UnnamedFunctionRef
 
 
 @dataclass(frozen=True)
@@ -460,6 +477,8 @@ def to_function_decl(self) -> FunctionDecl:
 @dataclass(frozen=True)
 class VarDecl:
     name: str
+    # Differentiate between let bound vars and vars created in rules so that they won't shadow in egglog, by adding a prefix
+    is_let: bool
 
 
 @dataclass(frozen=True)
@@ -566,6 +585,38 @@ def descendants(self) -> list[TypedExprDecl]:
         return l
 
 
+def replace_typed_expr(typed_expr: TypedExprDecl, replacements: Mapping[TypedExprDecl, TypedExprDecl]) -> TypedExprDecl:
+    """
+    Replace all the typed expressions in the given typed expression with the replacements.
+    """
+    # keep track of the traversed expressions for memoization
+    traversed: dict[TypedExprDecl, TypedExprDecl] = {}
+
+    def _inner(typed_expr: TypedExprDecl) -> TypedExprDecl:
+        if typed_expr in traversed:
+            return traversed[typed_expr]
+        if typed_expr in replacements:
+            res = replacements[typed_expr]
+        else:
+            match typed_expr.expr:
+                case (
+                    CallDecl(callable, args, bound_tp_params)
+                    | PartialCallDecl(CallDecl(callable, args, bound_tp_params))
+                ):
+                    new_args = tuple(_inner(a) for a in args)
+                    call_decl = CallDecl(callable, new_args, bound_tp_params)
+                    res = TypedExprDecl(
+                        typed_expr.tp,
+                        call_decl if isinstance(typed_expr.expr, CallDecl) else PartialCallDecl(call_decl),
+                    )
+                case _:
+                    res = typed_expr
+        traversed[typed_expr] = res
+        return res
+
+    return _inner(typed_expr)
+
+
 ##
 # Schedules
 ##