qip/
utils.rs

1#[cfg(feature = "parallel")]
2pub(crate) use rayon::prelude::*;
3
4use qip_iterators::into_iter;
5
6use std::sync::{Arc, Mutex};
7
8/// Mixes two bitstreams, `z_bits` and `o_bits`, takes one bit off the lowest position from
9/// each to construct the output, `selector` is used to choose which to use. 0 indices `z_bits` a
10///
11/// # Example
12/// ```
13/// use qip::utils::entwine_bits;
14///
15/// let n = 3;
16/// let off_bits = 0b01; // 2 bits from off
17/// let on_bits = 0b1;  // 1 bit from on
18/// let selector = 0b010; // first take from off, then on, then off
19/// assert_eq!(entwine_bits(n, selector, off_bits, on_bits), 0b011);
20/// ```
21pub fn entwine_bits(
22    n: usize,
23    mut selector: usize,
24    mut off_bits: usize,
25    mut on_bits: usize,
26) -> usize {
27    let mut result = 0;
28
29    for i in 0..n {
30        if selector & 1 == 0 {
31            let bit = off_bits & 1;
32            off_bits >>= 1;
33            result |= bit << i;
34        } else {
35            let bit = on_bits & 1;
36            on_bits >>= 1;
37            result |= bit << i;
38        }
39        selector >>= 1;
40    }
41
42    result
43}
44
45/// Extracts bits from a number in a particular order.
46///
47/// # Example
48///
49/// ```
50/// use qip::utils::extract_bits;
51///
52/// assert_eq!(extract_bits(0b1010, &[3, 0]), 0b01);
53/// ```
54#[inline]
55pub fn extract_bits(num: usize, indices: &[usize]) -> usize {
56    indices.iter().enumerate().fold(0, |acc, (i, index)| {
57        let bit = (num >> index) & 1;
58        acc | (bit << i)
59    })
60}
61
62/// Transpose a sparse matrix.
63pub fn transpose_sparse<T: Sync + Send>(sparse_mat: Vec<Vec<(usize, T)>>) -> Vec<Vec<(usize, T)>> {
64    let sparse_len = sparse_mat.len();
65    let flat_mat: Vec<_> = into_iter!(sparse_mat)
66        // .into_par_iter()
67        .enumerate()
68        .map(|(row, v)| {
69            let v: Vec<_> = v.into_iter().map(|(col, val)| (col, (row, val))).collect();
70            v
71        })
72        .flatten()
73        .collect();
74    let mut col_mat = <Vec<Arc<Mutex<Vec<(usize, T)>>>>>::new();
75    col_mat.resize_with(sparse_len, || Arc::new(Mutex::new(vec![])));
76    into_iter!(flat_mat).for_each(|(col, (row, val)): (usize, (usize, T))| {
77        col_mat[col].lock().unwrap().push((row, val))
78    });
79    let col_mat: Vec<_> = into_iter!(col_mat)
80        .map(|v| {
81            if let Ok(v) = Arc::try_unwrap(v) {
82                v.into_inner().unwrap()
83            } else {
84                panic!()
85            }
86        })
87        .collect();
88    into_iter!(col_mat)
89        .map(|mut v: Vec<(usize, T)>| {
90            v.sort_by_key(|(row, _)| *row);
91            v
92        })
93        .collect()
94}