feature: add barrier instruction (#1118)

Author: nilsquet

src/braket/circuits/circuit.py

@@ -729,6 +729,29 @@ def add_verbatim_box(
             self._has_compiler_directives = True
         return self
 
+    def barrier(self, target: QubitSetInput | None = None) -> Circuit:
+        """Add a barrier compiler directive to the circuit.
+
+        Args:
+            target (QubitSetInput | None): Target qubits for the barrier.
+            If None, applies to all qubits in the circuit.
+
+        Returns:
+            Circuit: self
+
+        Examples:
+            >>> circ = Circuit().h(0).barrier([0, 1]).cnot(0, 1)
+            >>> circ = Circuit().h(0).h(1).barrier()  # barrier on all qubits
+        """
+        target_qubits = self.qubits if target is None else QubitSet(target)
+
+        if target_qubits:
+            self.add_instruction(
+                Instruction(compiler_directives.Barrier(list(target_qubits)), target=target_qubits)
+            )
+            self._has_compiler_directives = True
+        return self
+
     def _add_measure(self, target_qubits: QubitSetInput) -> None:
         """Adds a measure instruction to the the circuit
 

src/braket/circuits/compiler_directives.py

@@ -16,6 +16,8 @@
 import braket.ir.jaqcd as ir
 
 from braket.circuits.compiler_directive import CompilerDirective
+from braket.circuits.serialization import IRType, SerializationProperties
+from braket.registers.qubit_set import QubitSet
 
 
 class StartVerbatimBox(CompilerDirective):
@@ -60,3 +62,36 @@ def _to_jaqcd(self, *args, **kwargs) -> Any:
 
     def _to_openqasm(self) -> str:
         return "}"
+
+
+class Barrier(CompilerDirective):
+    """Barrier compiler directive."""
+
+    def __init__(self, qubit_indices: list[int]):
+        super().__init__(["||"])
+        self._qubit_indices = qubit_indices
+
+    @property
+    def qubit_indices(self) -> list[int]:
+        return self._qubit_indices
+
+    @property
+    def qubit_count(self) -> int:
+        return len(self._qubit_indices)
+
+    def _to_jaqcd(self) -> Any:
+        raise NotImplementedError("Barrier is not supported in JAQCD")
+
+    def to_ir(
+        self,
+        target: QubitSet | None,
+        ir_type: IRType,
+        serialization_properties: SerializationProperties | None = None,
+        **kwargs,
+    ) -> Any:
+        if ir_type.name == "OPENQASM":
+            if target:
+                qubits = ", ".join(serialization_properties.format_target(int(q)) for q in target)
+                return f"barrier {qubits};"
+            return "barrier;"
+        return super().to_ir(target, ir_type, serialization_properties, **kwargs)

src/braket/circuits/moments.py

@@ -174,10 +174,17 @@ def add(self, instructions: Iterable[Instruction] | Instruction, noise_index: in
     def _add(self, instruction: Instruction, noise_index: int = 0) -> None:
         operator = instruction.operator
         if isinstance(operator, CompilerDirective):
-            time = self._update_qubit_times(self._qubits)
-            self._moments[MomentsKey(time, None, MomentType.COMPILER_DIRECTIVE, 0)] = instruction
+            qubit_range = instruction.target.union(instruction.control or QubitSet())
+            time = self._handle_compiler_directive(operator, qubit_range)
+            # For barriers without qubits, use empty qubit set for the key
+            key_qubits = (
+                QubitSet() if operator.name == "Barrier" and not qubit_range else qubit_range
+            )
+            self._moments[MomentsKey(time, key_qubits, MomentType.COMPILER_DIRECTIVE, 0)] = (
+                instruction
+            )
+            self._qubits.update(qubit_range)
             self._depth = time + 1
-            self._time_all_qubits = time
         elif isinstance(operator, Noise):
             self.add_noise(instruction)
         elif isinstance(operator, Gate) and operator.name == "GPhase":
@@ -282,6 +289,17 @@ def sort_moments(self) -> None:
 
         self._moments = sorted_moment
 
+    def _handle_compiler_directive(self, operator: CompilerDirective, qubit_range: QubitSet) -> int:
+        """Handle compiler directive and return the time slot."""
+        if operator.name == "Barrier" and not qubit_range:
+            time = self._get_qubit_times(self._qubits) + 1
+            self._time_all_qubits = time
+        else:
+            time = self._update_qubit_times(qubit_range or self._qubits)
+            if operator.name != "Barrier":
+                self._time_all_qubits = time
+        return time
+
     def _max_time_for_qubit(self, qubit: Qubit) -> int:
         # -1 if qubit is unoccupied because the first instruction will have an index of 0
         return self._max_times.get(qubit, -1)
@@ -307,7 +325,7 @@ def values(self) -> ValuesView[Instruction]:
         self.sort_moments()
         return self._moments.values()
 
-    def get(self, key: MomentsKey, default: Any | None = None) -> Instruction:
+    def get(self, key: MomentsKey, default: Any | None = None) -> Instruction | Any | None:
         """Get the instruction in self by key.
 
         Args:

src/braket/circuits/text_diagram_builders/ascii_circuit_diagram.py

@@ -70,6 +70,125 @@ def _duplicate_time_at_bottom(cls, lines: str) -> None:
         # duplicate times after an empty line
         lines.append(lines[0])
 
+    @classmethod
+    def _process_item_properties(
+        cls, item: Instruction | ResultType, circuit_qubits: QubitSet
+    ) -> tuple[QubitSet, QubitSet, QubitSet, QubitSet, list[str], dict | None]:
+        """Extract properties from an item, keeping original logic structure."""
+        if isinstance(item, ResultType) and not item.target:
+            target_qubits = circuit_qubits
+            control_qubits = QubitSet()
+            target_and_control = target_qubits.union(control_qubits)
+            qubits = circuit_qubits
+            ascii_symbols = [item.ascii_symbols[0]] * len(circuit_qubits)
+            map_control_qubit_states = None
+        elif isinstance(item, Instruction) and isinstance(item.operator, CompilerDirective):
+            if item.operator.name == "Barrier":
+                target_qubits = item.target
+                if not target_qubits:
+                    # Barrier without qubits - single barrier across all qubits
+                    target_qubits = circuit_qubits
+                    qubits = circuit_qubits
+                    ascii_symbols = [item.ascii_symbols[0]] * len(circuit_qubits)
+                else:
+                    # Barrier with specific qubits
+                    qubits = target_qubits
+                    ascii_symbols = [item.ascii_symbols[0]] * len(target_qubits)
+                target_and_control = target_qubits
+            else:
+                target_qubits = circuit_qubits
+                control_qubits = QubitSet()
+                target_and_control = target_qubits.union(control_qubits)
+                qubits = circuit_qubits
+                ascii_symbol = item.ascii_symbols[0]
+                marker = "*" * len(ascii_symbol)
+                num_after = len(circuit_qubits) - 1
+                after = ["|"] * (num_after - 1) + ([marker] if num_after else [])
+                ascii_symbols = [ascii_symbol, *after]
+            control_qubits = QubitSet()
+            map_control_qubit_states = None
+        elif (
+            isinstance(item, Instruction)
+            and isinstance(item.operator, Gate)
+            and item.operator.name == "GPhase"
+        ):
+            target_qubits = circuit_qubits
+            control_qubits = QubitSet()
+            target_and_control = QubitSet()
+            qubits = circuit_qubits
+            ascii_symbols = cls._qubit_line_character() * len(circuit_qubits)
+            map_control_qubit_states = None
+        else:
+            if isinstance(item.target, list):
+                target_qubits = reduce(QubitSet.union, map(QubitSet, item.target), QubitSet())
+            else:
+                target_qubits = item.target
+            control_qubits = getattr(item, "control", QubitSet())
+            control_state = getattr(item, "control_state", "1" * len(control_qubits))
+            map_control_qubit_states = dict(zip(control_qubits, control_state, strict=True))
+            target_and_control = target_qubits.union(control_qubits)
+            qubits = QubitSet(range(min(target_and_control), max(target_and_control) + 1))
+            ascii_symbols = item.ascii_symbols
+
+        return (
+            target_qubits,
+            control_qubits,
+            target_and_control,
+            qubits,
+            ascii_symbols,
+            map_control_qubit_states,
+        )
+
+    @classmethod
+    def _update_qubit_symbols_and_connections(
+        cls,
+        item: Instruction | ResultType,
+        qubit: int,
+        target_qubits: QubitSet,
+        control_qubits: QubitSet,
+        target_and_control: QubitSet,
+        ascii_symbols: list[str],
+        symbols: dict,
+        connections: dict,
+        map_control_qubit_states: dict | None,
+    ) -> None:
+        """Update symbols and connections for a qubit, keeping original logic."""
+        # Determine if the qubit is part of the item or in the middle of a
+        # multi qubit item.
+        if qubit in target_qubits:
+            item_qubit_index = next(index for index, q in enumerate(target_qubits) if q == qubit)
+            power_string = (
+                f"^{power}"
+                if (
+                    (power := getattr(item, "power", 1)) != 1
+                    # this has the limitation of not printing the power
+                    # when a user has a gate genuinely named C, but
+                    # is necessary to enable proper printing of custom
+                    # gates with built-in control qubits
+                    and ascii_symbols[item_qubit_index] != "C"
+                )
+                else ""
+            )
+            idx = item_qubit_index
+            symbols[qubit] = (
+                f"({ascii_symbols[idx]}{power_string})" if power_string else ascii_symbols[idx]
+            )
+        elif qubit in control_qubits:
+            symbols[qubit] = "C" if map_control_qubit_states[qubit] else "N"
+        else:
+            symbols[qubit] = "|"
+
+        # Set the margin to be a connector if not on the first qubit
+        if target_and_control and qubit != min(target_and_control):
+            is_barrier = (
+                isinstance(item, Instruction)
+                and isinstance(item.operator, CompilerDirective)
+                and item.operator.name == "Barrier"
+            )
+            # Add vertical lines for non-barriers or global barriers (no target)
+            if not is_barrier or not item.target:
+                connections[qubit] = "above"
+
     @classmethod
     def _create_diagram_column(
         cls,
@@ -91,78 +210,27 @@ def _create_diagram_column(
         connections = dict.fromkeys(circuit_qubits, "none")
 
         for item in items:
-            if isinstance(item, ResultType) and not item.target:
-                target_qubits = circuit_qubits
-                control_qubits = QubitSet()
-                target_and_control = target_qubits.union(control_qubits)
-                qubits = circuit_qubits
-                ascii_symbols = [item.ascii_symbols[0]] * len(circuit_qubits)
-            elif isinstance(item, Instruction) and isinstance(item.operator, CompilerDirective):
-                target_qubits = circuit_qubits
-                control_qubits = QubitSet()
-                target_and_control = target_qubits.union(control_qubits)
-                qubits = circuit_qubits
-                ascii_symbol = item.ascii_symbols[0]
-                marker = "*" * len(ascii_symbol)
-                num_after = len(circuit_qubits) - 1
-                after = ["|"] * (num_after - 1) + ([marker] if num_after else [])
-                ascii_symbols = [ascii_symbol, *after]
-            elif (
-                isinstance(item, Instruction)
-                and isinstance(item.operator, Gate)
-                and item.operator.name == "GPhase"
-            ):
-                target_qubits = circuit_qubits
-                control_qubits = QubitSet()
-                target_and_control = QubitSet()
-                qubits = circuit_qubits
-                ascii_symbols = cls._qubit_line_character() * len(circuit_qubits)
-            else:
-                if isinstance(item.target, list):
-                    target_qubits = reduce(QubitSet.union, map(QubitSet, item.target), QubitSet())
-                else:
-                    target_qubits = item.target
-                control_qubits = getattr(item, "control", QubitSet())
-                control_state = getattr(item, "control_state", "1" * len(control_qubits))
-                map_control_qubit_states = dict(zip(control_qubits, control_state, strict=True))
-
-                target_and_control = target_qubits.union(control_qubits)
-                qubits = QubitSet(range(min(target_and_control), max(target_and_control) + 1))
-
-                ascii_symbols = item.ascii_symbols
+            (
+                target_qubits,
+                control_qubits,
+                target_and_control,
+                qubits,
+                ascii_symbols,
+                map_control_qubit_states,
+            ) = cls._process_item_properties(item, circuit_qubits)
 
             for qubit in qubits:
-                # Determine if the qubit is part of the item or in the middle of a
-                # multi qubit item.
-                if qubit in target_qubits:
-                    item_qubit_index = [  # noqa: RUF015
-                        index for index, q in enumerate(target_qubits) if q == qubit
-                    ][0]
-                    power_string = (
-                        f"^{power}"
-                        if (
-                            (power := getattr(item, "power", 1)) != 1
-                            # this has the limitation of not printing the power
-                            # when a user has a gate genuinely named C, but
-                            # is necessary to enable proper printing of custom
-                            # gates with built-in control qubits
-                            and ascii_symbols[item_qubit_index] != "C"
-                        )
-                        else ""
-                    )
-                    symbols[qubit] = (
-                        f"({ascii_symbols[item_qubit_index]}{power_string})"
-                        if power_string
-                        else ascii_symbols[item_qubit_index]
-                    )
-                elif qubit in control_qubits:
-                    symbols[qubit] = "C" if map_control_qubit_states[qubit] else "N"
-                else:
-                    symbols[qubit] = "|"
-
-                # Set the margin to be a connector if not on the first qubit
-                if target_and_control and qubit != min(target_and_control):
-                    connections[qubit] = "above"
+                cls._update_qubit_symbols_and_connections(
+                    item,
+                    qubit,
+                    target_qubits,
+                    control_qubits,
+                    target_and_control,
+                    ascii_symbols,
+                    symbols,
+                    connections,
+                    map_control_qubit_states,
+                )
 
         return cls._create_output(symbols, connections, circuit_qubits, global_phase)