egglog 7.1.0__cp312-none-win_amd64.whl → 8.0.0__cp312-none-win_amd64.whl

egglog/bindings.pyi

@@ -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:
@@ -32,7 +33,7 @@ class EGraph:
         record: bool = False,
     ) -> None: ...
     def commands(self) -> str | None: ...
-    def parse_program(self, __input: str, /) -> list[_Command]: ...
+    def parse_program(self, __input: str, /, filename: str | None = None) -> list[_Command]: ...
     def run_program(self, *commands: _Command) -> list[str]: ...
     def extract_report(self) -> _ExtractReport | None: ...
     def run_report(self) -> RunReport | None: ...
@@ -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: ...
@@ -56,6 +56,25 @@ class EGraph:
 class EggSmolError(Exception):
     context: str
 
+##
+# Spans
+##
+
+@final
+class SrcFile:
+    def __init__(self, name: str, contents: str | None = None) -> None: ...
+    name: str
+    contents: str | None
+
+@final
+class Span:
+    def __init__(self, file: SrcFile, start: int, end: int) -> None: ...
+    file: SrcFile
+    start: int
+    end: int
+
+DUMMY_SPAN: Span = ...
+
 ##
 # Literals
 ##
@@ -92,17 +111,20 @@ _Literal: TypeAlias = Int | F64 | String | Bool | Unit
 
 @final
 class Lit:
-    def __init__(self, value: _Literal) -> None: ...
+    def __init__(self, span: Span, value: _Literal) -> None: ...
+    span: Span
     value: _Literal
 
 @final
 class Var:
-    def __init__(self, name: str) -> None: ...
+    def __init__(self, span: Span, name: str) -> None: ...
+    span: Span
     name: str
 
 @final
 class Call:
-    def __init__(self, name: str, args: list[_Expr]) -> None: ...
+    def __init__(self, span: Span, name: str, args: list[_Expr]) -> None: ...
+    span: Span
     name: str
     args: list[_Expr]
 
@@ -142,7 +164,8 @@ class TermDag:
 
 @final
 class Eq:
-    def __init__(self, exprs: list[_Expr]) -> None: ...
+    def __init__(self, span: Span, exprs: list[_Expr]) -> None: ...
+    span: Span
     exprs: list[_Expr]
 
 @final
@@ -172,43 +195,50 @@ _Change: TypeAlias = Delete | Subsume
 
 @final
 class Let:
-    def __init__(self, lhs: str, rhs: _Expr) -> None: ...
+    def __init__(self, span: Span, lhs: str, rhs: _Expr) -> None: ...
+    span: Span
     lhs: str
     rhs: _Expr
 
 @final
 class Set:
-    def __init__(self, lhs: str, args: list[_Expr], rhs: _Expr) -> None: ...
+    def __init__(self, span: Span, lhs: str, args: list[_Expr], rhs: _Expr) -> None: ...
+    span: Span
     lhs: str
     args: list[_Expr]
     rhs: _Expr
 
 @final
 class Change:
+    span: Span
     change: _Change
     sym: str
     args: list[_Expr]
-    def __init__(self, change: _Change, sym: str, args: list[_Expr]) -> None: ...
+    def __init__(self, span: Span, change: _Change, sym: str, args: list[_Expr]) -> None: ...
 
 @final
 class Union:
-    def __init__(self, lhs: _Expr, rhs: _Expr) -> None: ...
+    def __init__(self, span: Span, lhs: _Expr, rhs: _Expr) -> None: ...
+    span: Span
     lhs: _Expr
     rhs: _Expr
 
 @final
 class Panic:
-    def __init__(self, msg: str) -> None: ...
+    def __init__(self, span: Span, msg: str) -> None: ...
+    span: Span
     msg: str
 
 @final
 class Expr_:  # noqa: N801
-    def __init__(self, expr: _Expr) -> None: ...
+    def __init__(self, span: Span, expr: _Expr) -> None: ...
+    span: Span
     expr: _Expr
 
 @final
 class Extract:
-    def __init__(self, expr: _Expr, variants: _Expr) -> None: ...
+    def __init__(self, span: Span, expr: _Expr, variants: _Expr) -> None: ...
+    span: Span
     expr: _Expr
     variants: _Expr
 
@@ -256,17 +286,19 @@ class Schema:
 
 @final
 class Rule:
+    span: Span
     head: list[_Action]
     body: list[_Fact]
-    def __init__(self, head: list[_Action], body: list[_Fact]) -> None: ...
+    def __init__(self, span: Span, head: list[_Action], body: list[_Fact]) -> None: ...
 
 @final
 class Rewrite:
+    span: Span
     lhs: _Expr
     rhs: _Expr
     conditions: list[_Fact]
 
