ark_poly_commit/streaming_kzg/
data_structures.rs

1use crate::utils::ceil_div;
2use ark_ff::Field;
3#[cfg(not(feature = "std"))]
4use ark_std::vec::Vec;
5use ark_std::{borrow::Borrow, iterable::Iterable};
6
7/// A `Streamer` folding a vector of coefficients
8/// with the given challenges, and producing a stream of items
9/// `(i, v)` where `i` indicates the depth, and `v` is the next coefficient.
10/// The stream can produce all foldings in the tree with a single pass over the initial stream.
11#[derive(Clone, Copy)]
12pub struct FoldedPolynomialTree<'a, F, S> {
13    challenges: &'a [F],
14    coefficients: &'a S,
15}
16
17impl<'a, F, S> FoldedPolynomialTree<'a, F, S>
18where
19    S: Iterable,
20    F: Field,
21    S::Item: Borrow<F>,
22{
23    /// Initialize a new polynomial tree.
24    pub fn new(coefficients: &'a S, challenges: &'a [F]) -> Self {
25        Self {
26            coefficients,
27            challenges,
28        }
29    }
30
31    /// Outputs the depth of the polynomial tree.
32    #[inline]
33    pub fn depth(&self) -> usize {
34        self.challenges.len()
35    }
36}
37
38impl<'a, F, S> Iterable for FoldedPolynomialTree<'a, F, S>
39where
40    S: Iterable,
41    F: Field,
42    S::Item: Borrow<F>,
43{
44    type Item = (usize, F);
45
46    type Iter = FoldedPolynomialTreeIter<'a, F, S::Iter>;
47
48    fn iter(&self) -> Self::Iter {
49        FoldedPolynomialTreeIter::new(
50            self.coefficients.iter(),
51            self.coefficients.len(),
52            self.challenges,
53        )
54    }
55
56    fn len(&self) -> usize {
57        self.coefficients.len()
58    }
59}
60
61/// Iterator of the polynomial tree.
62pub struct FoldedPolynomialTreeIter<'a, F, I> {
63    challenges: &'a [F],
64    iterator: I,
65    stack: Vec<(usize, F)>,
66}
67
68fn init_stack<F: Field>(n: usize, challenges_len: usize) -> Vec<(usize, F)> {
69    let mut stack = Vec::with_capacity(challenges_len);
70
71    // generally we expect the size to be a power of two.
72    // If not, we are going to fill the stack as if the array was padded to zero up to the expected size.
73    let chunk_size = 1 << challenges_len;
74    if n % chunk_size != 0 {
75        let mut delta = chunk_size - n % chunk_size;
76        for i in (0..challenges_len).rev() {
77            if delta >= 1 << i {
78                stack.push((i, F::zero()));
79                delta -= 1 << i
80            }
81        }
82    }
83    stack
84}
85
86impl<'a, F, I> FoldedPolynomialTreeIter<'a, F, I>
87where
88    F: Field,
89    I: Iterator,
90    I::Item: Borrow<F>,
91{
92    fn new(iterator: I, n: usize, challenges: &'a [F]) -> Self {
93        let stack = init_stack(n, challenges.len());
94
95        Self {
96            challenges,
97            iterator,
98            stack,
99        }
100    }
101}
102
103impl<'a, F, I> Iterator for FoldedPolynomialTreeIter<'a, F, I>
104where
105    F: Field,
106    I: Iterator,
107    I::Item: Borrow<F>,
108{
109    type Item = (usize, F);
110
111    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
112        let len = self.stack.len();
113        let item = if len > 1 && self.stack[len - 1].0 == self.stack[len - 2].0 {
114            // pop the last two elements from the stack.
115            // we could also use .pop() twice but truncate is slightly faster.
116            let (_level, lhs) = self.stack[len - 1];
117            let (level, rhs) = self.stack[len - 2];
118            self.stack.truncate(len - 2);
119            // fold them producing the coefficient and the level `level+1`
120            let folded_coefficient = rhs * self.challenges[level] + lhs;
121            (level + 1, folded_coefficient)
122        } else {
123            (0, *self.iterator.next()?.borrow())
124        };
125
126        // do not add to the stack the coefficient of the max-depth folded polynomial.
127        if item.0 != self.challenges.len() {
128            self.stack.push(item)
129        }
130
131        // Skip the base polynomial, recursively calling itself to access the next level
132        if item.0 == 0 {
133            self.next()
134        } else {
135            Some(item)
136        }
137    }
138}
139
140/// Stream implementation of foleded polynomial.
141#[derive(Clone, Copy)]
142pub struct FoldedPolynomialStream<'a, F, S>(FoldedPolynomialTree<'a, F, S>);
143/// Iterator implementation of foleded polynomial.
144pub struct FoldedPolynomialStreamIter<'a, F, I> {
145    challenges: &'a [F],
146    iterator: I,
147    stack: Vec<(usize, F)>,
148}
149
150impl<'a, F, S> FoldedPolynomialStream<'a, F, S>
151where
152    S: Iterable,
153    F: Field,
154    S::Item: Borrow<F>,
155{
156    /// Initialize a new folded polynomial stream.
157    pub fn new(coefficients: &'a S, challenges: &'a [F]) -> Self {
158        let tree = FoldedPolynomialTree::new(coefficients, challenges);
159        Self(tree)
160    }
161}
162
163impl<'a, F, S> Iterable for FoldedPolynomialStream<'a, F, S>
164where
165    S: Iterable,
166    F: Field,
167    S::Item: Borrow<F>,
168{
169    type Item = F;
170    type Iter = FoldedPolynomialStreamIter<'a, F, S::Iter>;
171
172    fn iter(&self) -> Self::Iter {
173        let iterator = self.0.coefficients.iter();
174        let challenges = self.0.challenges;
175        let stack = init_stack(self.0.coefficients.len(), challenges.len());
176        FoldedPolynomialStreamIter {
177            iterator,
178            challenges,
179            stack,
180        }
181    }
182
183    fn len(&self) -> usize {
184        ceil_div(self.0.len(), 1 << self.0.challenges.len())
185    }
186}
187
188impl<'a, F, I> Iterator for FoldedPolynomialStreamIter<'a, F, I>
189where
190    F: Field,
191    I: Iterator,
192    I::Item: Borrow<F>,
193{
194    type Item = F;
195
196    fn next(&mut self) -> Option<Self::Item> {
197        let target_level = self.challenges.len();
198        loop {
199            let len = self.stack.len();
200            let (level, element) = if len > 1 && self.stack[len - 1].0 == self.stack[len - 2].0 {
201                let (_level, lhs) = self.stack[len - 1];
202                let (level, rhs) = self.stack[len - 2];
203                self.stack.truncate(len - 2);
204
205                let folded_coefficient = rhs * self.challenges[level] + lhs;
206                (level + 1, folded_coefficient)
207            } else if target_level > 0 && (len == 0 || (len > 0 && self.stack[len - 1].0 != 0)) {
208                // If the target level is strictly positive, there's no need to put elements of level zero in the stream.
209                // We can immediately read 2 elements from the stream and push an element of the form (1, folded_coefficient).
210                // Nota bene: this branch is not needed, but brings in a decent speed-up for the resulting implementation.
211                let rhs = self.iterator.next()?;
212                let lhs = self.iterator.next()?;
213
214                let folded_coefficient = self.challenges[0] * rhs.borrow() + lhs.borrow();
215                (1, folded_coefficient)
216            } else {
217                (0, *self.iterator.next()?.borrow())
218            };
219
220            // do not add to the stack the coefficient of the folded polynomial, but instead return it.
221            if level != target_level {
222                self.stack.push((level, element))
223            } else {
224                return Some(element);
225            }
226        }
227    }
228}
229
230#[test]
231fn test_folded_polynomial() {
232    use ark_bls12_381::Fr as F;
233    use ark_ff::One;
234
235    let two = F::one() + F::one();
236
237    let coefficients = vec![F::one(), two, F::one(), F::one()];
238    let challenges = vec![F::one(), two];
239    let coefficients_stream = coefficients.as_slice();
240    let foldstream = FoldedPolynomialTree::new(&coefficients_stream, challenges.as_slice());
241    let fold_stream = FoldedPolynomialStream(foldstream);
242    assert_eq!(fold_stream.len(), 1);
243    assert_eq!(
244        fold_stream.iter().next(),
245        Some(two + two * (F::one() + two))
246    );
247
248    let one = F::one();
249    let coefficients = vec![one; 12];
250    let challenges = vec![F::one(); 4];
251    let coefficients_stream = coefficients.as_slice();
252    let foldstream = FoldedPolynomialTree::new(&coefficients_stream, challenges.as_slice());
253    let fold_stream = FoldedPolynomialStream(foldstream).iter();
254    assert_eq!(fold_stream.last(), Some(coefficients.iter().sum()));
255}
256
257#[test]
258fn test_folded_polynomial_tree() {
259    use ark_bls12_381::Fr as F;
260    use ark_ff::One;
261
262    let two = F::one() + F::one();
263
264    let coefficients = vec![F::one(), two, F::one(), F::one()];
265    let challenges = vec![F::one(), two];
266    let coefficients_stream = coefficients.as_slice();
267    let fold_streamer = FoldedPolynomialTree::new(&coefficients_stream, challenges.as_slice());
268    let mut fold_iter = fold_streamer.iter();
269    // assert_eq!(fold_stream.next(), Some((0, F::one())));
270    // assert_eq!(fold_stream.next(), Some((0, two)));
271    assert_eq!(fold_iter.next(), Some((1, F::one() + two)));
272    // assert_eq!(fold_stream.next(), Some((0, F::one())));
273    // assert_eq!(fold_stream.next(), Some((0, F::one())));
274    assert_eq!(fold_iter.next(), Some((1, F::one() + F::one())));
275    assert_eq!(fold_iter.next(), Some((2, two + two * (F::one() + two))));
276
277    let one = F::one();
278    let coefficients = vec![one; 12];
279    let challenges = vec![F::one(); 4];
280    let coefficients_stream = coefficients.as_slice();
281    let fold_streamer = FoldedPolynomialTree::new(&coefficients_stream, challenges.as_slice());
282    let fold_init = fold_streamer.iter();
283    let mut fold_iter = fold_init.skip(5);
284    assert_eq!(fold_iter.next(), Some((1, two)));
285    assert_eq!(fold_iter.last(), Some((4, coefficients.iter().sum())));
286}
ark_poly_commit/streaming_kzg/data_structures.rs

ark_poly_commit/streaming_kzg/
data_structures.rs
