egglog 9.0.0__cp313-cp313-win_amd64.whl

egglog/exp/program_gen.py

@@ -0,0 +1,424 @@
+# mypy: disable-error-code="empty-body"
+"""
+Builds up imperative string expressions from a functional expression.
+"""
+
+from __future__ import annotations
+
+from typing import TypeAlias, Union
+
+from egglog import *
+
+ProgramLike: TypeAlias = Union["Program", StringLike]
+
+
+class Program(Expr):
+    """
+    Semanticallly represents an expression with a number of ordered statements that it depends on to run.
+
+    The expression and statements are all represented as strings.
+    """
+
+    def __init__(self, expr: StringLike, is_identifier: BoolLike = Bool(False)) -> None:
+        """
+        Create a program based on a string expression.
+        """
+
+    def __add__(self, other: ProgramLike) -> Program:
+        """
+        Concats the strings of the two expressions and also the statements.
+        """
+
+    @method(unextractable=True)
+    def statement(self, statement: ProgramLike) -> Program:
+        """
+        Uses the expression of the statement and adds that as a statement to the program.
+        """
+
+    def assign(self) -> Program:
+        """
+        Returns a new program with the expression assigned to a gensym.
+        """
+
+    def function_two(self, arg1: ProgramLike, arg2: ProgramLike, name: StringLike = String("__fn")) -> Program:
+        """
+        Returns a new program defining a function with two arguments.
+        """
+
+    def function_three(
+        self, arg1: ProgramLike, arg2: ProgramLike, arg3: ProgramLike, name: StringLike = String("__fn")
+    ) -> Program:
+        """
+        Returns a new program defining a function with three arguments.
+        """
+
+    def expr_to_statement(self) -> Program:
+        """
+        Returns a new program with the expression as a statement and the new expression empty.
+        """
+
+    @property
+    def expr(self) -> String:
+        """
+        Returns the expression of the program, if it's been compiled
+        """
+
+    @property
+    def statements(self) -> String:
+        """
+        Returns the statements of the program, if it's been compiled
+        """
+
+    @property
+    def next_sym(self) -> i64:
+        """
+        Returns the next gensym to use. This is set after calling `compile(i)` on a program.
+        """
+
+    # TODO: Replace w/ def next_sym(self) -> i64: ... ?
+    def compile(self, next_sym: i64 = i64(0)) -> Unit:
+        """
+        Triggers compilation of the program.
+        """
+
+    @method(merge=lambda old, _new: old)  # type: ignore[misc]
+    @property
+    def parent(self) -> Program:
+        """
+        Returns the parent of the program, if it's been compiled into the parent.
+
+        Only keeps the original parent, not any additional ones, so that each set of statements is only added once.
+        """
+
+    @property
+    def is_identifer(self) -> Bool:
+        """
+        Returns whether the expression is an identifier. Used so that we don't re-assign any identifiers.
+        """
+
+
+converter(String, Program, Program)
+
+
+class EvalProgram(Expr):
+    def __init__(self, program: Program, globals: object) -> None:
+        """
+        Evaluates the program and saves as the py_object
+        """
+
+    # Only allow it to be set once, b/c hash of functions not stable
+    @method(merge=lambda old, _new: old)  # type: ignore[misc]
+    @property
+    def as_py_object(self) -> PyObject:
+        """
+        Returns the python object of the program, if it's been evaluated.
+        """
+
+
+@ruleset
+def eval_program_rulseset(ep: EvalProgram, p: Program, expr: String, statements: String, g: PyObject):
+    # When we evaluate a program, we first want to compile to a string
+    yield rule(EvalProgram(p, g)).then(p.compile())
+    # Then we want to evaluate the statements/expr
+    yield rule(
+        eq(ep).to(EvalProgram(p, g)),
+        eq(p.statements).to(statements),
+        eq(p.expr).to(expr),
+    ).then(
+        set_(ep.as_py_object).to(
+            py_eval(
+                "l['___res']",
+                PyObject.dict(PyObject.from_string("l"), py_exec(join(statements, "\n", "___res = ", expr), g)),
+            )
+        ),
+    )
+
+
+@ruleset
+def program_gen_ruleset(
+    s: String,
+    s1: String,
+    s2: String,
+    s3: String,
+    s4: String,
+    s5: String,
+    s6: String,
+    p: Program,
+    p1: Program,
+    p2: Program,
+    p3: Program,
+    p4: Program,
+    i: i64,
+    i2: i64,
+    b: Bool,
+):
+    ##
+    # Expression
+    ##
+
+    # Compiling a string just gives that string
+    yield rule(eq(p).to(Program(s, b)), p.compile(i)).then(
+        set_(p.expr).to(s),
+        set_(p.statements).to(String("")),
+        set_(p.next_sym).to(i),
+        set_(p.is_identifer).to(b),
+    )
+
+    ##
+    # Statement
+    ##
+    # Compiling a statement means that we should use the expression of the statement as a statement and use the expression of the first
+    yield rewrite(p1.statement(p2)).to(p1 + p2.expr_to_statement())
+
+    ##
+    # Expr to statement
+    ##
+    stmt = eq(p).to(p1.expr_to_statement())
+    # 1. Set parent and is_identifier to false, since its empty
+    yield rule(stmt, p.compile(i)).then(set_(p1.parent).to(p), set_(p.is_identifer).to(Bool(False)))
+    # 2. Compile p1 if parent set
+    yield rule(stmt, p.compile(i), eq(p1.parent).to(p)).then(p1.compile(i))
+    # 3.a. If parent not set, set statements to expr
+    yield rule(
+        stmt,
+        p.compile(i),
+        ne(p1.parent).to(p),
+        eq(s1).to(p1.expr),
+    ).then(
+        set_(p.statements).to(join(s1, "\n")),
+        set_(p.next_sym).to(i),
+        set_(p.expr).to(String("")),
+    )
+    # 3.b. If parent set, set statements to expr + statements
+    yield rule(
+        stmt,
+        eq(p1.parent).to(p),
+        eq(s1).to(p1.expr),
+        eq(s2).to(p1.statements),
+        eq(i).to(p1.next_sym),
+    ).then(
+        set_(p.statements).to(join(s2, s1, "\n")),
+        set_(p.next_sym).to(i),
+        set_(p.expr).to(String("")),
+    )
+
+    ##
+    # Addition
+    ##
+
+    # Compiling an addition is the same as compiling one, then the other, then setting the expression as the addition
+    # of the two
+    program_add = eq(p).to(p1 + p2)
+
+    # If the resulting expression is either of the inputs, then its an identifer if those are
+    # Otherwise, if its not equal to either input, its not an identifier
+    yield rule(program_add, eq(p.expr).to(p1.expr), eq(b).to(p1.is_identifer)).then(set_(p.is_identifer).to(b))
+    yield rule(program_add, eq(p.expr).to(p2.expr), eq(b).to(p2.is_identifer)).then(set_(p.is_identifer).to(b))
+    yield rule(program_add, ne(p.expr).to(p1.expr), ne(p.expr).to(p2.expr)).then(set_(p.is_identifer).to(Bool(False)))
+
+    # Set parent of p1
+    yield rule(program_add, p.compile(i)).then(
+        set_(p1.parent).to(p),
+    )
+
+    # Compile p1, if p1 parent equal
+    yield rule(program_add, p.compile(i), eq(p1.parent).to(p)).then(p1.compile(i))
+
+    # Set parent of p2, once p1 compiled
+    yield rule(program_add, p.compile(i), p1.next_sym).then(set_(p2.parent).to(p))
+
+    # Compile p2, if p1 parent not equal, but p2 parent equal
+    yield rule(program_add, p.compile(i), ne(p1.parent).to(p), eq(p2.parent).to(p)).then(p2.compile(i))
+
+    # Compile p2, if p1 parent eqal
+    yield rule(program_add, p.compile(i2), eq(p1.parent).to(p), eq(i).to(p1.next_sym), eq(p2.parent).to(p)).then(
+        p2.compile(i)
+    )
+
+    # Set p expr to join of p1 and p2
+    yield rule(
+        program_add,
+        eq(s1).to(p1.expr),
+        eq(s2).to(p2.expr),
+    ).then(
+        set_(p.expr).to(join(s1, s2)),
+    )
+    # Set p statements to join and next sym to p2 if both parents set
+    yield rule(
+        program_add,
+        eq(p1.parent).to(p),
+        eq(p2.parent).to(p),
+        eq(s1).to(p1.statements),
+        eq(s2).to(p2.statements),
+        eq(i).to(p2.next_sym),
+    ).then(
+        set_(p.statements).to(join(s1, s2)),
+        set_(p.next_sym).to(i),
+    )
+    # Set p statements to empty and next sym to i if neither parents set
+    yield rule(
+        program_add,
+        p.compile(i),
+        ne(p1.parent).to(p),
+        ne(p2.parent).to(p),
+    ).then(
+        set_(p.statements).to(String("")),
+        set_(p.next_sym).to(i),
+    )
+    # Set p statements to p1 and next sym to p1 if p1 parent set and p2 parent not set
+    yield rule(
+        program_add,
+        eq(p1.parent).to(p),
+        ne(p2.parent).to(p),
+        eq(s1).to(p1.statements),
+        eq(i).to(p1.next_sym),
+    ).then(
+        set_(p.statements).to(s1),
+        set_(p.next_sym).to(i),
+    )
+    # Set p statements to p2 and next sym to p2 if p2 parent set and p1 parent not set
+    yield rule(
+        program_add,
+        eq(p2.parent).to(p),
+        ne(p1.parent).to(p),
+        eq(s2).to(p2.statements),
+        eq(i).to(p2.next_sym),
+    ).then(
+        set_(p.statements).to(s2),
+        set_(p.next_sym).to(i),
+    )
+
+    ##
+    # Assign
+    ##
+
+    # Compiling an assign is the same as compiling the expression, adding an assign statement, then setting the
+    # expression as the gensym, and setting is_identifier to true
+    program_assign = eq(p).to(p1.assign())
+    # Set parent
+    yield rule(program_assign, p.compile(i)).then(set_(p1.parent).to(p), set_(p.is_identifer).to(Bool(True)))
+    # If parent set, compile the expression
+    yield rule(program_assign, p.compile(i), eq(p1.parent).to(p)).then(p1.compile(i))
+
+    # 1. If p1 is not an identifier, then we must create a new one
+
+    # 1. a. If p1 parent is p, then use statements of p, next sym of p
+    symbol = join(String("_"), i.to_string())
+    yield rule(
+        program_assign,
+        eq(p1.parent).to(p),
+        eq(s1).to(p1.statements),
+        eq(i).to(p1.next_sym),
+        eq(s2).to(p1.expr),
+        eq(p1.is_identifer).to(Bool(False)),
+    ).then(
+        set_(p.statements).to(join(s1, symbol, " = ", s2, "\n")),
+        set_(p.expr).to(symbol),
+        set_(p.next_sym).to(i + 1),
+    )
+    # 1. b. If p1 parent is not p, then just use assign as statement, next sym of i
+    yield rule(
+        program_assign,
+        ne(p1.parent).to(p),
+        p.compile(i),
+        eq(s2).to(p1.expr),
+        eq(p1.is_identifer).to(Bool(False)),
+    ).then(
+        set_(p.statements).to(join(symbol, " = ", s2, "\n")),
+        set_(p.expr).to(symbol),
+        set_(p.next_sym).to(i + 1),
+    )
+
+    # 2. If p1 is an identifier, then program assign is a no-op
+
+    # 1. a. If p1 parent is p, then use statements of p, next sym of p
+    yield rule(
+        program_assign,
+        eq(p1.parent).to(p),
+        eq(s1).to(p1.statements),
+        eq(i).to(p1.next_sym),
+        eq(s2).to(p1.expr),
+        eq(p1.is_identifer).to(Bool(True)),
+    ).then(
+        set_(p.statements).to(s1),
+        set_(p.expr).to(s2),
+        set_(p.next_sym).to(i),
+    )
+    # 1. b. If p1 parent is not p, then just use assign as statement, next sym of i
+    yield rule(
+        program_assign,
+        ne(p1.parent).to(p),
+        p.compile(i),
+        eq(s2).to(p1.expr),
+        eq(p1.is_identifer).to(Bool(True)),
+    ).then(
+        set_(p.statements).to(String("")),
+        set_(p.expr).to(s2),
+        set_(p.next_sym).to(i),
+    )
+
+    ##
+    # Function two
+    ##
+
+    # When compiling a function, the two args, p2 and p3, should get compiled when we compile p1, and should just be vars.
+    fn_two = eq(p).to(p1.function_two(p2, p3, s1))
+    # 1. Set parents of both args to p and compile them
+    # Assumes that this if the first thing to compile, so no need to check, and assumes that compiling args doesn't result in any
+    # change in the next sym
+    yield rule(fn_two, p.compile(i)).then(
+        set_(p2.parent).to(p),
+        set_(p3.parent).to(p),
+        set_(p1.parent).to(p),
+        p2.compile(i),
+        p3.compile(i),
+        p1.compile(i),
+        set_(p.is_identifer).to(Bool(True)),
+    )
+    # 2. Set statements to function body and the next sym to i
+    yield rule(
+        fn_two,
+        p.compile(i),
+        eq(s2).to(p1.expr),
+        eq(s3).to(p1.statements),
+        eq(s4).to(p2.expr),
+        eq(s5).to(p3.expr),
+    ).then(
+        set_(p.statements).to(
+            join("def ", s1, "(", s4, ", ", s5, "):\n    ", s3.replace("\n", "\n    "), "return ", s2, "\n")
+        ),
+        set_(p.next_sym).to(i),
+        set_(p.expr).to(s1),
+    )
+
+    ##
+    # Function three
+    ##
+
+    fn_three = eq(p).to(p1.function_three(p2, p3, p4, s1))
+    yield rule(fn_three, p.compile(i)).then(
+        set_(p2.parent).to(p),
+        set_(p3.parent).to(p),
+        set_(p1.parent).to(p),
+        set_(p4.parent).to(p),
+        p2.compile(i),
+        p3.compile(i),
+        p1.compile(i),
+        p4.compile(i),
+        set_(p.is_identifer).to(Bool(True)),
+    )
+    yield rule(
+        fn_three,
+        p.compile(i),
+        eq(s2).to(p1.expr),
+        eq(s3).to(p1.statements),
+        eq(s4).to(p2.expr),
+        eq(s5).to(p3.expr),
+        eq(s6).to(p4.expr),
+    ).then(
+        set_(p.statements).to(
+            join("def ", s1, "(", s4, ", ", s5, ", ", s6, "):\n    ", s3.replace("\n", "\n    "), "return ", s2, "\n")
+        ),
+        set_(p.next_sym).to(i),
+        set_(p.expr).to(s1),
+    )

