1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use std::borrow::Borrow;
use std::collections::{HashMap, HashSet};
use std::convert::AsRef;

use quil_rs::{
    instruction::{FrameAttributes, FrameIdentifier},
    program::FrameSet,
};
use rigetti_pyo3::{
    impl_as_mut_for_wrapper, impl_repr, py_wrap_type,
    pyo3::{pymethods, PyResult, Python},
    PyTryFrom, PyWrapper, PyWrapperMut, ToPython,
};

use crate::{
    impl_eq,
    instruction::{PyFrameAttributes, PyFrameIdentifier, PyInstruction},
};

py_wrap_type! {
    #[derive(Debug, PartialEq, Eq)]
    PyFrameSet(FrameSet) as "FrameSet"
}
impl_repr!(PyFrameSet);
impl_as_mut_for_wrapper!(PyFrameSet);
impl_eq!(PyFrameSet);

#[pymethods]
impl PyFrameSet {
    #[new]
    pub fn new() -> Self {
        Self(FrameSet::new())
    }

    pub fn get(
        &self,
        py: Python<'_>,
        identifier: PyFrameIdentifier,
    ) -> PyResult<Option<PyFrameAttributes>> {
        self.as_inner()
            .get(&FrameIdentifier::py_try_from(py, &identifier)?)
            .to_python(py)
    }

    pub fn get_keys(&self, py: Python<'_>) -> PyResult<Vec<PyFrameIdentifier>> {
        self.as_inner().get_keys().to_python(py)
    }

    pub fn get_all_frames(
        &self,
        py: Python<'_>,
    ) -> PyResult<HashMap<PyFrameIdentifier, PyFrameAttributes>> {
        self.as_inner()
            .iter()
            .map(|(ident, attribs)| Ok((ident.to_python(py)?, attribs.to_python(py)?)))
            .collect()
    }

    pub fn insert(
        &mut self,
        py: Python<'_>,
        identifier: PyFrameIdentifier,
        attributes: PyFrameAttributes,
    ) -> PyResult<()> {
        self.as_inner_mut().insert(
            FrameIdentifier::py_try_from(py, &identifier)?,
            FrameAttributes::py_try_from(py, &attributes)?,
        );
        Ok(())
    }

    pub fn merge(&mut self, py: Python<'_>, other: PyFrameSet) -> PyResult<()> {
        self.as_inner_mut()
            .merge(FrameSet::py_try_from(py, &other)?);
        Ok(())
    }

    pub fn intersection(
        &self,
        py: Python<'_>,
        identifiers: HashSet<PyFrameIdentifier>,
    ) -> PyResult<Self> {
        Ok(Self(
            self.as_inner().intersection(
                &HashSet::<FrameIdentifier>::py_try_from(py, &identifiers)?
                    .iter()
                    .map(Borrow::borrow)
                    .collect(),
            ),
        ))
    }

    pub fn is_empty(&self) -> bool {
        self.as_inner().is_empty()
    }

    pub fn to_instructions(&self, py: Python<'_>) -> PyResult<Vec<PyInstruction>> {
        self.as_inner().to_instructions().to_python(py)
    }

    pub fn __len__(&self) -> usize {
        self.as_inner().len()
    }
}

impl Default for PyFrameSet {
    fn default() -> Self {
        Self::new()
    }
}