Merge pull request #23 from metadsl/upgrade

Upgrade egglog dependency

Author: saulshanabrook

Cargo.lock

@@ -11,7 +11,7 @@ crate-type = ["cdylib"]
 [dependencies]
 # Use unreleased version to depend on signature improvements https://github.com/PyO3/pyo3/pull/2702
 pyo3 = { version = "0.18.1", features = ["extension-module"] }
-egg-smol = { git = "https://github.com/egraphs-good/egglog", rev = "30feaaab88452ec4b6c5f7a199345298bac2dd0f" }
+egglog = { git = "https://github.com/egraphs-good/egglog", rev = "39b199d9bfce9cc47d0c54977279c5b04231e717" }
 # egg-smol = { path = "../egglog" }
 pyo3-log = "0.8.1"
 log = "0.4.17"

Cargo.toml

@@ -5,6 +5,10 @@
 - Renamed `config()` to `run()` to better match `egglog` command
 - Fixed `relation` type signature
 - Added default limit of 1 to `run()` to match `egglog` command and moved to second arg
+- Upgraded `egglog` dependency ([changes](https://github.com/egraphs-good/egglog/compare/30feaaab88452ec4b6c5f7a199345298bac2dd0f...39b199d9bfce9cc47d0c54977279c5b04231e717))
+  - Added `Set` sort and removed set method from `Map`
+  - Added `Vec` sort
+  - Added support for variable args for builtin functions, to use in creation of `Vec` and `Set` sorts.
 
 ## 0.4.0 (2023-05-03)
 

docs/changelog.md

@@ -95,6 +95,13 @@ class Set:
     args: list[_Expr]
     rhs: _Expr
 
+@final
+class SetNoTrack:
+    def __init__(self, lhs: str, args: list[_Expr], rhs: _Expr) -> None: ...
+    lhs: str
+    args: list[_Expr]
+    rhs: _Expr
+
 @final
 class Delete:
     sym: str
@@ -117,7 +124,7 @@ class Expr_:
     def __init__(self, expr: _Expr) -> None: ...
     expr: _Expr
 
-_Action = Let | Set | Delete | Union | Panic | Expr_
+_Action = Let | Set | SetNoTrack | Delete | Union | Panic | Expr_
 
 ##
 # Other Structs

python/egglog/bindings.pyi

@@ -19,6 +19,8 @@
     "StringLike",
     "Map",
     "Rational",
+    "Set",
+    "Vec",
 ]
 
 
@@ -159,41 +161,69 @@ def __init__(self, value: str):
 
 @BUILTINS.class_(egg_sort="Map")
 class Map(BaseExpr, Generic[T, V]):
-    @BUILTINS.method(egg_fn="empty")
+    @BUILTINS.method(egg_fn="map-empty")
     @classmethod
     def empty(cls) -> Map[T, V]:  # type: ignore[empty-body]
         ...
 
-    @BUILTINS.method(egg_fn="insert")
+    @BUILTINS.method(egg_fn="map-insert")
     def insert(self, key: T, value: V) -> Map[T, V]:  # type: ignore[empty-body]
         ...
 
-    @BUILTINS.method(egg_fn="get")
+    @BUILTINS.method(egg_fn="map-get")
     def __getitem__(self, key: T) -> V:  # type: ignore[empty-body]
         ...
 
-    @BUILTINS.method(egg_fn="not-contains")
+    @BUILTINS.method(egg_fn="map-not-contains")
     def not_contains(self, key: T) -> Unit:  # type: ignore[empty-body]
         ...
 
-    @BUILTINS.method(egg_fn="contains")
+    @BUILTINS.method(egg_fn="map-contains")
     def contains(self, key: T) -> Unit:  # type: ignore[empty-body]
         ...
 
+    @BUILTINS.method(egg_fn="map-remove")
+    def remove(self, key: T) -> Map[T, V]:  # type: ignore[empty-body]
+        ...
+
+
+@BUILTINS.class_(egg_sort="Set")
+class Set(BaseExpr, Generic[T]):
+    @BUILTINS.method(egg_fn="set-of")
+    def __init__(self, *args: T) -> None:
+        ...
+
+    @BUILTINS.method(egg_fn="set-empty")
+    @classmethod
+    def empty(cls) -> Set[T]:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="set-insert")
+    def insert(self, value: T) -> Set[T]:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="set-not-contains")
+    def not_contains(self, value: T) -> Unit:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="set-contains")
+    def contains(self, value: T) -> Unit:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="set-remove")
+    def remove(self, value: T) -> Set[T]:  # type: ignore[empty-body]
+        ...
+
     @BUILTINS.method(egg_fn="set-union")
-    def __or__(self, __t: Map[T, V]) -> Map[T, V]:  # type: ignore[empty-body]
+    def __or__(self, other: Set[T]) -> Set[T]:  # type: ignore[empty-body]
         ...
 
     @BUILTINS.method(egg_fn="set-diff")
