stwo_gpu/core/
queries.rs

1use core::ops::Deref;
2
3use itertools::Itertools;
4use std_shims::{BTreeSet, Vec};
5
6use super::channel::Channel;
7
8pub const UPPER_BOUND_QUERY_BYTES: usize = 4;
9
10/// Draws `n_queries` values in the range `[0, 2^log_domain_size)` from the channel.
11pub fn draw_queries(
12    channel: &mut impl Channel,
13    log_domain_size: u32,
14    n_queries: usize,
15) -> Vec<usize> {
16    let mut raw_positions = Vec::new();
17    let query_mask = (1 << log_domain_size) - 1;
18    loop {
19        let random_words = channel.draw_u32s();
20        for word in random_words {
21            let quotient_query = word & query_mask;
22            raw_positions.push(quotient_query.try_into().unwrap());
23            if raw_positions.len() == n_queries {
24                return raw_positions;
25            }
26        }
27    }
28}
29
30/// An ordered set of query positions.
31#[derive(Debug, Clone)]
32pub struct Queries {
33    /// Query positions sorted in ascending order.
34    pub positions: Vec<usize>,
35    /// Size of the domain from which the queries were sampled.
36    pub log_domain_size: u32,
37}
38
39impl Queries {
40    /// Creates a [Queries] instance from the given unsorted `raw_positions`.
41    pub fn new(raw_positions: &[usize], log_domain_size: u32) -> Self {
42        Self {
43            positions: BTreeSet::from_iter(raw_positions.iter())
44                .into_iter()
45                .cloned()
46                .collect(),
47            log_domain_size,
48        }
49    }
50
51    /// Calculates the matching query indices in a folded domain (i.e each domain point is doubled)
52    /// given `self` (the queries of the original domain) and the number of folds between domains.
53    pub fn fold(&self, n_folds: u32) -> Self {
54        assert!(n_folds <= self.log_domain_size);
55        Self {
56            positions: self.iter().map(|q| q >> n_folds).dedup().collect(),
57            log_domain_size: self.log_domain_size - n_folds,
58        }
59    }
60
61    #[cfg(test)]
62    pub fn from_positions(positions: Vec<usize>, log_domain_size: u32) -> Self {
63        assert!(positions.is_sorted());
64        assert!(positions.iter().all(|p| *p < (1 << log_domain_size)));
65        Self {
66            positions,
67            log_domain_size,
68        }
69    }
70}
71
72impl Deref for Queries {
73    type Target = Vec<usize>;
74
75    fn deref(&self) -> &Self::Target {
76        &self.positions
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use std_shims::Vec;
83
84    use crate::core::channel::Blake2sChannel;
85    use crate::core::poly::circle::CanonicCoset;
86    use crate::core::queries::{draw_queries, Queries};
87    use crate::core::utils::bit_reverse;
88
89    #[test]
90    fn test_generate_queries() {
91        let channel = &mut Blake2sChannel::default();
92        let log_query_size = 31;
93        let n_queries = 100;
94
95        let raw_positions = draw_queries(channel, log_query_size, n_queries);
96        let queries = Queries::new(&raw_positions, log_query_size);
97
98        assert!(queries.len() == n_queries);
99        assert!(queries.iter().is_sorted());
100        assert!(*queries.positions.last().unwrap() < 1 << log_query_size);
101    }
102
103    #[test]
104    pub fn test_folded_queries() {
105        let log_domain_size = 7;
106        let domain = CanonicCoset::new(log_domain_size).circle_domain();
107        let mut values = domain.iter().collect::<Vec<_>>();
108        bit_reverse(&mut values);
109
110        let log_folded_domain_size = 5;
111        let folded_domain = CanonicCoset::new(log_folded_domain_size).circle_domain();
112        let mut folded_values = folded_domain.iter().collect::<Vec<_>>();
113        bit_reverse(&mut folded_values);
114
115        // Generate all possible queries.
116        let queries = Queries {
117            positions: (0..1 << log_domain_size).collect(),
118            log_domain_size,
119        };
120        let n_folds = log_domain_size - log_folded_domain_size;
121        let ratio = 1 << n_folds;
122
123        let folded_queries = queries.fold(n_folds);
124        let repeated_folded_queries = folded_queries
125            .iter()
126            .flat_map(|q| core::iter::repeat_n(q, ratio));
127        for (query, folded_query) in queries.iter().zip(repeated_folded_queries) {
128            // Check only the x coordinate since folding might give you the conjugate point.
129            assert_eq!(
130                values[*query].repeated_double(n_folds).x,
131                folded_values[*folded_query].x
132            );
133        }
134    }
135}
stwo_gpu/core/queries.rs

stwo_gpu/core/
queries.rs
