Fix bug in qasm2 custom gate emit (#202)

* fix emit bug

* lint

Author: kaihsin

src/bloqade/qasm2/dialects/expr/_emit.py

@@ -1,7 +1,6 @@
 from typing import Literal
 
 from kirin import interp
-from kirin.emit.exceptions import EmitError
 
 from bloqade.qasm2.parse import ast
 from bloqade.qasm2.types import QubitType
@@ -19,17 +18,18 @@ def emit_func(
         self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.GateFunction
     ):
 
-        args, cparams, qparams = [], [], []
-        for arg in stmt.body.blocks[0].args[1:]:
-            name = frame.get_typed(arg, ast.Name)
-            args.append(name)
-            if not isinstance(name, ast.Name):
-                raise EmitError("expected ast.Name")
+        args: list[ast.Node] = []
+        cparams, qparams = [], []
+        for arg in stmt.body.blocks[0].args:
+            assert arg.name is not None
+
+            args.append(ast.Name(id=arg.name))
             if arg.type.is_subseteq(QubitType):
-                qparams.append(name.id)
+                qparams.append(arg.name)
             else:
-                cparams.append(name.id)
-        emit.run_ssacfg_region(frame, stmt.body, args)
+                cparams.append(arg.name)
+
+        emit.run_ssacfg_region(frame, stmt.body, tuple(args))
         emit.output = ast.Gate(
             name=stmt.sym_name,
             cparams=cparams,

src/bloqade/qasm2/emit/gate.py

@@ -86,21 +86,3 @@ def emit_err(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
     @interp.impl(func.ConstantNone)
     def ignore(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
         return ()
-
-    @interp.impl(func.Function)
-    def emit_func(
-        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: func.Function
-    ):
-        args_ssa = stmt.args
-        print(stmt.args)
-        emit.run_ssacfg_region(frame, stmt.body, frame.get_values(args_ssa))
-
-        cparams, qparams = [], []
-        for arg in args_ssa:
-            if arg.type.is_subseteq(QubitType):
-                qparams.append(frame.get(arg))
-            else:
-                cparams.append(frame.get(arg))
-
-        emit.output = ast.Gate(stmt.sym_name, cparams, qparams, frame.body)
-        return ()

src/bloqade/qasm2/emit/target.py

@@ -101,7 +101,7 @@ def emit(self, entry: ir.Method) -> ast.MainProgram:
 
         Py2QASM(entry.dialects)(entry)
         target_main = EmitQASM2Main(self.main_target)
-        target_main.run(entry, tuple(ast.Name(name) for name in entry.arg_names[1:]))
+        target_main.run(entry, ())
 
         main_program = target_main.output
         assert main_program is not None, f"failed to emit {entry.sym_name}"

test/qasm2/emit/test_qasm2.py

@@ -1,9 +1,6 @@
-import pytest
-
 from bloqade import qasm2
 
 
-@pytest.mark.skip(reason="broken gate emit!")
 def test_qasm2_custom_gate():
     @qasm2.gate
     def custom_gate(a: qasm2.Qubit, b: qasm2.Qubit):