egglog 7.2.0__cp311-none-win_amd64.whl → 8.0.1__cp311-none-win_amd64.whl

egglog/bindings.pyi

@@ -5,8 +5,6 @@ from typing import TypeAlias
 
 from typing_extensions import final
 
-HIGH_COST: int
-
 @final
 class SerializedEGraph:
     def inline_leaves(self) -> None: ...
@@ -14,12 +12,14 @@ class SerializedEGraph:
     def to_dot(self) -> str: ...
     def to_json(self) -> str: ...
     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 parse_program(__input: str, /, filename: str | None = None) -> list[_Command]: ...
 @final
 class EGraph:
     def __init__(
@@ -28,11 +28,9 @@ class EGraph:
         *,
         fact_directory: str | Path | None = None,
         seminaive: bool = True,
-        terms_encoding: bool = False,
         record: bool = False,
     ) -> None: ...
     def commands(self) -> str | None: ...
-    def parse_program(self, __input: str, /) -> list[_Command]: ...
     def run_program(self, *commands: _Command) -> list[str]: ...
     def extract_report(self) -> _ExtractReport | None: ...
     def run_report(self) -> RunReport | None: ...
@@ -43,7 +41,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 +53,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 +108,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 +161,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 +192,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
 
@@ -220,6 +247,7 @@ _Action: TypeAlias = Let | Set | Change | Union | Panic | Expr_ | Extract
 
 @final
 class FunctionDecl:
+    span: Span
     name: str
     schema: Schema
     default: _Expr | None
@@ -231,6 +259,7 @@ class FunctionDecl:
 
     def __init__(
         self,
+        span: Span,
         name: str,
         schema: Schema,
         default: _Expr | None = None,
@@ -243,7 +272,8 @@ class FunctionDecl:
 
 @final
 class Variant:
-    def __init__(self, name: str, types: list[str], cost: int | None = None) -> None: ...
+    def __init__(self, span: Span, name: str, types: list[str], cost: int | None = None) -> None: ...
+    span: Span
     name: str
     types: list[str]
     cost: int | None
@@ -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,27 +354,48 @@ _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
 
+##
+# Subdatatypes
+##
+
+@final
+class SubVariants:
+    def __init__(self, variants: list[Variant]) -> None: ...
+    variants: list[Variant]
+
+@final
+class NewSort:
+    def __init__(self, name: str, args: list[_Expr]) -> None: ...
+    name: str
+    args: list[_Expr]
+
+_Subdatatypes: TypeAlias = SubVariants | NewSort
+
 ##
 # Commands
 ##
@@ -355,21 +408,23 @@ class SetOption:
 
 @final
 class Datatype:
+    span: Span
     name: str
     variants: list[Variant]
-    def __init__(self, name: str, variants: list[Variant]) -> None: ...
+    def __init__(self, span: Span, name: str, variants: list[Variant]) -> None: ...
 
 @final
-class Declare:
-    name: str
-    sort: str
-    def __init__(self, name: str, sort: str) -> None: ...
+class Datatypes:
+    span: Span
+    datatypes: list[tuple[Span, str, _Subdatatypes]]
+    def __init__(self, span: Span, datatypes: list[tuple[Span, str, _Subdatatypes]]) -> None: ...
 
 @final
 class Sort:
+    span: Span
     name: str
     presort_and_args: tuple[str, list[_Expr]] | None
-    def __init__(self, name: str, presort_and_args: tuple[str, list[_Expr]] | None = None) -> None: ...
+    def __init__(self, span: Span, name: str, presort_and_args: tuple[str, list[_Expr]] | None = None) -> None: ...
 
 @final
 class Function:
@@ -415,49 +470,50 @@ class RunSchedule:
 
 @final
 class Simplify:
+    span: Span
     expr: _Expr
     schedule: _Schedule
-    def __init__(self, expr: _Expr, schedule: _Schedule) -> None: ...
-
-@final
-class Calc:
-    identifiers: list[IdentSort]
-    exprs: list[_Expr]
-    def __init__(self, identifiers: list[IdentSort], exprs: list[_Expr]) -> None: ...
+    def __init__(self, span: Span, expr: _Expr, schedule: _Schedule) -> None: ...
 
 @final
 class QueryExtract:
+    span: Span
     variants: int
     expr: _Expr
-    def __init__(self, variants: int, expr: _Expr) -> None: ...
+    def __init__(self, span: Span, variants: int, expr: _Expr) -> None: ...
 
 @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:
+    span: Span
     name: str | None
-    def __init__(self, name: str | None) -> None: ...
+    def __init__(self, span: Span, name: str | None) -> None: ...
 
 @final
 class Output:
+    span: Span
     file: str
     exprs: list[_Expr]
-    def __init__(self, file: str, exprs: list[_Expr]) -> None: ...
+    def __init__(self, span: Span, file: str, exprs: list[_Expr]) -> None: ...
 
 @final
 class Input:
+    span: Span
     name: str
     file: str
-    def __init__(self, name: str, file: str) -> None: ...
+    def __init__(self, span: Span, name: str, file: str) -> None: ...
 
 @final
 class Push:
@@ -466,29 +522,29 @@ class Push:
 
 @final
 class Pop:
+    span: Span
     length: int
-    def __init__(self, length: int) -> None: ...
+    def __init__(self, span: Span, length: int) -> None: ...
 
 @final
 class Fail:
+    span: Span
     command: _Command
-    def __init__(self, command: _Command) -> None: ...
+    def __init__(self, span: Span, command: _Command) -> None: ...
 
 @final
 class Include:
+    span: Span
     path: str
-    def __init__(self, path: str) -> None: ...
-
-@final
-class CheckProof:
-    def __init__(self) -> None: ...
+    def __init__(self, span: Span, path: str) -> None: ...
 
 @final
 class Relation:
+    span: Span
     constructor: str
     inputs: list[str]
 
-    def __init__(self, constructor: str, inputs: list[str]) -> None: ...
+    def __init__(self, span: Span, constructor: str, inputs: list[str]) -> None: ...
 
 @final
 class PrintOverallStatistics:
@@ -503,7 +559,7 @@ class UnstableCombinedRuleset:
 _Command: TypeAlias = (
     SetOption
     | Datatype
-    | Declare
+    | Datatypes
     | Sort
     | Function
     | AddRuleset
@@ -512,7 +568,6 @@ _Command: TypeAlias = (
     | BiRewriteCommand
     | ActionCommand
     | RunSchedule
-    | Calc
     | Simplify
     | QueryExtract
     | Check
@@ -524,7 +579,6 @@ _Command: TypeAlias = (
     | Pop
     | Fail
     | Include
-    | CheckProof
     | Relation
     | PrintOverallStatistics
     | UnstableCombinedRuleset

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:
@@ -172,6 +195,6 @@ def _get_tp(x: object) -> TypeName | type:
         return TypeName(x.__egg_typed_expr__.tp.name)
     tp = type(x)
     # If this value has a custom metaclass, let's use that as our index instead of the type
-    if type(tp) != type:
+    if type(tp) is not type:
         return type(tp)
     return tp