sp1_recursion_compiler/ir/
utils.rs

1use p3_field::{AbstractExtensionField, AbstractField};
2use std::ops::{Add, Mul, MulAssign};
3
4use super::{Array, Builder, Config, DslIr, Ext, Felt, SymbolicExt, Usize, Var, Variable};
5
6impl<C: Config> Builder<C> {
7    /// The generator for the field.
8    ///
9    /// Reference: [p3_baby_bear::BabyBear]
10    pub fn generator(&mut self) -> Felt<C::F> {
11        self.eval(C::F::from_canonical_u32(31))
12    }
13
14    /// Select a variable based on a condition.
15    pub fn select_v(&mut self, cond: Var<C::N>, a: Var<C::N>, b: Var<C::N>) -> Var<C::N> {
16        let c = self.uninit();
17        self.push_op(DslIr::CircuitSelectV(cond, a, b, c));
18        c
19    }
20
21    /// Select a felt based on a condition.
22    pub fn select_f(&mut self, cond: Var<C::N>, a: Felt<C::F>, b: Felt<C::F>) -> Felt<C::F> {
23        let c = self.uninit();
24        self.push_op(DslIr::CircuitSelectF(cond, a, b, c));
25        c
26    }
27
28    /// Select an extension based on a condition.
29    pub fn select_ef(
30        &mut self,
31        cond: Var<C::N>,
32        a: Ext<C::F, C::EF>,
33        b: Ext<C::F, C::EF>,
34    ) -> Ext<C::F, C::EF> {
35        let c = self.uninit();
36        self.push_op(DslIr::CircuitSelectE(cond, a, b, c));
37        c
38    }
39
40    /// Exponentiates a variable to a power of two.
41    pub fn exp_power_of_2<V: Variable<C>, E: Into<V::Expression>>(
42        &mut self,
43        e: E,
44        power_log: usize,
45    ) -> V
46    where
47        V::Expression: MulAssign<V::Expression> + Clone,
48    {
49        let mut e = e.into();
50        for _ in 0..power_log {
51            e *= e.clone();
52        }
53        self.eval(e)
54    }
55
56    /// Exponentializes a variable to an array of bits in little endian.
57    pub fn exp_bits<V>(&mut self, x: V, power_bits: &Array<C, Var<C::N>>) -> V
58    where
59        V::Expression: AbstractField,
60        V: Copy + Mul<Output = V::Expression> + Variable<C>,
61    {
62        let result = self.eval(V::Expression::one());
63        let power_f: V = self.eval(x);
64        self.range(0, power_bits.len()).for_each(|i, builder| {
65            let bit = builder.get(power_bits, i);
66            builder
67                .if_eq(bit, C::N::one())
68                .then(|builder| builder.assign(result, result * power_f));
69            builder.assign(power_f, power_f * power_f);
70        });
71        result
72    }
73
74    /// Exponentiates a felt to a list of bits in little endian.
75    pub fn exp_f_bits(&mut self, x: Felt<C::F>, power_bits: Vec<Var<C::N>>) -> Felt<C::F> {
76        let mut result = self.eval(C::F::one());
77        let mut power_f: Felt<_> = self.eval(x);
78        for i in 0..power_bits.len() {
79            let bit = power_bits[i];
80            let tmp = self.eval(result * power_f);
81            result = self.select_f(bit, tmp, result);
82            power_f = self.eval(power_f * power_f);
83        }
84        result
85    }
86
87    /// Exponentiates a extension to a list of bits in little endian.
88    pub fn exp_e_bits(
89        &mut self,
90        x: Ext<C::F, C::EF>,
91        power_bits: Vec<Var<C::N>>,
92    ) -> Ext<C::F, C::EF> {
93        let mut result = self.eval(SymbolicExt::from_f(C::EF::one()));
94        let mut power_f: Ext<_, _> = self.eval(x);
95        for i in 0..power_bits.len() {
96            let bit = power_bits[i];
97            let tmp = self.eval(result * power_f);
98            result = self.select_ef(bit, tmp, result);
99            power_f = self.eval(power_f * power_f);
100        }
101        result
102    }
103
104    /// Exponetiates a variable to a list of reversed bits with a given length.
105    ///
106    /// Reference: [p3_util::reverse_bits_len]
107    pub fn exp_reverse_bits_len<V>(
108        &mut self,
109        x: V,
110        power_bits: &Array<C, Var<C::N>>,
111        bit_len: impl Into<Usize<C::N>>,
112    ) -> V
113    where
114        V::Expression: AbstractField,
115        V: Copy + Mul<Output = V::Expression> + Variable<C>,
116    {
117        let result = self.eval(V::Expression::one());
118        let power_f: V = self.eval(x);
119        let bit_len = bit_len.into().materialize(self);
120        let bit_len_plus_one: Var<_> = self.eval(bit_len + C::N::one());
121
122        self.range(1, bit_len_plus_one).for_each(|i, builder| {
123            let index: Var<C::N> = builder.eval(bit_len - i);
124            let bit = builder.get(power_bits, index);
125            builder
126                .if_eq(bit, C::N::one())
127                .then(|builder| builder.assign(result, result * power_f));
128            builder.assign(power_f, power_f * power_f);
129        });
130        result
131    }
132
133    /// A version of `exp_reverse_bits_len` that uses the ExpReverseBitsLen precompile.
134    pub fn exp_reverse_bits_len_fast(
135        &mut self,
136        x: Felt<C::F>,
137        power_bits: &Array<C, Var<C::N>>,
138        bit_len: impl Into<Usize<C::N>>,
139    ) -> Felt<C::F> {
140        // Instantiate an array of length one and store the value of x.
141        let mut x_copy_arr: Array<C, Felt<C::F>> = self.dyn_array(1);
142        self.set(&mut x_copy_arr, 0, x);
143
144        // Get a pointer to the address holding x.
145        let x_copy_arr_ptr = match x_copy_arr {
146            Array::Dyn(ptr, _) => ptr,
147            _ => panic!("Expected a dynamic array"),
148        };
149
150        // Materialize the bit length as a Var.
151        let bit_len_var = bit_len.into().materialize(self);
152        // Get a pointer to the array of bits in the exponent.
153        let ptr = match power_bits {
154            Array::Dyn(ptr, _) => ptr,
155            _ => panic!("Expected a dynamic array"),
156        };
157
158        // Call the DslIR instruction ExpReverseBitsLen, which modifies the memory pointed to by
159        // `x_copy_arr_ptr`.
160        self.push_op(DslIr::ExpReverseBitsLen(x_copy_arr_ptr, ptr.address, bit_len_var));
161
162        // Return the value stored at the address pointed to by `x_copy_arr_ptr`.
163        self.get(&x_copy_arr, 0)
164    }
165
166    /// Exponentiates a variable to a list of bits in little endian.
167    pub fn exp_power_of_2_v<V>(
168        &mut self,
169        base: impl Into<V::Expression>,
170        power_log: impl Into<Usize<C::N>>,
171    ) -> V
172    where
173        V: Variable<C> + Copy + Mul<Output = V::Expression>,
174    {
175        let mut result: V = self.eval(base);
176        let power_log: Usize<_> = power_log.into();
177        match power_log {
178            Usize::Var(power_log) => {
179                self.range(0, power_log)
180                    .for_each(|_, builder| builder.assign(result, result * result));
181            }
182            Usize::Const(power_log) => {
183                for _ in 0..power_log {
184                    result = self.eval(result * result);
185                }
186            }
187        }
188        result
189    }
190
191    /// Exponentiates a variable to a list of bits in little endian inside a circuit.
192    pub fn exp_power_of_2_v_circuit<V>(
193        &mut self,
194        base: impl Into<V::Expression>,
195        power_log: usize,
196    ) -> V
197    where
198        V: Copy + Mul<Output = V::Expression> + Variable<C>,
199    {
200        let mut result: V = self.eval(base);
201        for _ in 0..power_log {
202            result = self.eval(result * result)
203        }
204        result
205    }
206
207    /// Multiplies `base` by `2^{log_power}`.
208    pub fn sll<V>(&mut self, base: impl Into<V::Expression>, shift: Usize<C::N>) -> V
209    where
210        V: Variable<C> + Copy + Add<Output = V::Expression>,
211    {
212        let result: V = self.eval(base);
213        self.range(0, shift).for_each(|_, builder| builder.assign(result, result + result));
214        result
215    }
216
217    /// Creates an ext from a slice of felts.
218    pub fn ext_from_base_slice(&mut self, arr: &[Felt<C::F>]) -> Ext<C::F, C::EF> {
219        assert!(arr.len() <= <C::EF as AbstractExtensionField::<C::F>>::D);
220        let mut res = SymbolicExt::from_f(C::EF::zero());
221        for i in 0..arr.len() {
222            res += arr[i] * SymbolicExt::from_f(C::EF::monomial(i));
223        }
224        self.eval(res)
225    }
226
227    pub fn felts2ext(&mut self, felts: &[Felt<C::F>]) -> Ext<C::F, C::EF> {
228        assert_eq!(felts.len(), 4);
229        let out: Ext<C::F, C::EF> = self.uninit();
230        self.push_op(DslIr::CircuitFelts2Ext(felts.try_into().unwrap(), out));
231        out
232    }
233
234    /// Converts an ext to a slice of felts.
235    pub fn ext2felt(&mut self, value: Ext<C::F, C::EF>) -> Array<C, Felt<C::F>> {
236        let result = self.dyn_array(4);
237        self.push_op(DslIr::HintExt2Felt(result.clone(), value));
238
239        // Verify that the decomposed extension element is correct.
240        let mut reconstructed_ext: Ext<C::F, C::EF> = self.constant(C::EF::zero());
241        for i in 0..4 {
242            let felt = self.get(&result, i);
243            let monomial: Ext<C::F, C::EF> = self.constant(C::EF::monomial(i));
244            reconstructed_ext = self.eval(reconstructed_ext + monomial * felt);
245        }
246
247        self.assert_ext_eq(reconstructed_ext, value);
248
249        result
250    }
251
252    /// Converts an ext to a slice of felts inside a circuit.
253    pub fn ext2felt_circuit(&mut self, value: Ext<C::F, C::EF>) -> [Felt<C::F>; 4] {
254        let a = self.uninit();
255        let b = self.uninit();
256        let c = self.uninit();
257        let d = self.uninit();
258        self.push_op(DslIr::CircuitExt2Felt([a, b, c, d], value));
259        [a, b, c, d]
260    }
261}
sp1_recursion_compiler/ir/utils.rs

sp1_recursion_compiler/ir/
utils.rs
