Merge branch 'main' into phil/ccx-rewrite

Author: Roger-luo

src/bloqade/qasm2/_qasm_loading.py

@@ -1,7 +1,12 @@
+import os
+import logging
+import pathlib
 from typing import Any
 
-from kirin import ir
+from kirin import ir, types
+from kirin.dialects import func
 
+from . import parse
 from .groups import main
 from .parse.lowering import QASM2
 
@@ -46,19 +51,40 @@ def loads(
     ''')
     ```
     """
-    return QASM2(dialects or main).loads(
-        qasm,
-        kernel_name=kernel_name,
-        globals=globals,
+    # TODO: add source info
+    stmt = parse.loads(qasm)
+    qasm2_lowering = QASM2(dialects or main)
+    body = qasm2_lowering.run(
+        stmt,
+        source=qasm,
         file=file,
+        globals=globals,
         lineno_offset=lineno_offset,
         col_offset=col_offset,
         compactify=compactify,
     )
+    return_value = func.ConstantNone()
+    body.blocks[0].stmts.append(return_value)
+    body.blocks[0].stmts.append(func.Return(value_or_stmt=return_value))
+
+    code = func.Function(
+        sym_name=kernel_name,
+        signature=func.Signature((), types.NoneType),
+        body=body,
+    )
+
+    return ir.Method(
+        mod=None,
+        py_func=None,
+        sym_name=kernel_name,
+        arg_names=[],
+        dialects=qasm2_lowering.dialects,
+        code=code,
+    )
 
 
 def loadfile(
-    qasm_file: str,
+    qasm_file: str | pathlib.Path,
     *,
     kernel_name: str = "main",
     dialects: ir.DialectGroup | None = None,
@@ -83,10 +109,27 @@ def loadfile(
         col_offset (int): The column number offset for error reporting. Defaults to 0.
         compactify (bool): Whether to compactify the output. Defaults to True.
     """
-    with open(qasm_file, "r") as f:
-        qasm = f.read()
+    if isinstance(file, pathlib.Path):
+        qasm_file_: pathlib.Path = qasm_file  # type: ignore
+    else:
+        qasm_file_ = pathlib.Path(*os.path.split(qasm_file))
+
+    if not qasm_file_.is_file():
+        raise FileNotFoundError(f"File {qasm_file_} does not exist")
+
+    if not qasm_file_.name.endswith(".qasm") or not qasm_file_.name.endswith(".qasm2"):
+        logging.warning(
+            f"File {qasm_file_} does not end with .qasm or .qasm2. "
+            "This may cause issues with loading the file."
+        )
+
+    kernel_name = file.name.replace(".qasm", "") if kernel_name is None else kernel_name
+
+    with qasm_file_.open("r") as f:
+        source = f.read()
+
     return loads(
-        qasm,
+        source,
         kernel_name=kernel_name,
         dialects=dialects,
         globals=globals,

src/bloqade/qasm2/parse/__init__.py

@@ -19,19 +19,22 @@ def loadfile(file: str | pathlib.Path):
         return loads(f.read())
 
 
-def pprint(node: ast.Node, *, console: Console | None = None):
+def pprint(node: ast.Node, *, console: Console | None = None, no_color: bool = False):
     if console:
-        return Printer(console).visit(node)
+        printer = Printer(console)
     else:
-        Printer().visit(node)
+        printer = Printer()
+    printer.console.no_color = no_color
+    printer.visit(node)
 
 
-def spprint(node: ast.Node, *, console: Console | None = None):
+def spprint(node: ast.Node, *, console: Console | None = None, no_color: bool = False):
     if console:
         printer = Printer(console)
     else:
         printer = Printer()
 
+    printer.console.no_color = no_color
     with printer.string_io() as stream:
         printer.visit(node)
         return stream.getvalue()

src/bloqade/qasm2/parse/lowering.py

@@ -1,5 +1,3 @@
-import os
-import pathlib
 from typing import Any
 from dataclasses import field, dataclass
 
@@ -19,119 +17,6 @@ class QASM2(lowering.LoweringABC[ast.Node]):
     hint_show_lineno: bool = field(default=True, kw_only=True)
     stacktrace: bool = field(default=True, kw_only=True)
 
-    def loads(
-        self,
-        source: str,
-        kernel_name: str,
-        *,
-        returns: str | None = None,
-        globals: dict[str, Any] | None = None,
-        file: str | None = None,
-        lineno_offset: int = 0,
-        col_offset: int = 0,
-        compactify: bool = True,
-    ) -> ir.Method:
-        from ..parse import loads
-
-        # TODO: add source info
-        stmt = loads(source)
-
-        state = lowering.State(
-            self,
-            file=file,
-            lineno_offset=lineno_offset,
-            col_offset=col_offset,
-        )
-        with state.frame(
-            [stmt],
-            globals=globals,
-            finalize_next=False,
-        ) as frame:
-            try:
-                self.visit(state, stmt)
-                # append return statement with the return values
-                if returns is not None:
-                    return_value = frame.get(returns)
-                    if return_value is None:
-                        raise lowering.BuildError(f"Cannot find return value {returns}")
-                else:
-                    return_value = func.ConstantNone()
-
-                return_node = frame.push(func.Return(value_or_stmt=return_value))
-
-            except lowering.BuildError as e:
-                hint = state.error_hint(
-                    e,
-                    max_lines=self.max_lines,
-                    indent=self.hint_indent,
-                    show_lineno=self.hint_show_lineno,
-                )
-                if self.stacktrace:
-                    raise Exception(
-                        f"{e.args[0]}\n\n{hint}",
-                        *e.args[1:],
-                    ) from e
-                else:
-                    e.args = (hint,)
-                    raise e
-
-            region = frame.curr_region
-
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
-
-            Walk(CFGCompactify()).rewrite(region)
-
-        code = func.Function(
-            sym_name=kernel_name,
-            signature=func.Signature((), return_node.value.type),
-            body=region,
-        )
-
-        return ir.Method(
-            mod=None,
-            py_func=None,
-            sym_name=kernel_name,
-            arg_names=[],
-            dialects=self.dialects,
-            code=code,
-        )
-
-    def loadfile(
-        self,
-        file: str | pathlib.Path,
-        *,
-        kernel_name: str | None = None,
-        returns: str | None = None,
-        globals: dict[str, Any] | None = None,
-        lineno_offset: int = 0,
-        col_offset: int = 0,
-        compactify: bool = True,
-    ) -> ir.Method:
-        if isinstance(file, str):
-            file = pathlib.Path(*os.path.split(file))
-
-        if not file.is_file() or not file.name.endswith(".qasm"):
-            raise ValueError("File must be a .qasm file")
-
-        kernel_name = (
-            file.name.replace(".qasm", "") if kernel_name is None else kernel_name
-        )
-
-        with file.open("r") as f:
-            source = f.read()
-
-        return self.loads(
-            source,
-            kernel_name,
-            returns=returns,
-            globals=globals,
-            file=str(file),
-            lineno_offset=lineno_offset,
-            col_offset=col_offset,
-            compactify=compactify,
-        )
-
     def run(
         self,
         stmt: ast.Node,
@@ -153,25 +38,9 @@ def run(
         with state.frame(
             [stmt],
             globals=globals,
+            finalize_next=False,
         ) as frame:
-            try:
-                self.visit(state, stmt)
-            except lowering.BuildError as e:
-                hint = state.error_hint(
-                    e,
-                    max_lines=self.max_lines,
-                    indent=self.hint_indent,
-                    show_lineno=self.hint_show_lineno,
-                )
-                if self.stacktrace:
-                    raise Exception(
-                        f"{e.args[0]}\n\n{hint}",
-                        *e.args[1:],
-                    ) from e
-                else:
-                    e.args = (hint,)
-                    raise e
-
+            self.visit(state, stmt)
             region = frame.curr_region
 
         if compactify:

test/qasm2/test_lowering.py

@@ -2,7 +2,7 @@
 import tempfile
 import textwrap
 
-from kirin import ir
+from kirin import ir, types
 from kirin.dialects import func
 
 from bloqade import qasm2
@@ -37,12 +37,7 @@ def test_loadfile():
             f.write(lines)
 
         file = pathlib.Path(f"{tmp_dir}/test.qasm")
-        kernel = QASM2(qasm2.main).loadfile(file, returns="c")
-
-        assert isinstance(
-            (ret := kernel.callable_region.blocks[0].last_stmt), func.Return
-        )
-        assert ret.value.type.is_equal(qasm2.types.CRegType)
+        qasm2.loadfile(file)
 
 
 def test_negative_lowering():
@@ -54,7 +49,7 @@ def test_negative_lowering():
     rz(-0.2) q[0];
     """
 
-    entry = QASM2(qasm2.main).loads(mwe, "entry", returns="q")
+    entry = qasm2.loads(mwe)
 
     body = ir.Region(
         ir.Block(
@@ -66,17 +61,19 @@ def test_negative_lowering():
                 (idx := qasm2.expr.ConstInt(value=0)),
                 (qubit := qasm2.core.QRegGet(qreg.result, idx.result)),
                 (qasm2.uop.RZ(qubit.result, theta.result)),
-                (func.Return(qreg.result)),
+                (none := func.ConstantNone()),
+                (func.Return(none.result)),
             ]
         )
     )
 
     code = func.Function(
-        sym_name="entry",
-        signature=func.Signature((), qasm2.QRegType),
+        sym_name="main",
+        signature=func.Signature((), types.NoneType),
         body=body,
     )
 
     code.print()
+    entry.print()
 
     assert entry.code.is_structurally_equal(code)