fix: preserve barriers and static array outputs in comptime

guppylang-internals/src/guppylang_internals/definition/custom.py

@@ -380,6 +380,11 @@ class CustomCallChecker(ABC):
     # overridden, this controls the default behaviour of `compute_input_flags`.
     input_flag_mode: ClassVar[InputFlagDefaultMode] = InputFlagDefaultMode.RAISE
 
+    # Whether inout arguments are guaranteed to return static Python values unchanged.
+    # During tracing, this permits packed dynamic values to be updated while leaving
+    # any static Python leaves intact.
+    preserves_static_inout_values: ClassVar[bool] = False
+
     _depth = 0
 
     @contextmanager

guppylang-internals/src/guppylang_internals/std/_internal/checker.py

@@ -515,6 +515,14 @@ def to_sized_iter(
 class BarrierChecker(CustomCallChecker):
     """Call checker for the `barrier` function."""
 
+    preserves_static_inout_values = True
+
+    @override
+    def compute_input_flags(self, args: list[ast.expr]) -> list[InputFlags]:
+        # Unlike ordinary borrows, barriers forward copyable values too. This ensures
+        # that subsequent operations depend on the barrier's output wires.
+        return [InputFlags.Inout] * len(args)
+
     @override
     def synthesize(self, args: list[ast.expr]) -> tuple[ast.expr, Type]:
         tys = [ExprSynthesizer(self.ctx).synthesize(val)[1] for val in args]

guppylang-internals/src/guppylang_internals/std/_internal/compiler/platform.py

@@ -13,6 +13,7 @@
     CustomCallChecker,
     CustomCallCompiler,
     CustomInoutCallCompiler,
+    DefaultCallChecker,
 )
 from guppylang_internals.definition.value import CallReturnWires
 from guppylang_internals.diagnostic import Error, Help, Note
@@ -37,6 +38,12 @@
 TAG_MAX_LEN = 200
 
 
+class ArrayOutputChecker(DefaultCallChecker):
+    """Call checker for array outputs, which leave their borrowed array unchanged."""
+
+    preserves_static_inout_values = True
+
+
 @dataclass(frozen=True)
 class OutputTagTooLongError(Error):
     title: ClassVar[str] = "Tag too long"

guppylang-internals/src/guppylang_internals/tracing/function.py

@@ -283,6 +283,14 @@ def assign_var(obj: GuppyObject, arg: Any) -> GuppyObject:
     # Record call to compile later
     call = state.recorder.record_call(call_node, new_vars)
 
+    # Some custom inout calls, such as barriers, return their arguments unchanged.
+    # This lets tracing reconnect dynamic values to the call's output wires while
+    # leaving immutable static Python values inside packed arguments untouched.
+    preserves_static_inout_values = (
+        isinstance(resolved_func, CustomFunctionDef)
+        and resolved_func.call_checker.preserves_static_inout_values
+    )
+
     # Since all inputs are GuppyObjects identifying ComptimeVariables (varieties
     # of Place), ExprCompiler will update the Places to map to the output wires
     # of the call. Here we just write the GuppyObjects with those Places back to
@@ -291,7 +299,12 @@ def assign_var(obj: GuppyObject, arg: Any) -> GuppyObject:
         if InputFlags.Inout in flags:
             ty = arg_obj._ty
             # This marks `arg_obj` as used, but clears usedness of `val`, as desired:
-            success = update_packed_value(val, arg_obj, state.recorder)
+            success = update_packed_value(
+                val,
+                arg_obj,
+                state.recorder,
+                preserve_static=preserves_static_inout_values,
+            )
 
             if not success:
                 # This means the user has passed an object that we cannot update,

guppylang-internals/src/guppylang_internals/tracing/unpacking.py