egglog/exp/siu_examples.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+
+# https://github.com/sklam/pyasir/blob/c363ff4f8f91177700ad4108dd5042b9b97d8289/pyasir/tests/test_fib.py
+
+# In progress - should be able to re-create this
+# @df.func
+# def fib_ir(n: pyasir.Int64) -> pyasir.Int64:
+#     @df.switch(n <= 1)
+#     def swt(n):
+#         @df.case(1)
+#         def case0(n):
+#             return 1
+
+#         @df.case(0)
+#         def case1(n):
+#             return fib_ir(n - 1) + fib_ir(n - 2)
+
+#         yield case0
+#         yield case1
+
+#     r = swt(n)
+#     return r
+
+
+# With something like this:
+# @egglog.function
+# def fib(n: Int) -> Int:
+#     return (n <= Int(1)).if_int(
+#         Int(1),
+#         fib(n - Int(1)) + fib(n - Int(2)),
+#     )

egglog/functionalize.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from functools import partial
+from inspect import Parameter, signature
+from typing import Any, TypeVar, cast
+
+__all__ = ["functionalize"]
+
+
+T = TypeVar("T", bound=Callable)
+
+
+# TODO: Add `to_lift` param so that we only transform those with vars in them to args
+
+
+def functionalize(f: T, get_annotation: Callable[[object], type | None]) -> T:
+    """
+    Takes a function and returns a new function with all names (co_names) and free variables (co_freevars) added as arguments
+    and then partially applied with their values. The second arg, get_annotation, will be applied to all values
+    to get their type annotation. If it is None, that arg will not be added as a parameter.
+
+    For example if you have:
+
+        def get_annotation(x): return int if x <= 10 else None
+
+        g = 10
+        def f(a, a2):
+            def h(b: Z):
+                return a + a2 + b + g
+
+            return functionalize(h, get_annotation)
+        res = f(9, 11)
+
+    It should be equivalent to (according to body, signature, and annotations) (Note that the new arguments will be positional only):
+
+        def h(a: get_annotation(a), g: get_annotation(g), b: Z):
+            return a + b + g
+        res = partial(h, a, g)
+    """
+    code = f.__code__
+    names = tuple(n for n in code.co_names if n in f.__globals__)
+    free_vars = code.co_freevars
+
+    global_values: list[Any] = [f.__globals__[name] for name in names]
+    free_var_values = [cell.cell_contents for cell in f.__closure__] if f.__closure__ else []
+    assert len(free_var_values) == len(free_vars), "Free vars and their values do not match"
+    global_values_filtered = [
+        (i, name, value, annotation)
+        for i, (name, value) in enumerate(zip(names, global_values, strict=True))
+        if (annotation := get_annotation(value)) is not None
+    ]
+    free_var_values_filtered = [
+        (i, name, value, annotation)
+        for i, (name, value) in enumerate(zip(free_vars, free_var_values, strict=True))
+        if (annotation := get_annotation(value)) is not None
+    ]
+    additional_arg_filtered = global_values_filtered + free_var_values_filtered
+
+    # Create a wrapper function
+    def wrapper(*args):
+        # Split args into names, free vars and other args
+        name_args, free_var_args, rest_args = (
+            args[: (n_names := len(global_values_filtered))],
+            args[n_names : (n_args := len(additional_arg_filtered))],
+            args[n_args:],
+        )
+        # Update globals with names
+        f.__globals__.update({
+            name: arg for (_, name, _, _), arg in zip(global_values_filtered, name_args, strict=False)
+        })
+        # update function free vars with free var args
+        for (i, _, _, _), value in zip(free_var_values_filtered, free_var_args, strict=True):
+            assert f.__closure__, "Function does not have closure"
+            f.__closure__[i].cell_contents = value
+        return f(*rest_args)
+
+    # Set the signature of the new function to a signature with the free vars and names added as arguments
+    orig_signature = signature(f)
+    wrapper.__signature__ = orig_signature.replace(  # type: ignore[attr-defined]
+        parameters=[
+            *[Parameter(n, Parameter.POSITIONAL_OR_KEYWORD) for _, n, _, _ in additional_arg_filtered],
+            *orig_signature.parameters.values(),
+        ]
+    )
+    # Set the annotations of the new function to the annotations of the original function + annotations of passed in values
+    wrapper.__annotations__ = f.__annotations__ | {n: a for _, n, _, a in additional_arg_filtered}
+    wrapper.__name__ = f.__name__
+
+    # Partially apply the wrapper function with the current values of the free vars
+    return cast("T", partial(wrapper, *(v for _, _, v, _ in additional_arg_filtered)))

