Skip to main content

tasm_lib/array/
horner_evaluation.rs

1use std::collections::HashMap;
2
3use itertools::Itertools;
4use triton_vm::prelude::*;
5
6use crate::data_type::ArrayType;
7use crate::prelude::*;
8use crate::traits::basic_snippet::Reviewer;
9use crate::traits::basic_snippet::SignOffFingerprint;
10
11/// Evaluate a polynomial in a point using the Horner method.
12///
13/// `HornerEvaluation` takes an array of coefficients (representing a
14/// polynomial) and a scalar (representing an indeterminate) and computes the
15/// value of the polynomial in that point. It can be used for univariate
16/// batching, whereby the object is to compute a random linear sum of a given
17/// set of points, and the weights are given by the powers of one challenge.
18///
19/// ### Behavior
20///
21/// ```text
22/// BEFORE: _ *coefficients [indeterminate: XFieldElement]
23/// AFTER:  _ [evaluation: XFieldElement]
24/// ```
25///
26/// ### Preconditions
27///
28/// None.
29///
30/// ### Postconditions
31///
32/// None.
33#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
34pub struct HornerEvaluation {
35    pub num_coefficients: usize,
36}
37
38impl HornerEvaluation {
39    pub fn new(num_coefficients: usize) -> Self {
40        Self { num_coefficients }
41    }
42}
43
44impl BasicSnippet for HornerEvaluation {
45    fn parameters(&self) -> Vec<(DataType, String)> {
46        let coefficients_ty = DataType::Array(Box::new(ArrayType {
47            element_type: DataType::Xfe,
48            length: self.num_coefficients,
49        }));
50
51        vec![
52            (coefficients_ty, "*coefficients".to_string()),
53            (DataType::Xfe, "indeterminate".to_string()),
54        ]
55    }
56
57    fn return_values(&self) -> Vec<(DataType, String)> {
58        vec![(DataType::Xfe, "evaluation".to_string())]
59    }
60
61    fn entrypoint(&self) -> String {
62        let n = self.num_coefficients;
63        format!("tasmlib_array_horner_evaluation_with_{n}_coefficients")
64    }
65
66    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
67        let update_running_evaluation = triton_asm! {
68            // BEFORE: _ *coefficients_end   [x: XFE] [v: XFE]
69            // AFTER : _ *coefficients_end-3 [x: XFE] [vx+c: XFE]
70            dup 5 dup 5 dup 5           // _ *coefficients_end [x] [v] [x]
71            xx_mul                      // _ *coefficients_end [x] [vx]
72            pick 6                      // _ [x] [vx] *coefficients_end
73            read_mem 3                  // _ [x] [vx] [c] *coefficients_end-3
74            place 9                     // _ *coefficients_end-3 [x] [vx] [c]
75            xx_add                      // _ *coefficients_end-3 [x] [vx+c]
76        };
77        let update_running_evaluation_for_each_coefficient = (0..self.num_coefficients)
78            .flat_map(|_| update_running_evaluation.clone())
79            .collect_vec();
80
81        triton_asm! {
82            // BEFORE: _ *coefficients [x: XFE]
83            // AFTER:  _ [v: XFE]
84            {self.entrypoint()}:
85                pick 3                  // _ [x: XFE] *coefficients
86                addi {(self.num_coefficients * 3).saturating_sub(1)}
87                place 3                 // _ *coefficients_end [x: XFE]
88                push 0 push 0 push 0    // _ *coefficients_end [x: XFE] [v: XFE]
89                {&update_running_evaluation_for_each_coefficient}
90                                        // _ *coefficients_end-3n [x: XFE] [v': XFE]
91                place 6 place 6 place 6 // _ [v': XFE] *coefficients_end-3n [x: XFE]
92                pop 4                   // _ [v': XFE]
93                return
94        }
95    }
96
97    fn sign_offs(&self) -> HashMap<Reviewer, SignOffFingerprint> {
98        let mut sign_offs = HashMap::new();
99        if self.num_coefficients == 4 {
100            sign_offs.insert(Reviewer("ferdinand"), 0xd6cdaea8d7f5385a.into());
101        }
102
103        sign_offs
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use crate::rust_shadowing_helper_functions::array::array_get;
111    use crate::rust_shadowing_helper_functions::array::insert_as_array;
112    use crate::test_prelude::*;
113    use crate::twenty_first::prelude::Polynomial;
114
115    impl Accessor for HornerEvaluation {
116        fn rust_shadow(
117            &self,
118            stack: &mut Vec<BFieldElement>,
119            memory: &HashMap<BFieldElement, BFieldElement>,
120        ) -> Result<(), RustShadowError> {
121            let indeterminate = pop_encodable(stack)?;
122            let coefficient_ptr = stack.pop().ok_or(RustShadowError::StackUnderflow)?;
123
124            let coefficients: Vec<XFieldElement> = (0..self.num_coefficients)
125                .map(|i| array_get(coefficient_ptr, i, memory, 3))
126                .map(|bfes| {
127                    bfes.try_into()
128                        .map(XFieldElement::new)
129                        .map_err(|_| RustShadowError::Other)
130                })
131                .try_collect()?;
132            let polynomial = Polynomial::new(coefficients);
133            let evaluation = polynomial.evaluate_in_same_field(indeterminate);
134
135            push_encodable(stack, &evaluation);
136            Ok(())
137        }
138
139        fn pseudorandom_initial_state(
140            &self,
141            seed: [u8; 32],
142            _: Option<BenchmarkCase>,
143        ) -> AccessorInitialState {
144            let mut rng = StdRng::from_seed(seed);
145
146            let coefficients = (0..self.num_coefficients)
147                .map(|_| rng.random::<XFieldElement>())
148                .collect();
149            let address = rng.random();
150
151            let mut memory = HashMap::new();
152            insert_as_array(address, &mut memory, coefficients);
153
154            let mut stack = self.init_stack_for_isolated_run();
155            stack.push(address);
156            push_encodable(&mut stack, &rng.random::<XFieldElement>()); // indeterminate
157
158            AccessorInitialState { stack, memory }
159        }
160    }
161
162    #[macro_rules_attr::apply(test)]
163    fn rust_shadow() {
164        for n in [0, 1, 4, 20, 587, 1000] {
165            ShadowedAccessor::new(HornerEvaluation::new(n)).test();
166        }
167    }
168}
169
170#[cfg(test)]
171mod benches {
172    use super::*;
173    use crate::test_prelude::*;
174
175    #[macro_rules_attr::apply(test)]
176    fn bench() {
177        for n in [100, 587] {
178            ShadowedAccessor::new(HornerEvaluation::new(n)).bench();
179        }
180    }
181}