Skip to main content

quil/program/
frame.rs

1use std::borrow::Borrow;
2use std::collections::{HashMap, HashSet};
3use std::convert::AsRef;
4
5use quil_rs::{
6    instruction::{FrameAttributes, FrameIdentifier},
7    program::FrameSet,
8};
9use rigetti_pyo3::{
10    impl_as_mut_for_wrapper, impl_repr, py_wrap_type,
11    pyo3::{pymethods, PyResult, Python},
12    PyTryFrom, PyWrapper, PyWrapperMut, ToPython,
13};
14
15use crate::{
16    impl_eq,
17    instruction::{PyFrameAttributes, PyFrameIdentifier, PyInstruction},
18};
19
20py_wrap_type! {
21    #[derive(Debug, PartialEq, Eq)]
22    PyFrameSet(FrameSet) as "FrameSet"
23}
24impl_repr!(PyFrameSet);
25impl_as_mut_for_wrapper!(PyFrameSet);
26impl_eq!(PyFrameSet);
27
28#[pymethods]
29impl PyFrameSet {
30    #[new]
31    pub fn new() -> Self {
32        Self(FrameSet::new())
33    }
34
35    pub fn get(
36        &self,
37        py: Python<'_>,
38        identifier: PyFrameIdentifier,
39    ) -> PyResult<Option<PyFrameAttributes>> {
40        self.as_inner()
41            .get(&FrameIdentifier::py_try_from(py, &identifier)?)
42            .to_python(py)
43    }
44
45    pub fn get_keys(&self, py: Python<'_>) -> PyResult<Vec<PyFrameIdentifier>> {
46        self.as_inner().get_keys().to_python(py)
47    }
48
49    pub fn get_all_frames(
50        &self,
51        py: Python<'_>,
52    ) -> PyResult<HashMap<PyFrameIdentifier, PyFrameAttributes>> {
53        self.as_inner()
54            .iter()
55            .map(|(ident, attribs)| Ok((ident.to_python(py)?, attribs.to_python(py)?)))
56            .collect()
57    }
58
59    pub fn insert(
60        &mut self,
61        py: Python<'_>,
62        identifier: PyFrameIdentifier,
63        attributes: PyFrameAttributes,
64    ) -> PyResult<()> {
65        self.as_inner_mut().insert(
66            FrameIdentifier::py_try_from(py, &identifier)?,
67            FrameAttributes::py_try_from(py, &attributes)?,
68        );
69        Ok(())
70    }
71
72    pub fn merge(&mut self, py: Python<'_>, other: PyFrameSet) -> PyResult<()> {
73        self.as_inner_mut()
74            .merge(FrameSet::py_try_from(py, &other)?);
75        Ok(())
76    }
77
78    pub fn intersection(
79        &self,
80        py: Python<'_>,
81        identifiers: HashSet<PyFrameIdentifier>,
82    ) -> PyResult<Self> {
83        Ok(Self(
84            self.as_inner().intersection(
85                &HashSet::<FrameIdentifier>::py_try_from(py, &identifiers)?
86                    .iter()
87                    .map(Borrow::borrow)
88                    .collect(),
89            ),
90        ))
91    }
92
93    pub fn is_empty(&self) -> bool {
94        self.as_inner().is_empty()
95    }
96
97    pub fn to_instructions(&self, py: Python<'_>) -> PyResult<Vec<PyInstruction>> {
98        self.as_inner().to_instructions().to_python(py)
99    }
100
101    pub fn __len__(&self) -> usize {
102        self.as_inner().len()
103    }
104}
105
106impl Default for PyFrameSet {
107    fn default() -> Self {
108        Self::new()
109    }
110}