-    def __sub__(self, __t: Map[T, V]) -> Map[T, V]:  # type: ignore[empty-body]
+    def __sub__(self, other: Set[T]) -> Set[T]:  # type: ignore[empty-body]
         ...
 
     @BUILTINS.method(egg_fn="set-intersect")
-    def __and__(self, __t: Map[T, V]) -> Map[T, V]:  # type: ignore[empty-body]
-        ...
-
-    @BUILTINS.method(egg_fn="map-remove")
-    def remove(self, key: T) -> Map[T, V]:  # type: ignore[empty-body]
+    def __and__(self, other: Set[T]) -> Set[T]:  # type: ignore[empty-body]
         ...
 
 
@@ -203,6 +233,10 @@ class Rational(BaseExpr):
     def __init__(self, num: i64Like, den: i64Like):
         ...
 
+    @BUILTINS.method(egg_fn="to-f64")
+    def to_f64(self) -> f64:  # type: ignore[empty-body]
+        ...
+
     @BUILTINS.method(egg_fn="+")
     def __add__(self, other: Rational) -> Rational:  # type: ignore[empty-body]
         ...
@@ -262,3 +296,35 @@ def sqrt(self) -> Rational:  # type: ignore[empty-body]
     @BUILTINS.method(egg_fn="cbrt")
     def cbrt(self) -> Rational:  # type: ignore[empty-body]
         ...
+
+
+@BUILTINS.class_(egg_sort="Vec")
+class Vec(BaseExpr, Generic[T]):
+    @BUILTINS.method(egg_fn="vec-of")
+    def __init__(self, *args: T) -> None:
+        ...
+
+    @BUILTINS.method(egg_fn="vec-empty")
+    @classmethod
+    def empty(cls) -> Vec[T]:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="vec-append")
+    def append(self, *others: Vec[T]) -> Vec[T]:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="vec-push")
+    def push(self, value: T) -> Vec[T]:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="vec-pop")
+    def pop(self) -> Vec[T]:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="vec-not-contains")
+    def not_contains(self, value: T) -> Unit:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="vec-contains")
+    def contains(self, value: T) -> Unit:  # type: ignore[empty-body]
+        ...

python/egglog/builtins.py

@@ -311,12 +311,15 @@ class FunctionDecl:
     # TODO: Add arg name to arg so can call with keyword arg
     arg_types: tuple[TypeOrVarRef, ...]
     return_type: TypeOrVarRef
+    var_arg_type: Optional[TypeOrVarRef] = None
     cost: Optional[int] = None
     default: Optional[ExprDecl] = None
     merge: Optional[ExprDecl] = None
     merge_action: tuple[ActionDecl, ...] = ()
 
     def to_commands(self, decls: Declarations, egg_name: str) -> Iterable[bindings._Command]:
+        if self.var_arg_type is not None:
+            raise NotImplementedError("egglog does not support variable arguments yet.")
         just_arg_types = [a.to_just() for a in self.arg_types]
         for a in just_arg_types:
             yield from decls._register_sort(a)
@@ -406,6 +409,7 @@ class CallDecl:
     callable: CallableRef
     args: tuple[ExprDecl, ...] = ()
     # type parameters that were bound to the callable, if it is a classmethod
+    # Used for pretty printing classmethod calls with type parameters
     bound_tp_params: Optional[tuple[JustTypeRef, ...]] = None
 
     def __post_init__(self):