-    def __init__(self, lhs: _Expr, rhs: _Expr, conditions: list[_Fact] = []) -> None: ...  # noqa: B006
+    def __init__(self, span: Span, lhs: _Expr, rhs: _Expr, conditions: list[_Fact] = []) -> None: ...  # noqa: B006
 
 @final
 class RunConfig:
@@ -322,24 +354,28 @@ _ExtractReport: TypeAlias = Variants | Best
 
 @final
 class Saturate:
+    span: Span
     schedule: _Schedule
-    def __init__(self, schedule: _Schedule) -> None: ...
+    def __init__(self, span: Span, schedule: _Schedule) -> None: ...
 
 @final
 class Repeat:
+    span: Span
     length: int
     schedule: _Schedule
-    def __init__(self, length: int, schedule: _Schedule) -> None: ...
+    def __init__(self, span: Span, length: int, schedule: _Schedule) -> None: ...
 
 @final
 class Run:
+    span: Span
     config: RunConfig
-    def __init__(self, config: RunConfig) -> None: ...
+    def __init__(self, span: Span, config: RunConfig) -> None: ...
 
 @final
 class Sequence:
+    span: Span
     schedules: list[_Schedule]
-    def __init__(self, schedules: list[_Schedule]) -> None: ...
+    def __init__(self, span: Span, schedules: list[_Schedule]) -> None: ...
 
 _Schedule: TypeAlias = Saturate | Repeat | Run | Sequence
 
@@ -361,9 +397,10 @@ class Datatype:
 
 @final
 class Declare:
+    span: Span
     name: str
     sort: str
-    def __init__(self, name: str, sort: str) -> None: ...
+    def __init__(self, span: Span, name: str, sort: str) -> None: ...
 
 @final
 class Sort:
@@ -421,9 +458,10 @@ class Simplify:
 
 @final
 class Calc:
+    span: Span
     identifiers: list[IdentSort]
     exprs: list[_Expr]
-    def __init__(self, identifiers: list[IdentSort], exprs: list[_Expr]) -> None: ...
+    def __init__(self, span: Span, identifiers: list[IdentSort], exprs: list[_Expr]) -> None: ...
 
 @final
 class QueryExtract:
@@ -433,14 +471,16 @@ class QueryExtract:
 
 @final
 class Check:
+    span: Span
     facts: list[_Fact]
-    def __init__(self, facts: list[_Fact]) -> None: ...
+    def __init__(self, span: Span, facts: list[_Fact]) -> None: ...
 
 @final
 class PrintFunction:
+    span: Span
     name: str
     length: int
-    def __init__(self, name: str, length: int) -> None: ...
+    def __init__(self, span: Span, name: str, length: int) -> None: ...
 
 @final
 class PrintSize:

egglog/builtins.py

@@ -6,17 +6,21 @@ Builtin sorts and function to egg.
 from __future__ import annotations
 
 from functools import partial
-from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union, overload
+from types import FunctionType
+from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union, cast, overload
 
 from typing_extensions import TypeVarTuple, Unpack
 
