Handle named args in squin apply lowering (#388)

Closes #380

Author: david-pl

src/bloqade/squin/lowering.py

@@ -15,13 +15,40 @@ class ApplyAnyCallLowering(lowering.FromPythonCall["qubit.ApplyAny"]):
     def lower(
         self, stmt: type["qubit.ApplyAny"], state: lowering.State, node: ast.Call
     ):
-        if len(node.args) < 2:
+        if len(node.args) + len(node.keywords) < 2:
             raise lowering.BuildError(
                 "Apply requires at least one operator and one qubit as arguments!"
             )
-        op, *qubits = node.args
+
+        op, qubits = self.unpack_arguments(node)
+
         op_ssa = state.lower(op).expect_one()
         qubits_lowered = [state.lower(qbit).expect_one() for qbit in qubits]
 
         s = stmt(op_ssa, tuple(qubits_lowered))
         return state.current_frame.push(s)
+
+    def unpack_arguments(self, node: ast.Call) -> tuple[ast.expr, list[ast.expr]]:
+        if len(node.keywords) == 0:
+            op, *qubits = node.args
+            return op, qubits
+
+        kwargs = {kw.arg: kw.value for kw in node.keywords}
+        if len(kwargs) > 2 or "qubits" not in kwargs:
+            raise lowering.BuildError(f"Got unsupported keyword argument {kwargs}")
+
+        qubits = kwargs["qubits"]
+        if len(kwargs) == 1:
+            if len(node.args) != 1:
+                raise lowering.BuildError("Missing operator argument")
+            op = node.args[0]
+        else:
+            try:
+                op = kwargs["operator"]
+            except KeyError:
+                raise lowering.BuildError(f"Got unsupported keyword argument {kwargs}")
+
+        if isinstance(qubits, ast.List):
+            return op, qubits.elts
+
+        return op, [qubits]

test/squin/test_sugar.py

@@ -104,3 +104,15 @@ def main():
     sim = StackMemorySimulator(min_qubits=3)
     ket = sim.state_vector(main)
     assert math.isclose(abs(ket[0]) ** 2, 0.5, abs_tol=1e-5)
+
+
+def test_apply_with_named_args():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(2)
+        h = squin.op.h()
+        squin.qubit.apply(h, qubits=[q[0]])
+        cx = squin.op.cx()
+        squin.qubit.apply(qubits=q, operator=cx)
+
+    main.print()