1use crate::iterators::{
2 ControlledOpIterator, MatrixOpIterator, SparseMatrixOpIterator, SwapOpIterator,
3};
4use num_traits::{One, Zero};
5use std::fmt;
6use std::iter::Sum;
7use std::ops::Mul;
8
9#[derive(Clone)]
11pub enum MatrixOp<P> {
12 Matrix(Vec<usize>, Vec<P>),
14 SparseMatrix(Vec<usize>, Vec<Vec<(usize, P)>>),
16 Swap(Vec<usize>, Vec<usize>),
18 Control(Vec<usize>, Vec<usize>, Box<MatrixOp<P>>),
20}
21
22impl<P> MatrixOp<P> {
23 pub fn num_indices(&self) -> usize {
25 match self {
26 MatrixOp::Matrix(indices, _) => indices.len(),
27 MatrixOp::SparseMatrix(indices, _) => indices.len(),
28 MatrixOp::Swap(a, b) => a.len() + b.len(),
29 MatrixOp::Control(cs, os, _) => cs.len() + os.len(),
30 }
31 }
32
33 pub fn new_matrix<Indx, Dat>(indices: Indx, data: Dat) -> Self
35 where
36 Indx: Into<Vec<usize>>,
37 Dat: Into<Vec<P>>,
38 {
39 Self::Matrix(indices.into(), data.into())
40 }
41
42 pub fn new_sparse<Indx, Dat>(indices: Indx, data: Dat) -> Self
44 where
45 Indx: Into<Vec<usize>>,
46 Dat: Into<Vec<Vec<(usize, P)>>>,
47 {
48 Self::SparseMatrix(indices.into(), data.into())
49 }
50
51 pub fn new_swap<IndxA, IndxB>(a: IndxA, b: IndxB) -> Self
53 where
54 IndxA: Into<Vec<usize>>,
55 IndxB: Into<Vec<usize>>,
56 {
57 Self::Swap(a.into(), b.into())
58 }
59
60 pub fn new_control<IndxA, IndxB>(c: IndxA, r: IndxB, op: MatrixOp<P>) -> Self
62 where
63 IndxA: Into<Vec<usize>>,
64 IndxB: Into<Vec<usize>>,
65 {
66 Self::Control(c.into(), r.into(), Box::new(op))
67 }
68}
69
70impl<P> MatrixOp<P>
71where
72 P: Clone + Zero + One + Mul,
73{
74 pub fn sum_for_op_cols<T, F>(&self, nindices: usize, row: usize, f: F) -> T
77 where
78 T: Sum,
79 F: Fn((usize, P)) -> T,
80 {
81 match &self {
82 MatrixOp::Matrix(_, data) => MatrixOpIterator::new(row, nindices, data).map(f).sum(),
83 MatrixOp::SparseMatrix(_, data) => SparseMatrixOpIterator::new(row, data.as_slice())
84 .map(f)
85 .sum(),
86 MatrixOp::Swap(_, _) => SwapOpIterator::new(row, nindices).map(f).sum(),
87 MatrixOp::Control(c_indices, o_indices, op) => {
88 let n_control_indices = c_indices.len();
89 let n_op_indices = o_indices.len();
90 op.sum_for_control_iterator(row, n_control_indices, n_op_indices, f)
91 }
92 }
93 }
94
95 fn sum_for_control_iterator<T, F>(
96 &self,
97 row: usize,
98 n_control_indices: usize,
99 n_op_indices: usize,
100 f: F,
101 ) -> T
102 where
103 T: Sum,
104 F: Fn((usize, P)) -> T,
105 {
106 match &self {
107 MatrixOp::Matrix(_, data) => {
108 let iter_builder = |row: usize| MatrixOpIterator::new(row, n_op_indices, data);
109 ControlledOpIterator::new(row, n_control_indices, n_op_indices, iter_builder)
110 .map(f)
111 .sum()
112 }
113 MatrixOp::SparseMatrix(_, data) => {
114 let iter_builder = |row: usize| SparseMatrixOpIterator::new(row, data);
115 ControlledOpIterator::new(row, n_control_indices, n_op_indices, iter_builder)
116 .map(f)
117 .sum()
118 }
119 MatrixOp::Swap(_, _) => {
120 let iter_builder = |row: usize| SwapOpIterator::new(row, n_op_indices);
121 ControlledOpIterator::new(row, n_control_indices, n_op_indices, iter_builder)
122 .map(f)
123 .sum()
124 }
125 MatrixOp::Control(c_indices, o_indices, op) => {
128 let n_control_indices = n_control_indices + c_indices.len();
129 let n_op_indices = o_indices.len();
130 op.sum_for_control_iterator(row, n_control_indices, n_op_indices, f)
131 }
132 }
133 }
134}
135
136impl<P> fmt::Debug for MatrixOp<P> {
137 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
138 let (name, indices) = match self {
139 MatrixOp::Matrix(indices, _) => ("Matrix".to_string(), indices.clone()),
140 MatrixOp::SparseMatrix(indices, _) => ("SparseMatrix".to_string(), indices.clone()),
141 MatrixOp::Swap(a_indices, b_indices) => {
142 let indices: Vec<_> = a_indices
143 .iter()
144 .cloned()
145 .chain(b_indices.iter().cloned())
146 .collect();
147 ("Swap".to_string(), indices)
148 }
149 MatrixOp::Control(indices, _, op) => {
150 let name = format!("C({:?})", *op);
151 (name, indices.clone())
152 }
153 };
154 let int_strings = indices
155 .iter()
156 .map(|x| x.clone().to_string())
157 .collect::<Vec<String>>();
158
159 write!(f, "{}[{}]", name, int_strings.join(", "))
160 }
161}