From 5755fd766bdc847f958459537ae6c0c2d90cd872 Mon Sep 17 00:00:00 2001
From: Seyon Sivarajah <seyon.sivarajah@quantinuum.com>
Date: Tue, 10 Dec 2024 15:24:31 +0000
Subject: [PATCH] feat: add `maybe_qubit` stdlib function

Closes #627
---
 guppylang/std/quantum.py          | 8 ++++++++
 tests/integration/test_quantum.py | 6 +++---
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/guppylang/std/quantum.py b/guppylang/std/quantum.py
index c4afa97b..d937da61 100644
--- a/guppylang/std/quantum.py
+++ b/guppylang/std/quantum.py
@@ -14,6 +14,7 @@
 from guppylang.std._internal.util import quantum_op
 from guppylang.std.angles import angle
 from guppylang.std.builtins import owned
+from guppylang.std.option import Option
 
 
 @guppy.type(ht.Qubit, linear=True)
@@ -38,6 +39,13 @@ def discard(self: "qubit" @ owned) -> None:
         discard(self)
 
 
+@guppy.hugr_op(quantum_op("TryQAlloc"))
+@no_type_check
+def maybe_qubit() -> Option[qubit]:
+    """Try to allocate a qubit, returning `some(qubit)`
+    if allocation succeeds or `nothing` if it fails."""
+
+
 @guppy.hugr_op(quantum_op("H"))
 @no_type_check
 def h(q: qubit) -> None: ...
diff --git a/tests/integration/test_quantum.py b/tests/integration/test_quantum.py
index 878328d1..b94efef0 100644
--- a/tests/integration/test_quantum.py
+++ b/tests/integration/test_quantum.py
@@ -8,7 +8,7 @@
 
 from guppylang.std.builtins import owned
 
-from guppylang.std.quantum import discard, measure, qubit
+from guppylang.std.quantum import discard, measure, qubit, maybe_qubit
 from guppylang.std.quantum_functional import (
     cx,
     cy,
@@ -43,7 +43,7 @@ def compile_quantum_guppy(fn) -> ModulePointer:
     ), "`@compile_quantum_guppy` does not support extra arguments."
 
     module = GuppyModule("module")
-    module.load(angle, qubit, discard, measure)
+    module.load(angle, qubit, discard, measure, maybe_qubit)
     module.load_all(quantum_functional)
     guppylang.decorator.guppy(module)(fn)
     return module.compile()
@@ -52,7 +52,7 @@ def compile_quantum_guppy(fn) -> ModulePointer:
 def test_alloc(validate):
     @compile_quantum_guppy
     def test() -> tuple[bool, bool]:
-        q1, q2 = qubit(), qubit()
+        q1, q2 = qubit(), maybe_qubit().unwrap()
         q1, q2 = cx(q1, q2)
         return (measure(q1), measure(q2))