1use num_traits::{One, Zero};
2use std::marker::PhantomData;
3
4use crate::utils::*;
5
6#[derive(Debug)]
8pub struct MatrixOpIterator<'a, P>
9where
10 P: Clone + Zero,
11{
12 n: usize,
13 data: &'a [P],
14 last_col: Option<usize>,
15}
16
17impl<'a, P> MatrixOpIterator<'a, P>
18where
19 P: Clone + Zero,
20{
21 pub fn new(row: usize, n: usize, data: &'a [P]) -> MatrixOpIterator<P> {
24 let lower = get_flat_index(n, row, 0);
25 let upper = get_flat_index(n, row, 1 << n);
26 MatrixOpIterator {
27 n,
28 data: &data[lower..upper],
29 last_col: None,
30 }
31 }
32}
33
34impl<'a, P> Iterator for MatrixOpIterator<'a, P>
35where
36 P: Clone + Zero,
37{
38 type Item = (usize, P);
39
40 fn next(&mut self) -> Option<Self::Item> {
41 let pos = if let Some(last_col) = self.last_col {
42 last_col + 1
43 } else {
44 0
45 };
46 self.last_col = None;
47 for col in pos..(1 << self.n) {
48 let val = &self.data[col];
49 if !val.is_zero() {
50 self.last_col = Some(col);
51 return Some((col, val.clone()));
52 }
53 }
54 None
55 }
56}
57
58#[derive(Debug)]
60pub struct SparseMatrixOpIterator<'a, P>
61where
62 P: Clone,
63{
64 data: &'a [(usize, P)],
65 last_index: Option<usize>,
66}
67
68impl<'a, P> SparseMatrixOpIterator<'a, P>
69where
70 P: Clone,
71{
72 pub fn new(row: usize, data: &'a [Vec<(usize, P)>]) -> SparseMatrixOpIterator<P> {
74 SparseMatrixOpIterator {
75 data: &data[row],
76 last_index: None,
77 }
78 }
79}
80
81impl<'a, P> Iterator for SparseMatrixOpIterator<'a, P>
82where
83 P: Clone,
84{
85 type Item = (usize, P);
86
87 fn next(&mut self) -> Option<Self::Item> {
88 let pos = if let Some(last_index) = self.last_index {
89 last_index + 1
90 } else {
91 0
92 };
93
94 if pos >= self.data.len() {
95 self.last_index = None;
96 None
97 } else {
98 self.last_index = Some(pos);
99 Some(self.data[pos].clone())
100 }
101 }
102}
103
104#[derive(Debug)]
106pub struct ControlledOpIterator<P, It>
107where
108 P: Clone + Zero,
109 It: Iterator<Item = (usize, P)>,
110{
111 row: usize,
112 index_threshold: usize,
113 op_iter: Option<It>,
114 last_col: Option<usize>,
115}
116
117impl<P, It> ControlledOpIterator<P, It>
118where
119 P: Clone + Zero,
120 It: Iterator<Item = (usize, P)>,
121{
122 pub fn new<F: FnOnce(usize) -> It>(
125 row: usize,
126 n_control_indices: usize,
127 n_op_indices: usize,
128 iter_builder: F,
129 ) -> ControlledOpIterator<P, It> {
130 let n_indices = n_control_indices + n_op_indices;
131 let index_threshold = (1 << n_indices) - (1 << n_op_indices);
132 let op_iter = if row >= index_threshold {
133 Some(iter_builder(row - index_threshold))
134 } else {
135 None
136 };
137 ControlledOpIterator {
138 row,
139 index_threshold,
140 op_iter,
141 last_col: None,
142 }
143 }
144}
145
146impl<P, It> Iterator for ControlledOpIterator<P, It>
147where
148 P: Clone + Zero + One,
149 It: Iterator<Item = (usize, P)>,
150{
151 type Item = (usize, P);
152
153 fn next(&mut self) -> Option<Self::Item> {
154 if let Some(it) = &mut self.op_iter {
155 let cval = it.next();
156 let ret_val = cval.map(|(i, val)| (i + self.index_threshold, val));
157 self.last_col = ret_val.as_ref().map(|(i, _)| *i);
158 ret_val
159 } else {
160 if let Some(last_col) = self.last_col {
161 if last_col < self.row {
162 self.last_col = Some(self.row);
163 } else {
164 self.last_col = None;
165 }
166 } else {
167 self.last_col = Some(self.row);
168 }
169 self.last_col.map(|c| (c, P::one()))
170 }
171 }
172}
173
174#[derive(Debug)]
176pub struct SwapOpIterator<P>
177where
178 P: One,
179{
180 row: usize,
181 half_n: usize,
182 last_col: Option<usize>,
183 phantom: PhantomData<P>,
184}
185
186impl<P> SwapOpIterator<P>
187where
188 P: One,
189{
190 pub fn new(row: usize, n_qubits: usize) -> SwapOpIterator<P> {
193 SwapOpIterator {
194 row,
195 half_n: n_qubits >> 1,
196 last_col: None,
197 phantom: PhantomData,
198 }
199 }
200}
201
202impl<P> Iterator for SwapOpIterator<P>
203where
204 P: One,
205{
206 type Item = (usize, P);
207
208 fn next(&mut self) -> Option<Self::Item> {
209 if self.last_col.is_none() {
210 let lower_mask: usize = !(std::usize::MAX << self.half_n);
211 let lower = self.row & lower_mask;
212 let upper = self.row >> self.half_n;
213 self.last_col = Some((lower << self.half_n) + upper);
214 } else {
215 self.last_col = None;
216 }
217 self.last_col.map(|col| (col, P::one()))
218 }
219}
220
221#[derive(Debug)]
223pub struct FunctionOpIterator<P> {
224 output_n: usize,
225 x: usize,
226 fx_xor_y: usize,
227 theta: P,
228 last_col: Option<usize>,
229 phantom: PhantomData<P>,
230}
231
232impl<P> FunctionOpIterator<P> {
233 pub fn new<F: Fn(usize) -> (usize, P)>(
236 row: usize,
237 input_n: usize,
238 output_n: usize,
239 f: F,
240 ) -> FunctionOpIterator<P> {
241 let x = row >> output_n;
242 let (fx, theta) = f(flip_bits(input_n, x));
243 let y = row & ((1 << output_n) - 1);
244 let fx_xor_y = y ^ flip_bits(output_n, fx);
245 FunctionOpIterator {
246 output_n,
247 x,
248 fx_xor_y,
249 theta,
250 last_col: None,
251 phantom: PhantomData,
252 }
253 }
254}
255
256impl<P> Iterator for FunctionOpIterator<P>
257where
258 P: Clone,
259{
260 type Item = (usize, P);
261
262 fn next(&mut self) -> Option<Self::Item> {
263 if self.last_col.is_some() {
264 self.last_col = None;
265 } else {
266 let colbits = (self.x << self.output_n) | self.fx_xor_y;
267 self.last_col = Some(colbits);
268 };
269 self.last_col.map(|col| (col, self.theta.clone()))
270 }
271}
272
273#[cfg(test)]
274mod iterator_tests {
275 use super::*;
276 use num_complex::Complex;
277 use num_traits::One;
278
279 fn from_reals<P: Zero + Clone>(data: &[P]) -> Vec<Complex<P>> {
281 data.iter()
282 .map(|x| Complex::<P> {
283 re: x.clone(),
284 im: P::zero(),
285 })
286 .collect()
287 }
288
289 #[test]
290 fn test_mat_iterator() {
291 let n = 1usize;
292 let mat: Vec<Vec<f64>> = (0..1 << n)
293 .map(|i| -> Vec<f64> {
294 let d = from_reals(&[0.0, 1.0, 1.0, 0.0]);
295 let it = MatrixOpIterator::new(i, n, &d);
296 let v: Vec<f64> = (0..1 << n).map(|_| 0.0).collect();
297 it.fold(v, |mut v, (indx, _)| {
298 v[indx] = 1.0;
299 v
300 })
301 })
302 .collect();
303
304 let expected = vec![vec![0.0, 1.0], vec![1.0, 0.0]];
305
306 assert_eq!(mat, expected);
307 }
308
309 #[test]
310 fn test_sparse_mat_iterator() {
311 let n = 1usize;
312 let one = Complex::<f64>::one();
313 let mat: Vec<Vec<f64>> = (0..1 << n)
314 .map(|i| -> Vec<f64> {
315 let d = vec![vec![(1, one)], vec![(0, one)]];
316 let it = SparseMatrixOpIterator::new(i, &d);
317 let v: Vec<f64> = (0..1 << n).map(|_| 0.0).collect();
318 it.fold(v, |mut v, (indx, _)| {
319 v[indx] = 1.0;
320 v
321 })
322 })
323 .collect();
324
325 let expected = vec![vec![0.0, 1.0], vec![1.0, 0.0]];
326
327 assert_eq!(mat, expected);
328 }
329
330 #[test]
331 fn test_swap_iterator() {
332 let n = 2usize;
333 let mat: Vec<Vec<f64>> = (0..1 << n)
334 .map(|i| -> Vec<f64> {
335 let it = SwapOpIterator::<f64>::new(i, n);
336 let v: Vec<f64> = (0..1 << n).map(|_| 0.0).collect();
337 it.fold(v, |mut v, (indx, _)| {
338 v[indx] = 1.0;
339 v
340 })
341 })
342 .collect();
343
344 let expected = vec![
345 vec![1.0, 0.0, 0.0, 0.0],
346 vec![0.0, 0.0, 1.0, 0.0],
347 vec![0.0, 1.0, 0.0, 0.0],
348 vec![0.0, 0.0, 0.0, 1.0],
349 ];
350
351 assert_eq!(mat, expected);
352 }
353
354 #[test]
355 fn test_c_iterator() {
356 let n = 2usize;
357
358 let d = from_reals(&[0.0, 1.0, 1.0, 0.0]);
359 let builder = |r: usize| MatrixOpIterator::new(r, n >> 1, &d);
360
361 let mat: Vec<Vec<f64>> = (0..1 << n)
362 .map(|i| -> Vec<f64> {
363 let it = ControlledOpIterator::new(i, n >> 1, n >> 1, builder);
364 let v: Vec<f64> = (0..1 << n).map(|_| 0.0).collect();
365 it.fold(v, |mut v, (indx, _)| {
366 v[indx] = 1.0;
367 v
368 })
369 })
370 .collect();
371
372 let expected = vec![
373 vec![1.0, 0.0, 0.0, 0.0],
374 vec![0.0, 1.0, 0.0, 0.0],
375 vec![0.0, 0.0, 0.0, 1.0],
376 vec![0.0, 0.0, 1.0, 0.0],
377 ];
378 assert_eq!(mat, expected);
379 }
380}