From 460bd6e78a65710687970079a0486a988ee4258e Mon Sep 17 00:00:00 2001
From: Sasha Lopoukhine <superlopuh@gmail.com>
Date: Tue, 7 Jan 2025 16:11:56 +0000
Subject: [PATCH] documentation: (marimo) move module_html helper to utils
 (#3714)

Unifies how we show modules in marimo notebooks.
---
 docs/marimo/linalg_snitch.py  | 21 +++++----------------
 docs/marimo/onnx/onnx_demo.py | 24 +++++++++---------------
 xdsl/utils/marimo.py          |  9 +++++++++
 3 files changed, 23 insertions(+), 31 deletions(-)

diff --git a/docs/marimo/linalg_snitch.py b/docs/marimo/linalg_snitch.py
index 2ba67ce787..b9409b45a6 100644
--- a/docs/marimo/linalg_snitch.py
+++ b/docs/marimo/linalg_snitch.py
@@ -140,7 +140,7 @@ def _(
     func,
     linalg,
     mo,
-    module_html,
+    xmo,
 ):
     a_type = MemRefType(f64, a_shape)
     b_type = MemRefType(f64, b_shape)
@@ -186,7 +186,7 @@ def _(
 
     Here is matrix multiplication defined in the `linalg` dialect, with the iteration space decoupled from the computation:
 
-    {module_html(linalg_module)}
+    {xmo.module_html(linalg_module)}
     """)
     return (
         a,
@@ -671,17 +671,6 @@ def format_row(key: str, *values: str):
     )
 
 
-@app.cell
-def _(ModuleOp):
-    import html as htmllib
-
-    def module_html(module: ModuleOp) -> str:
-        return f"""\
-        <div style="overflow-y: scroll; height:400px;"><small><code style="white-space: pre-wrap;">{htmllib.escape(str(module))}</code></small></div>
-        """
-    return htmllib, module_html
-
-
 @app.cell
 def _():
     from collections import Counter
@@ -689,7 +678,7 @@ def _():
 
 
 @app.cell
-def _(Counter, ModuleOp, ModulePass, PipelinePass, ctx, mo, module_html):
+def _(Counter, ModuleOp, ModulePass, PipelinePass, ctx, mo, xmo):
     def spec_str(p: ModulePass) -> str:
         if isinstance(p, PipelinePass):
             return ",".join(str(c.pipeline_pass_spec()) for c in p.passes)
@@ -711,12 +700,12 @@ def pipeline_accordion(
                 header = f"{spec} ({d_key_count[spec]})"
             else:
                 header = spec
-            html_res = module_html(res)
+            html_res = xmo.module_html(res)
             d.append(mo.vstack(
                 (
                     header,
                     text,
-                    mo.md(html_res),
+                    html_res,
                 )
             ))
         return (res, mo.carousel(d))
diff --git a/docs/marimo/onnx/onnx_demo.py b/docs/marimo/onnx/onnx_demo.py
index 8c84310cb9..fb8e42b4bf 100644
--- a/docs/marimo/onnx/onnx_demo.py
+++ b/docs/marimo/onnx/onnx_demo.py
@@ -1,6 +1,6 @@
 import marimo
 
-__generated_with = "0.10.0"
+__generated_with = "0.10.9"
 app = marimo.App()
 
 
@@ -132,13 +132,13 @@ def _(mo, model_def):
 
 
 @app.cell
-def _(html, init_module, mo):
+def _(init_module, mo, xmo):
     mo.md(f"""
     ### Converting to `linalg`
 
     Here is the xDSL representation of the function, it takes two `tensor` values of our chosen shape, passes them as operands to the `onnx.Add` operation, and returns it:
 
-    {html(init_module)}
+    {xmo.module_html(init_module)}
     """
     )
     return
@@ -260,15 +260,9 @@ def _():
 
 
 @app.cell(hide_code=True)
-def _(ModuleOp, mo):
-    import html as htmllib
-
-    def html(module: ModuleOp) -> mo.Html:
-        return f"""\
-        <small><code style="white-space: pre-wrap;">{htmllib.escape(str(module))}</code></small>
-        """
-        # return mo.as_html(str(module))
-    return html, htmllib
+def _():
+    import xdsl.utils.marimo as xmo
+    return (xmo,)
 
 
 @app.cell(hide_code=True)
@@ -278,7 +272,7 @@ def _():
 
 
 @app.cell(hide_code=True)
-def _(Counter, ModuleOp, ModulePass, PipelinePass, ctx, html, mo):
+def _(Counter, ModuleOp, ModulePass, PipelinePass, ctx, mo, xmo):
     def spec_str(p: ModulePass) -> str:
         if isinstance(p, PipelinePass):
             return ",".join(str(c.pipeline_pass_spec()) for c in p.passes)
@@ -298,11 +292,11 @@ def pipeline_accordion(passes: tuple[tuple[mo.Html, ModulePass], ...], module: M
                 header = f"{spec} ({d_key_count[spec]})"
             else:
                 header = spec
-            html_res = html(res)
+            html_res = xmo.module_html(res)
             d[header] = mo.vstack((
                 text,
                 # mo.plain_text(f"Pass: {p.pipeline_pass_spec()}"),
-                mo.md(html_res)
+                html_res
             ))
         return (res, mo.accordion(d))
     return pipeline_accordion, spec_str
diff --git a/xdsl/utils/marimo.py b/xdsl/utils/marimo.py
index 99bd216b20..a55a374282 100644
--- a/xdsl/utils/marimo.py
+++ b/xdsl/utils/marimo.py
@@ -1,8 +1,17 @@
 import marimo as mo
 
+from xdsl.dialects.builtin import ModuleOp
+
 
 def asm_html(asm: str) -> mo.Html:
     """
     Returns a Marimo-optimised representation of the assembly code passed in.
     """
     return mo.ui.code_editor(asm, language="python", disabled=True)
+
+
+def module_html(module: ModuleOp) -> mo.Html:
+    """
+    Returns a Marimo-optimised representation of the module passed in.
+    """
+    return mo.ui.code_editor(str(module), language="javascript", disabled=True)