Add support for returning values from defcal statements (#24)

* Added 'returns' statement to oqpy.Program for adding return statement to defcal and subroutine statements, added support for using Program.returns in subroutine decorator based on type hinting and added tests.

* Update oqpy/subroutines.py

Co-authored-by: Phil Reinhold <[email protected]>

* �Added new types: port, frame and waveform to represent the types of PortVar, FrameVar and WaveformVar

* port, waveform and frame now aliases of openpulse.ast PortType, WaveformType and FrameType

* Added missing period at the end of a docstring

Co-authored-by: Phil Reinhold <[email protected]>

Author: Swordcat

oqpy/program.py

@@ -405,6 +405,11 @@ def shift_scale(self, frame: AstConvertible, scale: AstConvertible) -> Program:
         self.function_call("shift_scale", [frame, scale])
         return self
 
+    def returns(self, expression: AstConvertible) -> Program:
+        """Return a statement from a function definition or a defcal statement."""
+        self._add_statement(ast.ReturnStatement(to_ast(self, expression)))
+        return self
+
     def gate(
         self, qubits: AstConvertible | Iterable[AstConvertible], name: str, *args: Any
     ) -> Program:

oqpy/pulse.py

@@ -24,7 +24,11 @@
 from oqpy.base import AstConvertible
 from oqpy.classical_types import OQFunctionCall, _ClassicalVar
 
-__all__ = ["PortVar", "WaveformVar", "FrameVar"]
+__all__ = ["PortVar", "WaveformVar", "FrameVar", "port", "waveform", "frame"]
+
+port = ast.PortType()
+waveform = ast.WaveformType()
+frame = ast.FrameType()
 
 
 class PortVar(_ClassicalVar):

oqpy/subroutines.py

@@ -99,6 +99,14 @@ def wrapper(program: oqpy.Program, *args: AstConvertible) -> OQFunctionCall:
             body.append(ast.ReturnStatement(to_ast(inner_prog, output)))
         elif output is None:
             return_type = None
+            if type_hints.get("return", False):
+                return_hint = type_hints["return"]()
+                if isinstance(return_hint, _ClassicalVar):
+                    return_type = return_hint
+                elif return_hint is not None:
+                    raise ValueError(
+                        f"Type hint for return variable on subroutine {name} is not an oqpy classical type."
+                    )
         else:
             raise ValueError(
                 "Output type of subroutine {name} was neither oqpy expression nor None."

tests/test_directives.py

@@ -407,6 +407,14 @@ def return1(prog: Program) -> float:
 
         return1(prog)
 
+    with pytest.raises(ValueError):
+
+        @subroutine
+        def return2(prog: Program) -> float:
+            prog.returns(1.0)
+
+        return2(prog)
+
     with pytest.raises(ValueError):
 
         @subroutine
@@ -661,6 +669,88 @@ def test_defcals():
     )
 
 
+def test_returns():
+    prog = Program()
+    constant = declare_waveform_generator("constant", [("length", duration), ("iq", complex128)])
+
+    rx_port = PortVar("rx_port")
+    tx_port = PortVar("tx_port")
+    rx_frame = FrameVar(rx_port, 5.752e9, name="rx_frame")
+    tx_frame = FrameVar(tx_port, 5.752e9, name="tx_frame")
+    capture_v2 = oqpy.declare_extern(
+        "capture_v2", [("output", oqpy.frame), ("duration", oqpy.duration)], oqpy.bit
+    )
+
+    q0 = PhysicalQubits[0]
+
+    with defcal(prog, q0, "measure_v1", return_type=oqpy.bit):
+        prog.play(tx_frame, constant(2.4e-6, 0.2))
+        prog.returns(capture_v2(rx_frame, 2.4e-6))
+
+    @subroutine
+    def increment_variable_return(prog: Program, i: IntVar) -> IntVar:
+        prog.increment(i, 1)
+        prog.returns(i)
+
+    j = IntVar(0, name="j")
+    k = IntVar(0, name="k")
+    prog.declare(j)
+    prog.declare(k)
+    prog.set(k, increment_variable_return(prog, j))
+
+    expected = textwrap.dedent(
+        """
+        OPENQASM 3.0;
+        extern constant(duration, complex[float[64]]) -> waveform;
+        extern capture_v2(frame, duration) -> bit;
+        def increment_variable_return(int[32] i) -> int[32] {
+            i += 1;
+            return i;
+        }
+        port rx_port;
+        port tx_port;
+        frame tx_frame = newframe(tx_port, 5752000000.0, 0);
+        frame rx_frame = newframe(rx_port, 5752000000.0, 0);
+        defcal measure_v1 $0 -> bit {
+            play(tx_frame, constant(2400.0ns, 0.2));
+            return capture_v2(rx_frame, 2400.0ns);
+        }
+        int[32] j = 0;
+        int[32] k = 0;
+        k = increment_variable_return(j);
+        """
+    ).strip()
+    print(prog.to_qasm())
+    assert prog.to_qasm() == expected
+
+    expected_defcal_measure_v1_q0 = textwrap.dedent(
+        """
+        defcal measure_v1 $0 -> bit {
+            play(tx_frame, constant(2400.0ns, 0.2));
+            return capture_v2(rx_frame, 2400.0ns);
+        }
+        """
+    ).strip()
+
+    assert (
+        dumps(prog.defcals[(("$0",), "measure_v1", ())], indent="    ").strip()
+        == expected_defcal_measure_v1_q0
+    )
+
+    expected_function_definition = textwrap.dedent(
+        """
+        def increment_variable_return(int[32] i) -> int[32] {
+            i += 1;
+            return i;
+        }
+        """
+    ).strip()
+    assert (
+        dumps(prog.subroutines["increment_variable_return"], indent="    ").strip()
+        == expected_function_definition
+    )
+
+
 def test_ramsey_example():
     prog = Program()
     constant = declare_waveform_generator("constant", [("length", duration), ("iq", complex128)])