@@ -161,7 +161,13 @@ def guppy_object_from_py(
             return GuppyObject(ty, recorder.record_load_val(v, ty, node))
 
 
-def update_packed_value(v: Any, obj: "GuppyObject", recorder: TraceRecorder) -> bool:
+def update_packed_value(
+    v: Any,
+    obj: "GuppyObject",
+    recorder: TraceRecorder,
+    *,
+    preserve_static: bool = False,
+) -> bool:
     """Given a Python value `v` and a `GuppyObject` `obj` that was constructed from `v`
     using `guppy_object_from_py`, tries to update the wires of any `GuppyObjects`
     contained in `v` to the new wires specified by `obj`.
@@ -170,6 +176,10 @@ def update_packed_value(v: Any, obj: "GuppyObject", recorder: TraceRecorder) ->
     the object available again since it now corresponds to a fresh wire.
 
     Returns `True` if all wires could be updated, otherwise `False`.
+
+    If `preserve_static` is true, Python values without backing wires are left
+    unchanged. This is only safe for inout calls that guarantee they return static
+    values unchanged, such as barriers.
     """
     match v:
         case GuppyObject() as v_obj:
@@ -179,13 +189,21 @@ def update_packed_value(v: Any, obj: "GuppyObject", recorder: TraceRecorder) ->
                 state = get_tracing_state()
                 state.unused_undroppable_objs[v_obj._id] = v_obj
             v_obj._used = None
+        case GuppyEnumObject(_ty=ty) as v_obj:
+            assert ty == obj._ty
+            object.__setattr__(v_obj, "_wire", obj._use_wire(None))
         case None:
             assert isinstance(obj._ty, NoneType)
         case tuple(vs):
             assert isinstance(obj._ty, TupleType)
             wires = recorder.record_untuple(obj._ty.element_types, obj._use_wire(None))
             for v, ty, out_wire in zip(vs, obj._ty.element_types, wires, strict=True):
-                success = update_packed_value(v, GuppyObject(ty, out_wire), recorder)
+                success = update_packed_value(
+                    v,
+                    GuppyObject(ty, out_wire),
+                    recorder,
+                    preserve_static=preserve_static,
+                )
                 if not success:
                     return False
         case GuppyStructObject(_ty=ty, _field_values=values):
@@ -196,7 +214,10 @@ def update_packed_value(v: Any, obj: "GuppyObject", recorder: TraceRecorder) ->
             for field, out_wire in zip(ty.fields, wires, strict=True):
                 v = values[field.name]
                 success = update_packed_value(
-                    v, GuppyObject(field.ty, out_wire), recorder
+                    v,
+                    GuppyObject(field.ty, out_wire),
+                    recorder,
+                    preserve_static=preserve_static,
                 )
                 if not success:
                     values[field.name] = obj
@@ -205,9 +226,17 @@ def update_packed_value(v: Any, obj: "GuppyObject", recorder: TraceRecorder) ->
             wires = unpack_array(recorder, obj._ty, obj._use_wire(None))
             elem_ty = get_element_type(obj._ty)
             for i, (v, wire) in enumerate(zip(vs, wires, strict=True)):
-                success = update_packed_value(v, GuppyObject(elem_ty, wire), recorder)
+                success = update_packed_value(
+                    v,
+                    GuppyObject(elem_ty, wire),
+                    recorder,
+                    preserve_static=preserve_static,
+                )
                 if not success:
                     vs[i] = obj
         case _:
+            if preserve_static:
+                obj._use_wire(None)
+                return True
             return False
     return True

guppylang/src/guppylang/std/platform.py

@@ -13,6 +13,7 @@
     BarrierChecker,
 )
 from guppylang_internals.std._internal.compiler.platform import (
+    ArrayOutputChecker,
     ArrayOutputCompiler,
     MeasurementOutputChecker,
     OutputCompiler,
@@ -52,19 +53,25 @@ def _output_float(tag: str @ comptime, value: float) -> None: ...
 def _output_measurement(tag: str @ comptime, value: Measurement) -> None: ...
 
 
-@custom_function(ArrayOutputCompiler("result_array_int", with_int_width=True))
+@custom_function(
+    ArrayOutputCompiler("result_array_int", with_int_width=True),
+    checker=ArrayOutputChecker(),
+)
 def _output_int_array(tag: str @ comptime, value: array[int, n]) -> None: ...
 
 
-@custom_function(ArrayOutputCompiler("result_array_uint", with_int_width=True))
+@custom_function(
+    ArrayOutputCompiler("result_array_uint", with_int_width=True),
+    checker=ArrayOutputChecker(),
+)
 def _output_nat_array(tag: str @ comptime, value: array[nat, n]) -> None: ...
 
 
-@custom_function(ArrayOutputCompiler("result_array_bool"))
+@custom_function(ArrayOutputCompiler("result_array_bool"), checker=ArrayOutputChecker())
 def _output_bool_array(tag: str @ comptime, value: array[bool, n]) -> None: ...
 
 
-@custom_function(ArrayOutputCompiler("result_array_f64"))
+@custom_function(ArrayOutputCompiler("result_array_f64"), checker=ArrayOutputChecker())
 def _output_float_array(tag: str @ comptime, value: array[float, n]) -> None: ...
 
 

tests/integration/test_output.py

@@ -118,6 +118,32 @@ def main() -> None:
     validate(main.compile_function())
 
 
+def test_comptime_static_array_preserved_after_output(validate):
+    @guppy.comptime
+    def main() -> int:
+        xs = array(1, 2)
+        output("xs", xs)
+
+        assert xs == [1, 2]
+        return xs[0] + xs[1]
+
+    validate(main.compile_function())
+
+
+def test_comptime_dynamic_array_preserved_after_output(validate):
+    @guppy
+    def add_one(x: int) -> int:
+        return x + 1
+
+    @guppy.comptime
+    def main(x: int) -> int:
+        xs = array(add_one(x), 2)
+        output("xs", xs)
+        return xs[0] + xs[1]
+
+    validate(main.compile_function())
+
+
 def test_deprecated_result_alias_still_compiles(validate):
     @compile_guppy
     def main(x: int) -> None:

tests/integration/test_quantum.py

@@ -2,6 +2,8 @@
 
 from typing import no_type_check
 
+from hugr import ops
+
 from guppylang.std.angles import angle
 
 from guppylang.std.builtins import owned, array, panic, output
@@ -252,3 +254,142 @@ def test() -> None:
         output("c2", measure(q1).read())
 
     validate(test.compile_function())
+
+
+def test_comptime_barrier(validate):
+    """Barrier borrows linear arguments during comptime tracing."""
+
+    @guppy.comptime
+    def test() -> None:
+        q1, q2, q3, q4 = qubit(), qubit(), qubit(), qubit()
+
+        q.h(q1)
+        q.h(q2)
+        barrier(q1, q2, q3)
+        q.h(q3)
+
+        q.cx(q1, q2)
+        barrier(q2, q3)
+        q.cx(q3, q4)
+
+        discard(q1)
+        discard(q2)
+        barrier()  # does nothing
+        discard(q3)
+        discard(q4)
+
+    validate(test.compile_function())
+
+
+def test_comptime_barrier_array(validate):
+    """Barrier on comptime array/struct access."""
+
+    @guppy.comptime
+    @no_type_check
+    def test() -> None:
+        qs = array(qubit() for _ in range(4))
+        q.h(qs[0])
+        q.h(qs[1])
+        barrier(qs[0], qs[1], qs[2])
+        barrier(qs[0])
+        q.h(qs[2])
+        barrier(*qs, 42)
+        q.cx(qs[0], qs[1])
+        barrier(qs[1], qs[2])
+        q.cx(qs[2], qs[3])
+        barrier(qs)
+        discard_array(qs)
+
+    validate(test.compile_function())
+
+
+def test_comptime_barrier_misc(validate):
+    """Comptime barrier on classical and non-place."""
+
+    @guppy.comptime
+    @no_type_check
+    def test() -> None:
+        q1 = qubit()
+        q.h(q1)
+        x = 1
+        barrier(q1, array(1, 2, 3), 2 + 3, x)
+
+        output("c", x)
+        output("c2", measure(q1).read())
+
+    validate(test.compile_function())
+
+
+def test_comptime_barrier_struct(validate):
+    """Barrier on array/struct access at comptime."""
+
+    @guppy.struct
+    class S:
+        q1: qubit
+        q2: qubit
+        q3: qubit
+        q4: qubit
+
+    @guppy.comptime
+    @no_type_check
+    def test() -> None:
+        qs = S(qubit(), qubit(), qubit(), qubit())
+        q.h(qs.q1)
+        q.h(qs.q2)
+        barrier(qs.q1, qs.q2, qs.q3)
+        barrier(qs.q1)
+        q.h(qs.q3)
+
+        q.cx(qs.q1, qs.q2)
+        barrier(qs.q2, qs.q3)
+        q.cx(qs.q3, qs.q4)
+
+        discard(qs.q1)
+        discard(qs.q2)
+        discard(qs.q3)
+        discard(qs.q4)
+
+    validate(test.compile_function())
+
+
+def test_comptime_barrier_dynamic_copyable_values(validate):
+    """Comptime barriers preserve dependencies for dynamic copyable values."""
+
+    @guppy
+    def add_one(x: int) -> int:
+        return x + 1
+
+    @guppy.struct
+    class S:
+        dynamic: int
+        static: int
+
+    @guppy.enum
+    class E:
+        Variant = {"value": int}
+
+    @guppy.comptime
+    def test(x: int) -> tuple[int, E]:
+        dynamic = add_one(x)
+        xs = array(add_one(x), add_one(x))
+        pair = (add_one(x), 42)
+        struct = S(add_one(x), 42)
+        enum = E.Variant(add_one(x))
+
+        barrier(dynamic)
+        barrier(xs)
+        barrier(pair)
+        barrier(struct)
+        barrier(enum)
+
+        return dynamic + xs[0] + pair[0] + struct.dynamic, enum
+
+    package = test.compile_function()
+    barriers = [
+        data.op
+        for _, data in package.modules[0].nodes()
+        if isinstance(data.op, ops.ExtOp)
+        and data.op.op_def().qualified_name() == "prelude.Barrier"
+    ]
+    assert len(barriers) == 5
+    validate(package)