cleanup loads/loadfile in lowering (#242)

follow up #240 I removed the `returns` option cuz it seems no one is
using it. We can add something back if there is an actual user asking
for it.

I'm not sure why it was copy pasting `QASM2.run` method in the old
`loads` implementation, maybe there was some issue with the lowering so
this was a workaround? @weinbe58

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/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)