From e89bab9662e62a549b980fe05d5a8314ce6dc2a3 Mon Sep 17 00:00:00 2001 From: swap357 Date: Thu, 5 Jun 2025 16:16:00 -0700 Subject: [PATCH 1/5] add egglog_to_inference.py with functions for tokenizing, parsing, and rendering s-expressions to latex --- sealir-tutorials/egglog_to_inference.py | 162 ++++++++++++++++++++++++ 1 file changed, 162 insertions(+) create mode 100644 sealir-tutorials/egglog_to_inference.py diff --git a/sealir-tutorials/egglog_to_inference.py b/sealir-tutorials/egglog_to_inference.py new file mode 100644 index 0000000..264a68a --- /dev/null +++ b/sealir-tutorials/egglog_to_inference.py @@ -0,0 +1,162 @@ +from typing import List, Union + + +def tokenize(egglog_str: str) -> List[str]: + """ + Splits an Egglog S-expression string into a flat list of tokens. + Tokens are either "(" or ")", or atoms (any sequence of non-whitespace, non-parenthesis chars). + """ + tokens = [] + i = 0 + while i < len(egglog_str): + c = egglog_str[i] + if c.isspace(): + i += 1 + continue + if c in ("(", ")"): + tokens.append(c) + i += 1 + else: + j = i + while j < len(egglog_str) and not egglog_str[j].isspace() and egglog_str[j] not in ("(", ")"): + j += 1 + tokens.append(egglog_str[i:j]) + i = j + return tokens + + +def parse_sexps(tokens: List[str]) -> List[Union[str, list]]: + """ + Parses a flat list of tokens into a nested list of S-expression forms. + Each form is either an atom (string) or a list whose first element is the head. + Returns a flat list of top-level S-expressions (each itself a nested list). + """ + stack: List[List] = [] + current: List[Union[str, list]] = [] + for tok in tokens: + if tok == "(": + stack.append(current) + current = [] + elif tok == ")": + completed = current + current = stack.pop() + current.append(completed) + else: + current.append(tok) + # If the entire parse wrapped everything in a single list, unwrap it: + if len(current) == 1 and isinstance(current[0], list): + return current[0] + return current + +def sexp_to_string(sexp): + """Convert parsed S-expression back to original string format""" + if isinstance(sexp, str): + return sexp + elif isinstance(sexp, list): + inner = ' '.join(sexp_to_string(item) for item in sexp) + return f"({inner})" + else: + return str(sexp) + +LATEX_ESCAPE = str.maketrans({ + "_": r"\_", + "#": r"\#", + "$": r"\$", + "%": r"\%", + "&": r"\&", + "{": r"\{", + "}": r"\}", + "~": r"\textasciitilde{}", + "^": r"\^{}", + "\\": r"\\", +}) + +def _atom_tex(a: str) -> str: + try: + float(a) # leave numerics bare + return a + except ValueError: + return r"\text{" + a.translate(LATEX_ESCAPE) + "}" + +INFIX_OPS = {"=", "!=", "<", "<=", ">", ">=", "+", "-", "*", "/", "%", "**"} + +def _sexp_tex(x) -> str: + if isinstance(x, str): + return _atom_tex(x) + + head, *args = x + + # infix pretty-printing for common binary ops + if head in INFIX_OPS and len(args) == 2: + return f"{_sexp_tex(args[0])} {head} {_sexp_tex(args[1])}" + + return ( + r"\text{" + head.translate(LATEX_ESCAPE) + "}" + + "(" + ", ".join(_sexp_tex(a) for a in args) + ")" + ) + +def _is_set_expr(x): + # Detects if x is a set-like S-expression: ['set', lhs, rhs] + return isinstance(x, list) and len(x) == 3 and x[0] == "set" + +def _set_tex(x): + # Renders set(lhs, rhs) as lhs \to rhs + return f"{_sexp_tex(x[1])} \\to {_sexp_tex(x[2])}" + +def to_latex(sexp): + """ + Render (rewrite …) or (rule …) as KaTeX-safe LaTeX. + """ + if not (isinstance(sexp, list) and sexp): + return None + + tag = sexp[0] + + # ─────────────── REWRITE ──────────────── + if tag == "rewrite" and len(sexp) >= 3: + lhs, rhs = sexp[1], sexp[2] + + # harvest optional :when clause (list of conditions) + when_conds = [] + i = 3 + while i < len(sexp): + if sexp[i] == ":when" and i + 1 < len(sexp): + when_conds = sexp[i + 1] # list of cond S-exps + break + i += 1 # <- step only ONE token + + lhs_tex = _sexp_tex(lhs) + rhs_tex = _sexp_tex(rhs) + + cond_tex = "" + if when_conds: + joined = r",\; ".join(_sexp_tex(c) for c in when_conds) + cond_tex = rf",\; {joined}" + + num = rf"\text{{expr}} = {lhs_tex}{cond_tex}" + den = rf"\text{{expr}} \to {rhs_tex}" + + return rf"\frac{{{num}}}{{{den}}}" + + # ──────────────── RULE ───────────────── + if tag == "rule" and len(sexp) >= 3: + premises, conclusions = sexp[1], sexp[2] + + def render_stack(exprs): + lines = [] + for e in exprs: + if _is_set_expr(e): + lines.append(_set_tex(e)) + else: + lines.append(_sexp_tex(e)) + return r"\\ ".join(lines) + + prem_tex = render_stack(premises) + concl_tex = render_stack(conclusions) + + num = rf"\begin{{array}}{{c}}{prem_tex}\end{{array}}" + den = rf"\begin{{array}}{{c}}{concl_tex}\end{{array}}" + + return rf"\frac{{{num}}}{{{den}}}" + + return None From 3cdc04ce8ec48e6f8504f4696f551e5ebe36d8f0 Mon Sep 17 00:00:00 2001 From: swap357 Date: Thu, 5 Jun 2025 16:17:23 -0700 Subject: [PATCH 2/5] add visualization of ruleset structure in chapter 3 using functions from egglog_to_inference for latex rendering, conditional display for jupyter nb --- .../ch03_egraph_program_rewrites.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/sealir-tutorials/ch03_egraph_program_rewrites.py b/sealir-tutorials/ch03_egraph_program_rewrites.py index 4ab4247..8d4192d 100644 --- a/sealir-tutorials/ch03_egraph_program_rewrites.py +++ b/sealir-tutorials/ch03_egraph_program_rewrites.py @@ -24,6 +24,8 @@ from sealir import rvsdg from sealir.eqsat import rvsdg_eqsat from sealir.eqsat.rvsdg_eqsat import GraphRoot, Term, TermList +from egglog_to_inference import tokenize, parse_sexps, to_latex, sexp_to_string +from utils import IN_NOTEBOOK # We'll be extending from chapter 2. from ch02_egraph_basic import ( @@ -181,6 +183,23 @@ def ruleset_const_fold_if_else(a: Term, b: Term, c: Term, operands: TermList): | ruleset_const_fold_if_else # <-- the new rule for if-else ) + # Visualize the ruleset structure + demo_egraph = EGraph(save_egglog_string=True) + demo_egraph.run(my_ruleset) + egglog_str = demo_egraph.as_egglog_string + tokens = tokenize(egglog_str) + sexps = parse_sexps(tokens) + + if IN_NOTEBOOK: + from IPython.display import display, Math + + for sexp in sexps: + tex = to_latex(sexp) + if tex: + print(sexp_to_string(sexp)) + display(Math(tex)) + print() + jt = compiler_pipeline(ifelse_fold, verbose=True, ruleset=my_ruleset) run_test(ifelse_fold, jt, (12, 34)) From 3788e4ebcd03484bcff9665842432446344975c2 Mon Sep 17 00:00:00 2001 From: swap357 Date: Thu, 5 Jun 2025 16:41:03 -0700 Subject: [PATCH 3/5] rename util appropriate- egglog_to_latex --- sealir-tutorials/ch03_egraph_program_rewrites.py | 2 +- sealir-tutorials/{egglog_to_inference.py => egglog_to_latex.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename sealir-tutorials/{egglog_to_inference.py => egglog_to_latex.py} (100%) diff --git a/sealir-tutorials/ch03_egraph_program_rewrites.py b/sealir-tutorials/ch03_egraph_program_rewrites.py index 8d4192d..f92a5fc 100644 --- a/sealir-tutorials/ch03_egraph_program_rewrites.py +++ b/sealir-tutorials/ch03_egraph_program_rewrites.py @@ -24,7 +24,7 @@ from sealir import rvsdg from sealir.eqsat import rvsdg_eqsat from sealir.eqsat.rvsdg_eqsat import GraphRoot, Term, TermList -from egglog_to_inference import tokenize, parse_sexps, to_latex, sexp_to_string +from egglog_to_latex import tokenize, parse_sexps, to_latex, sexp_to_string from utils import IN_NOTEBOOK # We'll be extending from chapter 2. diff --git a/sealir-tutorials/egglog_to_inference.py b/sealir-tutorials/egglog_to_latex.py similarity index 100% rename from sealir-tutorials/egglog_to_inference.py rename to sealir-tutorials/egglog_to_latex.py From 62a840bfed59a7b2d94150268c51f45f9a1b8ca6 Mon Sep 17 00:00:00 2001 From: swap357 Date: Fri, 13 Jun 2025 12:56:14 -0700 Subject: [PATCH 4/5] add visualize_ruleset_latex() that takes entire egglog ruleset and renders latex --- sealir-tutorials/egglog_to_latex.py | 51 +++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/sealir-tutorials/egglog_to_latex.py b/sealir-tutorials/egglog_to_latex.py index 264a68a..0c81e93 100644 --- a/sealir-tutorials/egglog_to_latex.py +++ b/sealir-tutorials/egglog_to_latex.py @@ -1,4 +1,5 @@ from typing import List, Union +from egglog import EGraph def tokenize(egglog_str: str) -> List[str]: @@ -48,6 +49,7 @@ def parse_sexps(tokens: List[str]) -> List[Union[str, list]]: return current[0] return current + def sexp_to_string(sexp): """Convert parsed S-expression back to original string format""" if isinstance(sexp, str): @@ -58,6 +60,7 @@ def sexp_to_string(sexp): else: return str(sexp) + LATEX_ESCAPE = str.maketrans({ "_": r"\_", "#": r"\#", @@ -71,6 +74,7 @@ def sexp_to_string(sexp): "\\": r"\\", }) + def _atom_tex(a: str) -> str: try: float(a) # leave numerics bare @@ -78,8 +82,10 @@ def _atom_tex(a: str) -> str: except ValueError: return r"\text{" + a.translate(LATEX_ESCAPE) + "}" + INFIX_OPS = {"=", "!=", "<", "<=", ">", ">=", "+", "-", "*", "/", "%", "**"} + def _sexp_tex(x) -> str: if isinstance(x, str): return _atom_tex(x) @@ -95,14 +101,17 @@ def _sexp_tex(x) -> str: + "(" + ", ".join(_sexp_tex(a) for a in args) + ")" ) + def _is_set_expr(x): # Detects if x is a set-like S-expression: ['set', lhs, rhs] return isinstance(x, list) and len(x) == 3 and x[0] == "set" + def _set_tex(x): # Renders set(lhs, rhs) as lhs \to rhs return f"{_sexp_tex(x[1])} \\to {_sexp_tex(x[2])}" + def to_latex(sexp): """ Render (rewrite …) or (rule …) as KaTeX-safe LaTeX. @@ -160,3 +169,45 @@ def render_stack(exprs): return rf"\frac{{{num}}}{{{den}}}" return None + + +def visualize_ruleset_latex(ruleset, verbose=True): + """ + Visualize an egglog ruleset by converting it to LaTeX representation. + Only works in notebook environments. + + Args: + ruleset: The egglog ruleset to visualize + verbose: If True, prints the original S-expression before LaTeX display + + Returns: + None, but displays LaTeX representation if in notebook environment + """ + try: + shell = get_ipython().__class__.__name__ + is_notebook = shell == "ZMQInteractiveShell" + except NameError: + is_notebook = False + + if not is_notebook: + return + + # Create demo egraph and run ruleset + demo_egraph = EGraph(save_egglog_string=True) + demo_egraph.run(ruleset) + egglog_str = demo_egraph.as_egglog_string + + # Parse into S-expressions + tokens = tokenize(egglog_str) + sexps = parse_sexps(tokens) + + from IPython.display import display, Math + + for sexp in sexps: + tex = to_latex(sexp) + if tex: + if verbose: + print(sexp_to_string(sexp)) + display(Math(tex)) + if verbose: + print() From 297bd11113f95a1c3a0d6f0d0ca71fd10ba071e6 Mon Sep 17 00:00:00 2001 From: swap357 Date: Fri, 13 Jun 2025 12:57:31 -0700 Subject: [PATCH 5/5] simplify visualize_ruleset_latex(), show latex after ruleset definitions --- .../ch03_egraph_program_rewrites.py | 27 ++++++------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/sealir-tutorials/ch03_egraph_program_rewrites.py b/sealir-tutorials/ch03_egraph_program_rewrites.py index f92a5fc..360eeb3 100644 --- a/sealir-tutorials/ch03_egraph_program_rewrites.py +++ b/sealir-tutorials/ch03_egraph_program_rewrites.py @@ -24,8 +24,8 @@ from sealir import rvsdg from sealir.eqsat import rvsdg_eqsat from sealir.eqsat.rvsdg_eqsat import GraphRoot, Term, TermList -from egglog_to_latex import tokenize, parse_sexps, to_latex, sexp_to_string from utils import IN_NOTEBOOK +from egglog_to_latex import visualize_ruleset_latex # We'll be extending from chapter 2. from ch02_egraph_basic import ( @@ -118,8 +118,11 @@ def ruleset_const_propagate(a: Term, ival: i64): IsConstantFalse(a) ) +if IN_NOTEBOOK: + # Visualize the constant propagation ruleset + visualize_ruleset_latex(ruleset_const_propagate) -# Now, we’ll test our newly defined ruleset. This complete ruleset combines a +# Now, we'll test our newly defined ruleset. This complete ruleset combines a # few built-in RVSDG rules with our recently crafted simple constant-propagation # rules. @@ -175,6 +178,9 @@ def ruleset_const_fold_if_else(a: Term, b: Term, c: Term, operands: TermList): IsConstantFalse(a), ) +if IN_NOTEBOOK: + # Visualize the if-else folding ruleset + visualize_ruleset_latex(ruleset_const_fold_if_else) if __name__ == "__main__": my_ruleset = ( @@ -183,23 +189,6 @@ def ruleset_const_fold_if_else(a: Term, b: Term, c: Term, operands: TermList): | ruleset_const_fold_if_else # <-- the new rule for if-else ) - # Visualize the ruleset structure - demo_egraph = EGraph(save_egglog_string=True) - demo_egraph.run(my_ruleset) - egglog_str = demo_egraph.as_egglog_string - tokens = tokenize(egglog_str) - sexps = parse_sexps(tokens) - - if IN_NOTEBOOK: - from IPython.display import display, Math - - for sexp in sexps: - tex = to_latex(sexp) - if tex: - print(sexp_to_string(sexp)) - display(Math(tex)) - print() - jt = compiler_pipeline(ifelse_fold, verbose=True, ruleset=my_ruleset) run_test(ifelse_fold, jt, (12, 34))