@@ -424,7 +428,7 @@ def from_egg(cls, decls: Declarations, call: bindings.Call) -> tuple[JustTypeRef
         for callable_ref in decls._egg_fn_to_callable_refs[call.name]:
             tcs = TypeConstraintSolver()
             fn_decl = decls._get_function_decl(callable_ref)
-            return_tp = tcs.infer_return_type(fn_decl.arg_types, fn_decl.return_type, arg_types)
+            return_tp = tcs.infer_return_type(fn_decl.arg_types, fn_decl.return_type, fn_decl.var_arg_type, arg_types)
             return return_tp, cls(callable_ref, arg_decls)
         raise ValueError(f"Could not find callable ref for call {call}")
 

python/egglog/declarations.py

@@ -561,10 +561,22 @@ def _generate_function_decl(
             if first.annotation != Parameter.empty:
                 raise ValueError(f"First arg of a method must not have an annotation, not {first.annotation}")
 
+        # Check that all the params are positional or keyword, and that there is only one var arg at the end
+        found_var_arg = False
         for param in params:
-            if param.kind != Parameter.POSITIONAL_OR_KEYWORD:
+            if found_var_arg:
+                raise ValueError("Can only have a single var arg at the end")
+            kind = param.kind
+            if kind == Parameter.VAR_POSITIONAL:
+                found_var_arg = True
+            elif kind != Parameter.POSITIONAL_OR_KEYWORD:
                 raise ValueError(f"Can only register functions with positional or keyword args, not {param.kind}")
 
+        if found_var_arg:
+            var_arg_param, *params = params
+            var_arg_type = self._resolve_type_annotation(hints[var_arg_param.name], cls_typevars, cls_type_and_name)
+        else:
+            var_arg_type = None
         arg_types = tuple(self._resolve_type_annotation(hints[t.name], cls_typevars, cls_type_and_name) for t in params)
         # If the first arg is a self, and this not an __init__ fn, add this as a typeref
         if isinstance(first_arg, (ClassTypeVarRef, TypeRefWithVars)) and not is_init:
@@ -594,6 +606,7 @@ def _generate_function_decl(
         )
         decl = FunctionDecl(
             return_type=return_type,
+            var_arg_type=var_arg_type,
             arg_types=arg_types,
             cost=cost,
             default=default_decl,

python/egglog/egraph.py

@@ -79,7 +79,7 @@ def if_(c: Term, t: Term, f: Term) -> Term:
     ...
 
 
-StringSet = Map[Var, i64]
+StringSet = Set[Var]
 
 
 @egraph.function(merge=lambda old, new: old & new)
@@ -95,7 +95,7 @@ def freer(t: Term) -> StringSet:
 egraph.register(
     # freer
     rule(eq(t).to(Term.val(v))).then(set_(freer(t)).to(StringSet.empty())),
-    rule(eq(t).to(Term.var(x))).then(set_(freer(t)).to(StringSet.empty().insert(x, i64(1)))),
+    rule(eq(t).to(Term.var(x))).then(set_(freer(t)).to(StringSet.empty().insert(x))),
     rule(eq(t).to(t1 + t2), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),
     rule(eq(t).to(t1 == t2), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),
     rule(eq(t).to(t1(t2)), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),

python/egglog/examples/lambda.py

@@ -67,7 +67,9 @@ def __dir__(self) -> list[str]:
             possible_methods.append("__call__")
         return possible_methods
 
-    def __getitem__(self, args: tuple[RuntimeTypeArgType, ...]) -> RuntimeParamaterizedClass:
+    def __getitem__(self, args: tuple[RuntimeTypeArgType, ...] | RuntimeTypeArgType) -> RuntimeParamaterizedClass:
+        if not isinstance(args, tuple):
+            args = (args,)
         tp = JustTypeRef(self.__egg_name__, tuple(class_to_ref(arg) for arg in args))
         return RuntimeParamaterizedClass(self.__egg_decls__, tp)
 
@@ -160,7 +162,7 @@ def _call(
         tcs = TypeConstraintSolver()
 
     if fn_decl is not None:
-        return_tp = tcs.infer_return_type(fn_decl.arg_types, fn_decl.return_type, arg_types)
+        return_tp = tcs.infer_return_type(fn_decl.arg_types, fn_decl.return_type, fn_decl.var_arg_type, arg_types)
     else:
         return_tp = JustTypeRef("unit")
 

python/egglog/runtime.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Collection
+from itertools import chain, repeat
+from typing import Collection, Optional
 
 from .declarations import *
 
@@ -38,24 +39,28 @@ def infer_return_type(
         self,
         fn_args: Collection[TypeOrVarRef],
         fn_return: TypeOrVarRef,
+        fn_var_args: Optional[TypeOrVarRef],
         args: Collection[JustTypeRef],
     ) -> JustTypeRef:
         # Infer the type of each type variable based on the actual types of the arguments
-        self._infer_typevars_zip(fn_args, args)
+        self._infer_typevars_zip(fn_args, fn_var_args, args)
         # Substitute the type variables with their inferred types
         return self._subtitute_typevars(fn_return)
 
-    def _infer_typevars_zip(self, fn_args: Collection[TypeOrVarRef], args: Collection[JustTypeRef]) -> None:
-        if len(fn_args) != len(args):
+    def _infer_typevars_zip(
+        self, fn_args: Collection[TypeOrVarRef], fn_var_args: Optional[TypeOrVarRef], args: Collection[JustTypeRef]
+    ) -> None:
+        if len(fn_args) != len(args) if fn_var_args is None else len(fn_args) > len(args):
             raise TypeConstraintError(f"Expected {len(fn_args)} args, got {len(args)}")
-        for fn_arg, arg in zip(fn_args, args):
+        all_fn_args = fn_args if fn_var_args is None else chain(fn_args, repeat(fn_var_args))
+        for fn_arg, arg in zip(all_fn_args, args):
             self._infer_typevars(fn_arg, arg)
 
     def _infer_typevars(self, fn_arg: TypeOrVarRef, arg: JustTypeRef) -> None:
         if isinstance(fn_arg, TypeRefWithVars):
             if fn_arg.name != arg.name:
                 raise TypeConstraintError(f"Expected {fn_arg.name}, got {arg.name}")
-            self._infer_typevars_zip(fn_arg.args, arg.args)
+            self._infer_typevars_zip(fn_arg.args, None, arg.args)
         elif fn_arg.index not in self._cls_typevar_index_to_type:
             self._cls_typevar_index_to_type[fn_arg.index] = arg
         elif self._cls_typevar_index_to_type[fn_arg.index] != arg: