openvm_native_compiler/ir/
utils.rs

1use std::ops::{Add, Mul};
2
3use openvm_native_compiler_derive::iter_zip;
4use openvm_stark_backend::p3_field::{FieldAlgebra, FieldExtensionAlgebra, PrimeField};
5
6use super::{
7    Array, ArrayLike, Builder, CanSelect, Config, DslIr, Ext, Felt, MemIndex, RVar, SymbolicExt,
8    Var, Variable,
9};
10
11pub const NUM_LIMBS: usize = 32;
12pub const LIMB_BITS: usize = 8;
13
14/// Converts a prime field element to a usize.
15pub fn prime_field_to_usize<F: PrimeField>(x: F) -> usize {
16    let biguint = x.as_canonical_biguint();
17    let digits = biguint.to_u64_digits();
18    if digits.is_empty() {
19        return 0;
20    }
21    assert!(digits.len() == 1, "Prime field element is too large");
22    digits[0] as usize
23}
24
25impl<C: Config> Builder<C> {
26    /// The generator for the field.
27    ///
28    /// Reference: [`openvm_stark_sdk::p3_baby_bear::BabyBear`]
29    pub fn generator(&mut self) -> Felt<C::F> {
30        self.eval(C::F::from_canonical_u32(31))
31    }
32
33    /// Select a variable based on a condition.
34    pub fn select_v(&mut self, cond: Var<C::N>, a: Var<C::N>, b: Var<C::N>) -> Var<C::N> {
35        let c = self.uninit();
36        if self.flags.static_only {
37            self.operations.push(DslIr::CircuitSelectV(cond, a, b, c));
38        } else {
39            self.if_eq(cond, C::N::ONE).then_or_else(
40                |builder| builder.assign(&c, a),
41                |builder| builder.assign(&c, b),
42            );
43        }
44        c
45    }
46
47    /// Select a felt based on a condition.
48    pub fn select_f(&mut self, cond: Var<C::N>, a: Felt<C::F>, b: Felt<C::F>) -> Felt<C::F> {
49        let c = self.uninit();
50        if self.flags.static_only {
51            self.operations.push(DslIr::CircuitSelectF(cond, a, b, c));
52        } else {
53            self.if_eq(cond, C::N::ONE).then_or_else(
54                |builder| builder.assign(&c, a),
55                |builder| builder.assign(&c, b),
56            );
57        }
58        c
59    }
60
61    /// Select an extension based on a condition.
62    pub fn select_ef(
63        &mut self,
64        cond: Var<C::N>,
65        a: Ext<C::F, C::EF>,
66        b: Ext<C::F, C::EF>,
67    ) -> Ext<C::F, C::EF> {
68        let c = self.uninit();
69        if self.flags.static_only {
70            self.operations.push(DslIr::CircuitSelectE(cond, a, b, c));
71        } else {
72            self.if_eq(cond, C::N::ONE).then_or_else(
73                |builder| builder.assign(&c, a),
74                |builder| builder.assign(&c, b),
75            );
76        }
77        c
78    }
79
80    /// Exponentiates a variable to a list of big endian bits with a given length.
81    ///
82    /// Example: if power_bits = [1, 0, 1, 0], then the result should be x^8 * x^2 = x^10.
83    pub fn exp_bits_big_endian<V>(&mut self, x: V, power_bits: &Array<C, Var<C::N>>) -> V
84    where
85        V::Expression: FieldAlgebra,
86        V: Copy + Mul<Output = V::Expression> + Variable<C> + CanSelect<C>,
87    {
88        let result: V = self.eval(V::Expression::ONE);
89        let power_f: V = self.eval(x);
90        let one_var: V = self.eval(V::Expression::ONE);
91
92        // Implements a square-and-multiply algorithm.
93        iter_zip!(self, power_bits).for_each(|ptr_vec, builder| {
94            let bit = builder.iter_ptr_get(power_bits, ptr_vec[0]);
95            builder.assign(&result, result * result);
96            let mul = V::select(builder, bit, power_f, one_var);
97            builder.assign(&result, result * mul);
98        });
99
100        result
101    }
102
103    /// Exponentiates a variable to a list of bits in little endian.
104    pub fn exp_power_of_2_v<V>(
105        &mut self,
106        base: impl Into<V::Expression>,
107        power_log: impl Into<RVar<C::N>>,
108    ) -> V
109    where
110        V: Variable<C> + Copy + Mul<Output = V::Expression>,
111    {
112        let result: V = self.eval(base);
113        let power_log = power_log.into();
114        self.range(0, power_log)
115            .for_each(|_, builder| builder.assign(&result, result * result));
116        result
117    }
118
119    /// Multiplies `base` by `2^{log_power}`.
120    pub fn sll<V>(&mut self, base: impl Into<V::Expression>, shift: RVar<C::N>) -> V
121    where
122        V: Variable<C> + Clone + Add<Output = V::Expression>,
123    {
124        let result: V = self.eval(base);
125        self.range(0, shift)
126            .for_each(|_, builder| builder.assign(&result, result.clone() + result.clone()));
127        result
128    }
129
130    /// Creates an ext from a slice of felts.
131    pub fn ext_from_base_slice(&mut self, arr: &[Felt<C::F>]) -> Ext<C::F, C::EF> {
132        assert!(arr.len() <= <C::EF as FieldExtensionAlgebra<C::F>>::D);
133        let mut res = SymbolicExt::from_f(C::EF::ZERO);
134        for i in 0..arr.len() {
135            res += arr[i] * SymbolicExt::from_f(C::EF::monomial(i));
136        }
137        self.eval(res)
138    }
139
140    pub fn felts2ext(&mut self, felts: &[Felt<C::F>]) -> Ext<C::F, C::EF> {
141        assert_eq!(felts.len(), 4);
142        let out: Ext<C::F, C::EF> = self.uninit();
143        self.push(DslIr::CircuitFelts2Ext(felts.try_into().unwrap(), out));
144        out
145    }
146
147    /// Converts an ext to a slice of felts.
148    pub fn ext2felt(&mut self, value: Ext<C::F, C::EF>) -> Array<C, Felt<C::F>> {
149        if self.flags.static_only {
150            let felts = self.ext2felt_circuit(value);
151            self.vec(felts.to_vec())
152        } else {
153            let result = self.array(C::EF::D);
154            let index = MemIndex {
155                index: RVar::zero(),
156                offset: 0,
157                size: C::EF::D,
158            };
159            if let Array::Dyn(ptr, _) = &result {
160                self.store(*ptr, index, value);
161            } else {
162                unreachable!()
163            }
164            result
165        }
166    }
167
168    /// Converts an ext to a slice of felts inside a circuit.
169    pub fn ext2felt_circuit(&mut self, value: Ext<C::F, C::EF>) -> [Felt<C::F>; 4] {
170        let a = self.uninit();
171        let b = self.uninit();
172        let c = self.uninit();
173        let d = self.uninit();
174        self.operations
175            .push(DslIr::CircuitExt2Felt([a, b, c, d], value));
176        [a, b, c, d]
177    }
178    pub fn felt_reduce_circuit(&mut self, value: Felt<C::F>) {
179        self.operations.push(DslIr::CircuitFeltReduce(value));
180    }
181    pub fn ext_reduce_circuit(&mut self, value: Ext<C::F, C::EF>) {
182        self.operations.push(DslIr::CircuitExtReduce(value));
183    }
184}