feat: Add memory_values to QPUResultData (#393)

* feat: Add memory_values to QPUResultData

* implement python half

* clippy

* make memory accessible on py execution result

* add missing import

* dont wrap memory_values in Option

* update type hints

* clippy

* update py tests

* fix test

* update optional type on executionresult

crates/lib/src/execution_data.rs

@@ -352,19 +352,51 @@ mod describe_register_map {
     use maplit::hashmap;
     use ndarray::prelude::*;
 
+    use crate::qpu::result_data::MemoryValues;
     use crate::qpu::QpuResultData;
     use crate::qvm::QvmResultData;
 
     use super::{RegisterData, RegisterMap};
     use qcs_api_client_grpc::models::controller::readout_values::Values;
-    use qcs_api_client_grpc::models::controller::{IntegerReadoutValues, ReadoutValues};
+    use qcs_api_client_grpc::models::controller::{
+        self, BinaryDataValue, DataValue as ControllerMemoryValue, IntegerDataValue,
+        IntegerReadoutValues, ReadoutValues, RealDataValue,
+    };
 
     fn dummy_readout_values(v: Vec<i32>) -> ReadoutValues {
         ReadoutValues {
             values: Some(Values::IntegerValues(IntegerReadoutValues { values: v })),
         }
     }
 
+    fn dummy_memory_values(
+        binary: Vec<u8>,
+        int: Vec<i64>,
+        real: Vec<f64>,
+    ) -> (
+        ControllerMemoryValue,
+        ControllerMemoryValue,
+        ControllerMemoryValue,
+    ) {
+        (
+            ControllerMemoryValue {
+                value: Some(controller::data_value::Value::Binary(BinaryDataValue {
+                    data: binary,
+                })),
+            },
+            ControllerMemoryValue {
+                value: Some(controller::data_value::Value::Integer(IntegerDataValue {
+                    data: int,
+                })),
+            },
+            ControllerMemoryValue {
+                value: Some(controller::data_value::Value::Real(RealDataValue {
+                    data: real,
+                })),
+            },
+        )
+    }
+
     #[test]
     fn it_converts_rectangular_qpu_result_data_to_register_map() {
         let readout_mappings = hashmap! {
@@ -383,8 +415,19 @@ mod describe_register_map {
             String::from("qE") => dummy_readout_values(vec![2, 3]),
         };
 
-        let qpu_result_data =
-            QpuResultData::from_controller_mappings_and_values(&readout_mappings, &readout_values);
+        let (binary_values, integer_values, real_values) =
+            dummy_memory_values(vec![0, 1, 2], vec![3, 4, 5], vec![6.0, 7.0, 8.0]);
+        let memory_values = hashmap! {
+            String::from("binary") => binary_values,
+            String::from("int") => integer_values,
+            String::from("real") => real_values,
+        };
+
+        let qpu_result_data = QpuResultData::from_controller_mappings_and_values(
+            &readout_mappings,
+            &readout_values,
+            &memory_values,
+        );
 
         let register_map = RegisterMap::from_qpu_result_data(&qpu_result_data)
             .expect("Should be able to create RegisterMap from rectangular QPU readout");
@@ -408,6 +451,13 @@ mod describe_register_map {
         let expected_bar = arr2(&[[2, 0], [3, 1]]);
 
         assert_eq!(bar, expected_bar);
+
+        let expected_memory_values = hashmap! {
+            String::from("binary") => MemoryValues::Binary(vec![0, 1, 2]),
+            String::from("int") => MemoryValues::Integer(vec![3, 4, 5]),
+            String::from("real") => MemoryValues::Real(vec![6.0, 7.0, 8.0]),
+        };
+        assert_eq!(qpu_result_data.memory_values, expected_memory_values);
     }
 
     #[test]
@@ -427,8 +477,11 @@ mod describe_register_map {
             String::from("qE") => dummy_readout_values(vec![44]),
         };
 
-        let qpu_result_data =
-            QpuResultData::from_controller_mappings_and_values(&readout_mappings, &readout_values);
+        let qpu_result_data = QpuResultData::from_controller_mappings_and_values(
+            &readout_mappings,
+            &readout_values,
+            &hashmap! {},
+        );
 
         RegisterMap::from_qpu_result_data(&qpu_result_data)
             .expect_err("Should not be able to create RegisterMap from QPU readout with missing indices for a register");
@@ -448,8 +501,19 @@ mod describe_register_map {
             String::from("qC") => dummy_readout_values(vec![3]),
         };
 
-        let qpu_result_data =
-            QpuResultData::from_controller_mappings_and_values(&readout_mappings, &readout_values);
+        let (binary_values, integer_values, real_values) =
+            dummy_memory_values(vec![0, 1, 2], vec![3, 4, 5], vec![6.0, 7.0, 8.0]);
+        let memory_values = hashmap! {
+            String::from("binary") => binary_values,
+            String::from("int") => integer_values,
+            String::from("real") => real_values,
+        };
+
+        let qpu_result_data = QpuResultData::from_controller_mappings_and_values(
+            &readout_mappings,
+            &readout_values,
+            &memory_values,
+        );
 
         RegisterMap::from_qpu_result_data(&qpu_result_data)
             .expect_err("Should not be able to create RegisterMap from QPU readout with jagged data for a register");

crates/lib/src/qpu/execution.rs

@@ -283,6 +283,7 @@ impl<'a> Execution<'a> {
             result_data: ResultData::Qpu(QpuResultData::from_controller_mappings_and_values(
                 job_handle.readout_map(),
                 &response.readout_values,
+                &response.memory_values,
             )),
             duration: Some(Duration::from_micros(
                 response.execution_duration_microseconds,

crates/lib/src/qpu/result_data.rs

@@ -6,7 +6,8 @@ use quil_rs::instruction::MemoryReference;
 use std::collections::HashMap;
 
 use qcs_api_client_grpc::models::controller::{
-    readout_values as controller_readout_values, ReadoutValues as ControllerReadoutValues,
+    self, data_value as controller_memory_value, readout_values as controller_readout_values,
+    DataValue as ControllerMemoryValues, ReadoutValues as ControllerReadoutValues,
 };
 
 /// A row of readout values from the QPU. Each row contains all the values emitted to a
@@ -21,12 +22,25 @@ pub enum ReadoutValues {
     Complex(Vec<Complex64>),
 }
 
+/// A row of data containing the contents of each memory region at the end of a job.
+#[derive(Debug, Clone, EnumAsInner, PartialEq)]
+pub enum MemoryValues {
+    /// Values that correspond to a memory region declared with the BIT or OCTET data type.
+    Binary(Vec<u8>),
+    /// Values that correspond to a memory region declared with the INTEGER data type.
+    Integer(Vec<i64>),
+    /// Values that correspond to a memory region declared with the REAL data type.
+    Real(Vec<f64>),
+}
+
 /// This struct encapsulates data returned from the QPU after executing a job.
 #[allow(clippy::module_name_repetitions)]
 #[derive(Debug, Clone, PartialEq)]
 pub struct QpuResultData {
     pub(crate) mappings: HashMap<String, String>,
     pub(crate) readout_values: HashMap<String, ReadoutValues>,
+    /// The final contents of each memory region, keyed on region name.
+    pub(crate) memory_values: HashMap<String, MemoryValues>,
 }
 
 impl QpuResultData {
@@ -36,21 +50,24 @@ impl QpuResultData {
     pub fn from_mappings_and_values(
         mappings: HashMap<String, String>,
         readout_values: HashMap<String, ReadoutValues>,
+        memory_values: HashMap<String, MemoryValues>,
     ) -> Self {
         Self {
             mappings,
             readout_values,
+            memory_values,
         }
     }
 
     /// Creates a new [`QpuResultData`] using data returned from controller service.
     pub(crate) fn from_controller_mappings_and_values(
         mappings: &HashMap<String, String>,
-        values: &HashMap<String, ControllerReadoutValues>,
+        readout_values: &HashMap<String, ControllerReadoutValues>,
+        memory_values: &HashMap<String, ControllerMemoryValues>,
     ) -> Self {
         Self {
             mappings: mappings.clone(),
-            readout_values: values
+            readout_values: readout_values
                 .iter()
                 .map(|(key, readout_values)| {
                     (
@@ -74,6 +91,27 @@ impl QpuResultData {
                     )
                 })
                 .collect(),
+            memory_values: memory_values
+                .iter()
+                .filter_map(|(key, memory_values)| {
+                    memory_values.value.as_ref().map(|value| {
+                        (
+                            key.clone(),
+                            match value {
+                                controller_memory_value::Value::Binary(
+                                    controller::BinaryDataValue { data: v },
+                                ) => MemoryValues::Binary(v.clone()),
+                                controller_memory_value::Value::Integer(
+                                    controller::IntegerDataValue { data: v },
+                                ) => MemoryValues::Integer(v.clone()),
+                                controller_memory_value::Value::Real(
+                                    controller::RealDataValue { data: v },
+                                ) => MemoryValues::Real(v.clone()),
+                            },
+                        )
+                    })
+                })
+                .collect(),
         }
     }
 
@@ -100,4 +138,10 @@ impl QpuResultData {
     pub fn readout_values(&self) -> &HashMap<String, ReadoutValues> {
         &self.readout_values
     }
+
+    /// Get mapping of a memory region (ie. "ro") to the final contents of that memory region.
+    #[must_use]
+    pub fn memory_values(&self) -> &HashMap<String, MemoryValues> {
+        &self.memory_values
+    }
 }

crates/python/qcs_sdk/qpu/__init__.pyi

@@ -1,4 +1,4 @@
-from typing import Dict, List, Mapping, Sequence, Optional, Union, final, Any
+from typing import Dict, List, Mapping, Sequence, Optional, Union, final
 
 from qcs_sdk.client import QCSClient
 
@@ -51,6 +51,46 @@ class ReadoutValues:
     @staticmethod
     def from_complex(inner: Sequence[complex]) -> "ReadoutValues": ...
 
+@final
+class MemoryValues:
+    """
+    A row of data containing the contents of a memory region at the end of a job. There
+    is a variant for each possible type the memory region could be.
+
+    Variants:
+        - ``binary``: Corresponds to the Quil `BIT` and `OCTET` types.
+        - ``integer``: Corresponds to the Quil `INTEGER` types.
+        - ``real``: Corresponds to the Quil `REAL` type.
+
+    Methods (each per variant):
+        - ``is_*``: if the underlying values are that type.
+        - ``as_*``: if the underlying values are that type, then those values, otherwise ``None``.
+        - ``to_*``: the underlying values as that type, raises ``ValueError`` if they are not.
+        - ``from_*``: wrap underlying values as this enum type.
+    """
+    def __new__(cls, values: Union[List[int], List[float]]):
+        """Construct a new ReadoutValues from a list of values."""
+        ...
+    def inner(self) -> Union[List[int], List[float]]:
+        """Return the inner list of readout values."""
+        ...
+    def is_integer(self) -> bool: ...
+    def is_binary(self) -> bool: ...
+    def is_real(self) -> bool: ...
+    def as_integer(self) -> Optional[List[int]]: ...
+    def as_binary(self) -> Optional[List[int]]: ...
+    def as_real(self) -> Optional[List[float]]: ...
+    def to_integer(self) -> List[int]: ...
+    def to_binary(self) -> List[int]: ...
+    def to_real(self) -> List[float]: ...
+    @staticmethod
+    def from_integer(inner: Sequence[int]) -> "ReadoutValues": ...
+    @staticmethod
+    def from_binary(inner: Sequence[int]) -> "ReadoutValues": ...
+    @staticmethod
+    def from_real(inner: Sequence[float]) -> "ReadoutValues": ...
+
+
 @final
 class QPUResultData:
     """
@@ -62,7 +102,7 @@ class QPUResultData:
     across all shots.
     """
 
-    def __new__(cls, mappings: Mapping[str, str], readout_values: Mapping[str, ReadoutValues]): ...
+    def __new__(cls, mappings: Mapping[str, str], readout_values: Mapping[str, ReadoutValues], memory_values: Mapping[str, MemoryValues]): ...
     @property
     def mappings(self) -> Dict[str, str]:
         """
@@ -75,13 +115,19 @@ class QPUResultData:
         Get the mappings of a readout values identifier (ie. "q0") to a set of ``ReadoutValues``
         """
         ...
+    @property
+    def memory_values(self) -> Dict[str, MemoryValues]:
+        """
+        Get mapping of a memory region (ie. "ro") to the final contents of that memory region.
+        """
+        ...
     def to_raw_readout_data(
         self,
     ) -> RawQPUReadoutData:
         """
         Get a copy of this result data flattened into a ``RawQPUReadoutData``. This reduces
         the contained data down to primitive types, offering a simpler structure at the
-        cost of the type safety provided by ``ReadoutValues``.
+        cost of the type safety provided by ``ReadoutValues`` and ``MemoryValues``.
         """
         ...
 
@@ -105,6 +151,12 @@ class RawQPUReadoutData:
         Get the mappings of a readout values identifier (ie. "q0") to a list of those readout values
         """
         ...
+    @property
+    def memory_values(self) -> Dict[str, Optional[Union[List[int], List[float]]]]:
+        """
+        Get mapping of a memory region (ie. "ro") to the final contents of that memory region.
+        """
+        ...
 
 class ListQuantumProcessorsError(RuntimeError):
     """A request to list available Quantum Processors failed."""

crates/python/qcs_sdk/qpu/api.pyi

@@ -1,7 +1,7 @@
-from enum import Enum
-from typing import Callable, Dict, List, Sequence, Mapping, Optional, Union, final
+from typing import Dict, List, Sequence, Mapping, Optional, Union, final
 
 from qcs_sdk.client import QCSClient
+from qcs_sdk.qpu import MemoryValues
 
 class SubmissionError(RuntimeError):
     """There was a problem submitting the program to QCS for execution."""
@@ -77,6 +77,11 @@ class ExecutionResults:
         """
         ...
     @property
+    def memory(self) -> Dict[str, MemoryValues]:
+        """
+        The final state of memory for parameters that were read from and written to during the exectuion of the program.
+        """
+    @property
     def execution_duration_microseconds(self) -> Optional[int]:
         """The time spent executing the program."""
         ...

crates/python/src/execution_data.rs

@@ -21,7 +21,7 @@ py_wrap_union_enum! {
         qvm: Qvm => PyQvmResultData
     }
 }
-impl_repr!(PyQpuResultData);
+impl_repr!(PyResultData);
 
 wrap_error!(RustRegisterMatrixConversionError(
     qcs::RegisterMatrixConversionError