Add support for capture_quantum_state

backend/src/lib.rs

@@ -18,6 +18,7 @@ mod simulator;
 
 use bitvec::prelude::*;
 use nearly_zero::NearlyZero;
+use num_bigint::BigUint;
 use num_complex::Complex64;
 use simulator::QuantumSim;
 use std::cell::RefCell;
@@ -946,6 +947,15 @@ pub extern "C" fn __quantum__rt__qubit_to_string(qubit: *mut c_void) -> *const C
     }
 }
 
+/// Rust API for getting a snapshot of current quantum state.
+#[must_use]
+pub fn capture_quantum_state() -> Vec<(BigUint, Complex64)> {
+    SIM_STATE.with(|sim_state| {
+        let mut state = sim_state.borrow_mut();
+        state.sim.get_state()
+    })
+}
+
 /// QIR API for dumping full internal simulator state.
 #[no_mangle]
 pub extern "C" fn __quantum__qis__dumpmachine__body(location: *mut c_void) {
@@ -975,10 +985,11 @@ mod tests {
         __quantum__qis__s__adj, __quantum__qis__s__body, __quantum__qis__x__body,
         __quantum__rt__qubit_allocate, __quantum__rt__qubit_allocate_array,
         __quantum__rt__qubit_release, __quantum__rt__qubit_release_array,
-        __quantum__rt__result_equal, map_to_z_basis, qubit_is_zero,
+        __quantum__rt__result_equal, capture_quantum_state, map_to_z_basis, qubit_is_zero,
         result_bool::__quantum__rt__result_get_one, result_bool::__quantum__rt__result_get_zero,
         unmap_from_z_basis, SIM_STATE,
     };
+    use num_bigint::BigUint;
     use qir_stdlib::{
         arrays::{
             __quantum__rt__array_create_1d, __quantum__rt__array_get_element_ptr_1d,
@@ -1262,4 +1273,25 @@ mod tests {
         ));
         assert!(qubit_is_zero(qubit));
     }
+
+    #[test]
+    fn test_capture_quantum_state() {
+        let qubit = __quantum__rt__qubit_allocate();
+        let state = capture_quantum_state();
+        assert_eq!(state.len(), 1);
+        assert_eq!(state[0].0, BigUint::from(0u32));
+        __quantum__qis__x__body(qubit);
+        let state = capture_quantum_state();
+        assert_eq!(state.len(), 1);
+        assert_eq!(state[0].0, BigUint::from(1u32));
+        __quantum__qis__h__body(qubit);
+        let state = capture_quantum_state();
+        assert_eq!(state.len(), 2);
+        assert_eq!(state[0].1, -state[1].1);
+        __quantum__qis__h__body(qubit);
+        __quantum__qis__x__body(qubit);
+        let state = capture_quantum_state();
+        assert_eq!(state.len(), 1);
+        assert_eq!(state[0].0, BigUint::from(0u32));
+    }
 }

backend/src/simulator.rs

@@ -69,6 +69,31 @@ impl QuantumSim {
         }
     }
 
+    #[must_use]
+    pub(crate) fn get_state(&mut self) -> Vec<(BigUint, Complex64)> {
+        // Swap all the entries in the state to be ordered by qubit identifier. This makes
+        // interpreting the state easier for external consumers that don't have access to the id map.
+        let mut sorted_keys: Vec<usize> = self.id_map.keys().copied().collect();
+        self.flush_queue(&sorted_keys, FlushLevel::HRxRy);
+
+        sorted_keys.sort_unstable();
+        sorted_keys.iter().enumerate().for_each(|(index, &key)| {
+            if index != self.id_map[&key] {
+                self.swap_qubit_state(self.id_map[&key], index);
+                let swapped_key = *self
+                    .id_map
+                    .iter()
+                    .find(|(_, &value)| value == index)
+                    .unwrap()
+                    .0;
+                *(self.id_map.get_mut(&swapped_key).unwrap()) = self.id_map[&key];
+                *(self.id_map.get_mut(&key).unwrap()) = index;
+            }
+        });
+
+        self.state.clone().drain().collect()
+    }
+
     /// Allocates a fresh qubit, returning its identifier. Note that this will use the lowest available
     /// identifier, and may result in qubits being allocated "in the middle" of an existing register
     /// if those identifiers are available.