From 9fc2fa044301868f8527a1b554095797e88aa324 Mon Sep 17 00:00:00 2001
From: Ryan Shaffer <3620100+rmshaffer@users.noreply.github.com>
Date: Tue, 23 Apr 2024 12:00:44 -0400
Subject: [PATCH] Use BasisStateInput for control_state

---
 .../autoqasm/instructions/instructions.py     | 22 ++++++++---------
 .../autoqasm/test_gate_definitions.py         | 24 +++++++++++++++++--
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/src/braket/experimental/autoqasm/instructions/instructions.py b/src/braket/experimental/autoqasm/instructions/instructions.py
index 6545f351e..3411b1a71 100644
--- a/src/braket/experimental/autoqasm/instructions/instructions.py
+++ b/src/braket/experimental/autoqasm/instructions/instructions.py
@@ -21,6 +21,7 @@
 
 import oqpy
 
+from braket.circuits.basis_state import BasisState, BasisStateInput
 from braket.experimental.autoqasm import program as aq_program
 from braket.experimental.autoqasm import types as aq_types
 from braket.experimental.autoqasm.instructions.qubits import _qubit
@@ -33,7 +34,7 @@ def _qubit_instruction(
     *args: Any,
     is_unitary: bool = True,
     control: QubitIdentifierType | Iterable[QubitIdentifierType] | None = None,
-    control_state: str | None = None,
+    control_state: BasisStateInput | None = None,
     power: float | None = None,
 ) -> None:
     program_conversion_context = aq_program.get_program_conversion_context()
@@ -57,7 +58,7 @@ def _qubit_instruction(
 
 def _get_pos_neg_control(
     control: QubitIdentifierType | Iterable[QubitIdentifierType] | None = None,
-    control_state: str | None = None,
+    control_state: BasisStateInput | None = None,
 ) -> tuple[list[oqpy.Qubit], list[oqpy.Qubit]]:
     if control is None and control_state is not None:
         raise ValueError(control_state, "control_state provided without control qubits")
@@ -68,17 +69,16 @@ def _get_pos_neg_control(
     if aq_types.is_qubit_identifier_type(control):
         control = [control]
 
-    if control_state is not None and len(control) != len(control_state):
+    if control_state is None:
+        return [_qubit(q) for q in control], []
+
+    control_state = BasisState(control_state).as_tuple
+
+    if len(control) != len(control_state):
         raise ValueError(control_state, "control and control_state must have same length")
 
-    pos_control = [
-        _qubit(q) for i, q in enumerate(control) if control_state is None or control_state[i] == "1"
-    ]
-    neg_control = [
-        _qubit(q)
-        for i, q in enumerate(control)
-        if control_state is not None and control_state[i] == "0"
-    ]
+    pos_control = [_qubit(q) for i, q in enumerate(control) if control_state[i] == 1]
+    neg_control = [_qubit(q) for i, q in enumerate(control) if control_state[i] == 0]
     return pos_control, neg_control
 
 
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_gate_definitions.py b/test/unit_tests/braket/experimental/autoqasm/test_gate_definitions.py
index 3bc1e2dac..7dcb51e82 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_gate_definitions.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_gate_definitions.py
@@ -157,12 +157,14 @@ def test_gates(gate, qubits, params, expected_qasm) -> None:
         (x, [1], [], [0], None, None, "\nctrl @ x __qubits__[0], __qubits__[1];"),
         (x, [1], [], 0, "1", None, "\nctrl @ x __qubits__[0], __qubits__[1];"),
         (x, [1], [], [0], "0", None, "\nnegctrl @ x __qubits__[0], __qubits__[1];"),
+        (x, [1], [], [0], 0, None, "\nnegctrl @ x __qubits__[0], __qubits__[1];"),
+        (x, [1], [], [0], [0], None, "\nnegctrl @ x __qubits__[0], __qubits__[1];"),
         (
             x,
             [2],
             [],
             [0, 1],
-            "11",
+            "11",  # BasisStateInput as str
             None,
             "\nctrl(2) @ x __qubits__[0], __qubits__[1], __qubits__[2];",
         ),
@@ -171,7 +173,25 @@ def test_gates(gate, qubits, params, expected_qasm) -> None:
             [2],
             [],
             [0, 1],
-            "10",
+            [1, 1],  # BasisStateInput as list[int]
+            None,
+            "\nctrl(2) @ x __qubits__[0], __qubits__[1], __qubits__[2];",
+        ),
+        (
+            x,
+            [2],
+            [],
+            [0, 1],
+            3,  # BasisStateInput as int
+            None,
+            "\nctrl(2) @ x __qubits__[0], __qubits__[1], __qubits__[2];",
+        ),
+        (
+            x,
+            [2],
+            [],
+            [0, 1],
+            "10",  # BasisStateInput as str
             None,
             "\nctrl @ negctrl @ x __qubits__[0], __qubits__[1], __qubits__[2];",
         ),