-from .conversion import converter
-from .egraph import Expr, Unit, function, method
-from .runtime import RuntimeFunction
+from .conversion import converter, get_type_args
+from .egraph import Expr, Unit, function, get_current_ruleset, method
+from .functionalize import functionalize
+from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction
+from .thunk import Thunk
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+
 __all__ = [
     "i64",
     "i64Like",
@@ -80,7 +84,7 @@ class Bool(Expr, egg_sort="bool", builtin=True):
 converter(bool, Bool, Bool)
 
 # The types which can be convertered into an i64
-i64Like = Union["i64", int]  # noqa: N816
+i64Like: TypeAlias = Union["i64", int]  # noqa: N816, PYI042
 
 
 class i64(Expr, builtin=True):  # noqa: N801
@@ -182,7 +186,7 @@ converter(int, i64, i64)
 def count_matches(s: StringLike, pattern: StringLike) -> i64: ...
 
 
-f64Like = Union["f64", float]  # noqa: N816
+f64Like: TypeAlias = Union["f64", float]  # noqa: N816, PYI042
 
 
 class f64(Expr, builtin=True):  # noqa: N801
@@ -404,6 +408,12 @@ class Vec(Expr, Generic[T], builtin=True):
     @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: ...
@@ -501,3 +511,36 @@ class UnstableFn(Expr, Generic[T, Unpack[TS]], builtin=True):
 
 converter(RuntimeFunction, UnstableFn, UnstableFn)
 converter(partial, UnstableFn, lambda p: UnstableFn(p.func, *p.args))
+
+
+def _convert_function(a: FunctionType) -> UnstableFn:
+    """
+    Converts a function type to an unstable function
+
+    Would just be UnstableFn(function(a)) but we have to look for any nonlocals and globals
+    which are runtime expressions with `var`s in them and add them as args to the function
+    """
+    # Update annotations of a to be the type we are trying to convert to
+    return_tp, *arg_tps = get_type_args()
+    a.__annotations__ = {
+        "return": return_tp,
+        # The first varnames should always be the arg names
+        **dict(zip(a.__code__.co_varnames, arg_tps, strict=False)),
+    }
+    # Modify name to make it unique
+    # a.__name__ = f"{a.__name__} {hash(a.__code__)}"
+    transformed_fn = functionalize(a, value_to_annotation)
+    assert isinstance(transformed_fn, partial)
+    return UnstableFn(
+        function(ruleset=get_current_ruleset(), use_body_as_name=True)(transformed_fn.func), *transformed_fn.args
+    )
+
+
+def value_to_annotation(a: object) -> type | None:
+    # only lift runtime expressions (which could contain vars) not any other nonlocals/globals we use in the function
+    if not isinstance(a, RuntimeExpr):
+        return None
+    return cast(type, RuntimeClass(Thunk.value(a.__egg_decls__), a.__egg_typed_expr__.tp.to_var()))
+
+
+converter(FunctionType, UnstableFn, _convert_function)

egglog/conversion.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
+from contextvars import ContextVar
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, NewType, TypeVar, cast
 
@@ -9,9 +11,8 @@ from .runtime import *
 from .thunk import *
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Callable, Generator
 
-    from .declarations import HasDeclerations
     from .egraph import Expr
 
 __all__ = ["convert", "converter", "resolve_literal", "convert_to_same_type"]
@@ -84,7 +85,7 @@ def convert(source: object, target: type[V]) -> V:
     Convert a source object to a target type.
     """
     assert isinstance(target, RuntimeClass)
-    return cast(V, resolve_literal(target.__egg_tp__, source))
+    return cast(V, resolve_literal(target.__egg_tp__, source, target.__egg_decls_thunk__))
 
 
 def convert_to_same_type(source: object, target: RuntimeExpr) -> RuntimeExpr:
@@ -92,7 +93,7 @@ def convert_to_same_type(source: object, target: RuntimeExpr) -> RuntimeExpr:
     Convert a source object to the same type as the target.
     """
     tp = target.__egg_typed_expr__.tp
-    return resolve_literal(tp.to_var(), source)
+    return resolve_literal(tp.to_var(), source, Thunk.value(target.__egg_decls__))
 
 
 def process_tp(tp: type | RuntimeClass) -> TypeName | type:
@@ -140,7 +141,28 @@ def identity(x: object) -> object:
     return x
 
 
-def resolve_literal(tp: TypeOrVarRef, arg: object) -> RuntimeExpr:
+TYPE_ARGS = ContextVar[tuple[RuntimeClass, ...]]("TYPE_ARGS")
+
+
+def get_type_args() -> tuple[RuntimeClass, ...]:
+    """
+    Get the type args for the type being converted.
+    """
+    return TYPE_ARGS.get()
+
+
+@contextmanager
+def with_type_args(args: tuple[JustTypeRef, ...], decls: Callable[[], Declarations]) -> Generator[None, None, None]:
+    token = TYPE_ARGS.set(tuple(RuntimeClass(decls, a.to_var()) for a in args))
+    try:
+        yield
+    finally:
+        TYPE_ARGS.reset(token)
+
+
+def resolve_literal(
+    tp: TypeOrVarRef, arg: object, decls: Callable[[], Declarations] = CONVERSIONS_DECLS
+) -> RuntimeExpr:
     arg_type = _get_tp(arg)
 
     # If we have any type variables, dont bother trying to resolve the literal, just return the arg
@@ -148,7 +170,7 @@ def resolve_literal(tp: TypeOrVarRef, arg: object) -> RuntimeExpr:
         tp_just = tp.to_just()
     except NotImplementedError:
         # If this is a var, it has to be a runtime expession
-        assert isinstance(arg, RuntimeExpr)
+        assert isinstance(arg, RuntimeExpr), f"Expected a runtime expression, got {arg}"
         return arg
     tp_name = TypeName(tp_just.name)
     if arg_type == tp_name:
@@ -158,13 +180,14 @@ def resolve_literal(tp: TypeOrVarRef, arg: object) -> RuntimeExpr:
     # Try all parent types as well, if we are converting from a Python type
     for arg_type_instance in arg_type.__mro__ if isinstance(arg_type, type) else [arg_type]:
         try:
-            fn = CONVERSIONS[(cast(TypeName | type, arg_type_instance), tp_name)][1]
+            fn = CONVERSIONS[(arg_type_instance, tp_name)][1]
         except KeyError:
             continue
         break
     else:
         raise ConvertError(f"Cannot convert {arg_type} to {tp_name}")
-    return fn(arg)
+    with with_type_args(tp_just.args, decls):
+        return fn(arg)
 
 
 def _get_tp(x: object) -> TypeName | type: