twenty_first/math/
zerofier_tree.rs

1use std::collections::VecDeque;
2use std::ops::MulAssign;
3
4use num_traits::One;
5
6use super::b_field_element::BFieldElement;
7use super::polynomial::Polynomial;
8use super::traits::FiniteField;
9
10#[derive(Debug, Clone, PartialEq)]
11pub struct Leaf<'c, FF: FiniteField + MulAssign<BFieldElement>> {
12    pub(crate) points: Vec<FF>,
13    zerofier: Polynomial<'c, FF>,
14}
15
16impl<FF> Leaf<'static, FF>
17where
18    FF: FiniteField + MulAssign<BFieldElement>,
19{
20    pub fn new(points: Vec<FF>) -> Leaf<'static, FF> {
21        let zerofier = Polynomial::zerofier(&points);
22        Self { points, zerofier }
23    }
24}
25
26#[derive(Debug, Clone, PartialEq)]
27pub struct Branch<'c, FF: FiniteField + MulAssign<BFieldElement>> {
28    zerofier: Polynomial<'c, FF>,
29    pub(crate) left: ZerofierTree<'c, FF>,
30    pub(crate) right: ZerofierTree<'c, FF>,
31}
32
33impl<'c, FF> Branch<'c, FF>
34where
35    FF: FiniteField + MulAssign<BFieldElement> + 'static,
36{
37    pub fn new(left: ZerofierTree<'c, FF>, right: ZerofierTree<'c, FF>) -> Self {
38        let zerofier = left.zerofier().multiply(&right.zerofier());
39        Self {
40            zerofier,
41            left,
42            right,
43        }
44    }
45}
46
47/// A zerofier tree is a balanced binary tree of vanishing polynomials.
48/// Conceptually, every leaf corresponds to a single point, and the value of
49/// that leaf is the monic linear polynomial that evaluates to zero there and
50/// no-where else. Every non-leaf node is the product of its two children.
51/// In practice, it makes sense to truncate the tree depth, in which case every
52/// leaf contains a chunk of points whose size is upper-bounded and more or less
53/// equal to some constant threshold.
54#[derive(Debug, Clone, PartialEq)]
55pub enum ZerofierTree<'c, FF: FiniteField + MulAssign<BFieldElement>> {
56    Leaf(Leaf<'c, FF>),
57    Branch(Box<Branch<'c, FF>>),
58    Padding,
59}
60
61impl<FF: FiniteField + MulAssign<BFieldElement>> ZerofierTree<'static, FF> {
62    /// Regulates the depth at which the tree is truncated. Phrased differently,
63    /// regulates the number of points contained by each leaf.
64    const RECURSION_CUTOFF_THRESHOLD: usize = 16;
65
66    pub fn new_from_domain(domain: &[FF]) -> Self {
67        let mut nodes = domain
68            .chunks(Self::RECURSION_CUTOFF_THRESHOLD)
69            .map(|chunk| {
70                let leaf = Leaf::new(chunk.to_vec());
71                ZerofierTree::Leaf(leaf)
72            })
73            .collect::<VecDeque<_>>();
74        nodes.resize(nodes.len().next_power_of_two(), ZerofierTree::Padding);
75        while nodes.len() > 1 {
76            let right = nodes.pop_back().unwrap();
77            let left = nodes.pop_back().unwrap();
78            if left == ZerofierTree::Padding {
79                nodes.push_front(ZerofierTree::Padding);
80            } else {
81                let new_node = Branch::new(left, right);
82                nodes.push_front(ZerofierTree::Branch(Box::new(new_node)));
83            }
84        }
85        nodes.pop_front().unwrap()
86    }
87}
88
89impl<'c, FF> ZerofierTree<'c, FF>
90where
91    FF: FiniteField + MulAssign<BFieldElement> + 'static,
92{
93    pub fn zerofier(&self) -> Polynomial<'c, FF> {
94        match self {
95            ZerofierTree::Leaf(leaf) => leaf.zerofier.clone(),
96            ZerofierTree::Branch(branch) => branch.zerofier.clone(),
97            ZerofierTree::Padding => Polynomial::one(),
98        }
99    }
100}
101
102#[cfg(test)]
103#[cfg_attr(coverage_nightly, coverage(off))]
104mod tests {
105    use num_traits::ConstZero;
106    use num_traits::Zero;
107    use proptest::collection::vec;
108    use proptest::prop_assert_eq;
109    use proptest_arbitrary_adapter::arb;
110
111    use crate::math::zerofier_tree::ZerofierTree;
112    use crate::prelude::BFieldElement;
113    use crate::prelude::Polynomial;
114    use crate::tests::proptest;
115    use crate::tests::test;
116
117    #[macro_rules_attr::apply(test)]
118    fn zerofier_tree_can_be_empty() {
119        ZerofierTree::<BFieldElement>::new_from_domain(&[]);
120    }
121
122    #[macro_rules_attr::apply(proptest)]
123    fn zerofier_tree_root_is_multiple_of_children(
124        #[strategy(vec(arb(), 2*ZerofierTree::<BFieldElement>::RECURSION_CUTOFF_THRESHOLD))]
125        points: Vec<BFieldElement>,
126    ) {
127        let zerofier_tree = ZerofierTree::new_from_domain(&points);
128        let ZerofierTree::Branch(ref branch) = zerofier_tree else {
129            panic!("not enough leafs");
130        };
131        prop_assert_eq!(
132            Polynomial::zero(),
133            zerofier_tree.zerofier().reduce(&branch.left.zerofier())
134        );
135        prop_assert_eq!(
136            Polynomial::zero(),
137            zerofier_tree.zerofier().reduce(&branch.right.zerofier())
138        );
139    }
140
141    #[macro_rules_attr::apply(proptest)]
142    fn zerofier_tree_root_has_right_degree(
143        #[strategy(vec(arb(), 1..(1<<10)))] points: Vec<BFieldElement>,
144    ) {
145        let zerofier_tree = ZerofierTree::new_from_domain(&points);
146        prop_assert_eq!(points.len(), zerofier_tree.zerofier().degree() as usize);
147    }
148
149    #[macro_rules_attr::apply(proptest)]
150    fn zerofier_tree_root_zerofies(
151        #[strategy(vec(arb(), 1..(1<<10)))] points: Vec<BFieldElement>,
152        #[strategy(0usize..#points.len())] index: usize,
153    ) {
154        let zerofier_tree = ZerofierTree::new_from_domain(&points);
155        prop_assert_eq!(
156            BFieldElement::ZERO,
157            zerofier_tree.zerofier().evaluate(points[index])
158        );
159    }
160
161    #[macro_rules_attr::apply(proptest)]
162    fn zerofier_tree_and_polynomial_agree_on_zerofiers(
163        #[strategy(vec(arb(), 1..(1<<10)))] points: Vec<BFieldElement>,
164    ) {
165        let zerofier_tree = ZerofierTree::new_from_domain(&points);
166        let polynomial_zerofier = Polynomial::zerofier(&points);
167        prop_assert_eq!(polynomial_zerofier, zerofier_tree.zerofier());
168    }
169}
twenty_first/math/zerofier_tree.rs

twenty_first/math/
zerofier_tree.rs