egglog/ipython_magic.py

@@ -0,0 +1,41 @@
+from .bindings import EGraph
+
+EGRAPH_VAR = "_MAGIC_EGRAPH"
+
+try:
+    get_ipython()  # type: ignore[name-defined]
+    IN_IPYTHON = True
+except NameError:
+    IN_IPYTHON = False
+
+if IN_IPYTHON:
+    import graphviz
+    from IPython.core.magic import needs_local_scope, register_cell_magic
+
+    @needs_local_scope
+    @register_cell_magic
+    def egglog(line, cell, local_ns):
+        """
+        Run an egglog program.
+
+        Usage:
+
+            %%egglog [output] [continue] [graph]
+            (egglog program)
+
+        If `output` is specified, the output of the program will be printed.
+        If `continue` is specified, the program will be run in the same EGraph as the previous cell.
+        If `graph` is specified, the EGraph will be displayed as a graph.
+        """
+        if EGRAPH_VAR in local_ns and "continue" in line:
+            e = local_ns[EGRAPH_VAR]
+        else:
+            e = EGraph()
+            local_ns[EGRAPH_VAR] = e
+        cmds = e.parse_program(cell)
+        res = e.run_program(*cmds)
+        if "output" in line:
+            print("\n".join(res))
+        if "graph" in line:
+            return graphviz.Source(e.to_graphviz_string())
+        return None