Skip to main content

sp1_hypercube/ir/
expr_impl.rs

1use std::{
2    iter::{Product, Sum},
3    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
4};
5
6use slop_algebra::{extension::BinomialExtensionField, AbstractExtensionField, AbstractField};
7use sp1_primitives::SP1Field;
8
9use crate::ir::{BinOp, ExprExtRef, ExprRef, GLOBAL_AST};
10
11pub(crate) type F = SP1Field;
12pub(crate) type EF = BinomialExtensionField<SP1Field, 4>;
13pub(crate) type Expr = ExprRef<F>;
14pub(crate) type ExprExt = ExprExtRef<EF>;
15
16impl AbstractField for Expr {
17    type F = F;
18
19    fn zero() -> Self {
20        F::zero().into()
21    }
22    fn one() -> Self {
23        F::one().into()
24    }
25    fn two() -> Self {
26        F::two().into()
27    }
28    fn neg_one() -> Self {
29        F::neg_one().into()
30    }
31
32    fn from_f(f: Self::F) -> Self {
33        f.into()
34    }
35    fn from_bool(b: bool) -> Self {
36        F::from_bool(b).into()
37    }
38    fn from_canonical_u8(n: u8) -> Self {
39        F::from_canonical_u8(n).into()
40    }
41    fn from_canonical_u16(n: u16) -> Self {
42        F::from_canonical_u16(n).into()
43    }
44    fn from_canonical_u32(n: u32) -> Self {
45        F::from_canonical_u32(n).into()
46    }
47    fn from_canonical_u64(n: u64) -> Self {
48        F::from_canonical_u64(n).into()
49    }
50    fn from_canonical_usize(n: usize) -> Self {
51        F::from_canonical_usize(n).into()
52    }
53    fn from_wrapped_u32(n: u32) -> Self {
54        F::from_wrapped_u32(n).into()
55    }
56    fn from_wrapped_u64(n: u64) -> Self {
57        F::from_wrapped_u64(n).into()
58    }
59
60    fn generator() -> Self {
61        F::generator().into()
62    }
63}
64
65impl Add for Expr {
66    type Output = Self;
67
68    fn add(self, rhs: Self) -> Self::Output {
69        let mut ast = GLOBAL_AST.lock().unwrap();
70        ast.bin_op(BinOp::Add, self, rhs)
71    }
72}
73
74impl Sub for Expr {
75    type Output = Self;
76
77    fn sub(self, rhs: Self) -> Self::Output {
78        let mut ast = GLOBAL_AST.lock().unwrap();
79        ast.bin_op(BinOp::Sub, self, rhs)
80    }
81}
82
83impl Mul for Expr {
84    type Output = Self;
85
86    fn mul(self, rhs: Self) -> Self::Output {
87        let mut ast = GLOBAL_AST.lock().unwrap();
88        ast.bin_op(BinOp::Mul, self, rhs)
89    }
90}
91
92impl Neg for Expr {
93    type Output = Self;
94
95    fn neg(self) -> Self::Output {
96        let mut ast = GLOBAL_AST.lock().unwrap();
97        ast.negate(self)
98    }
99}
100
101impl Add<F> for Expr {
102    type Output = Self;
103
104    fn add(self, rhs: F) -> Self::Output {
105        self + Expr::from(rhs)
106    }
107}
108
109impl Sub<F> for Expr {
110    type Output = Self;
111
112    fn sub(self, rhs: F) -> Self::Output {
113        self - Expr::from(rhs)
114    }
115}
116
117impl Mul<F> for Expr {
118    type Output = Self;
119
120    fn mul(self, rhs: F) -> Self::Output {
121        self * Expr::from(rhs)
122    }
123}
124
125impl Sum for Expr {
126    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
127        iter.fold(F::zero().into(), Add::add)
128    }
129}
130
131impl Product for Expr {
132    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
133        iter.fold(F::one().into(), Mul::mul)
134    }
135}
136
137impl From<F> for Expr {
138    fn from(f: F) -> Self {
139        Expr::IrVar(crate::ir::IrVar::Constant(f))
140    }
141}
142
143impl Default for Expr {
144    fn default() -> Self {
145        F::zero().into()
146    }
147}
148
149impl AddAssign for Expr {
150    fn add_assign(&mut self, rhs: Self) {
151        *self = *self + rhs;
152    }
153}
154
155impl SubAssign for Expr {
156    fn sub_assign(&mut self, rhs: Self) {
157        *self = *self - rhs;
158    }
159}
160
161impl MulAssign for Expr {
162    fn mul_assign(&mut self, rhs: Self) {
163        *self = *self * rhs;
164    }
165}
166
167impl AbstractField for ExprExt {
168    type F = EF;
169
170    fn zero() -> Self {
171        EF::zero().into()
172    }
173    fn one() -> Self {
174        EF::one().into()
175    }
176    fn two() -> Self {
177        EF::two().into()
178    }
179    fn neg_one() -> Self {
180        EF::neg_one().into()
181    }
182
183    fn from_f(f: Self::F) -> Self {
184        f.into()
185    }
186    fn from_bool(b: bool) -> Self {
187        EF::from_bool(b).into()
188    }
189    fn from_canonical_u8(n: u8) -> Self {
190        EF::from_canonical_u8(n).into()
191    }
192    fn from_canonical_u16(n: u16) -> Self {
193        EF::from_canonical_u16(n).into()
194    }
195    fn from_canonical_u32(n: u32) -> Self {
196        EF::from_canonical_u32(n).into()
197    }
198    fn from_canonical_u64(n: u64) -> Self {
199        EF::from_canonical_u64(n).into()
200    }
201    fn from_canonical_usize(n: usize) -> Self {
202        EF::from_canonical_usize(n).into()
203    }
204    fn from_wrapped_u32(n: u32) -> Self {
205        EF::from_wrapped_u32(n).into()
206    }
207    fn from_wrapped_u64(n: u64) -> Self {
208        EF::from_wrapped_u64(n).into()
209    }
210
211    fn generator() -> Self {
212        EF::generator().into()
213    }
214}
215
216impl AbstractExtensionField<Expr> for ExprExt {
217    const D: usize = <EF as AbstractExtensionField<F>>::D;
218
219    fn from_base(b: Expr) -> Self {
220        let mut ast = GLOBAL_AST.lock().unwrap();
221        ast.ext_from_base(b)
222    }
223
224    fn from_base_slice(_: &[Expr]) -> Self {
225        todo!()
226    }
227
228    fn from_base_fn<F: FnMut(usize) -> Expr>(_: F) -> Self {
229        todo!()
230    }
231
232    fn as_base_slice(&self) -> &[Expr] {
233        todo!()
234    }
235}
236
237impl From<Expr> for ExprExt {
238    fn from(e: Expr) -> Self {
239        let mut ast = GLOBAL_AST.lock().unwrap();
240        ast.ext_from_base(e)
241    }
242}
243
244impl Add<Expr> for ExprExt {
245    type Output = Self;
246
247    fn add(self, rhs: Expr) -> Self::Output {
248        let mut ast = GLOBAL_AST.lock().unwrap();
249        ast.bin_op_base_ext(BinOp::Add, self, rhs)
250    }
251}
252
253impl Sub<Expr> for ExprExt {
254    type Output = Self;
255
256    fn sub(self, rhs: Expr) -> Self::Output {
257        let mut ast = GLOBAL_AST.lock().unwrap();
258        ast.bin_op_base_ext(BinOp::Sub, self, rhs)
259    }
260}
261
262impl Mul<Expr> for ExprExt {
263    type Output = Self;
264
265    fn mul(self, rhs: Expr) -> Self::Output {
266        let mut ast = GLOBAL_AST.lock().unwrap();
267        ast.bin_op_base_ext(BinOp::Mul, self, rhs)
268    }
269}
270
271impl MulAssign<Expr> for ExprExt {
272    fn mul_assign(&mut self, rhs: Expr) {
273        *self = *self * rhs;
274    }
275}
276
277impl AddAssign<Expr> for ExprExt {
278    fn add_assign(&mut self, rhs: Expr) {
279        *self = *self + rhs;
280    }
281}
282
283impl SubAssign<Expr> for ExprExt {
284    fn sub_assign(&mut self, rhs: Expr) {
285        *self = *self - rhs;
286    }
287}
288
289impl Add for ExprExt {
290    type Output = Self;
291
292    fn add(self, rhs: Self) -> Self::Output {
293        let mut ast = GLOBAL_AST.lock().unwrap();
294        ast.bin_op_ext(BinOp::Add, self, rhs)
295    }
296}
297
298impl Sub for ExprExt {
299    type Output = Self;
300
301    fn sub(self, rhs: Self) -> Self::Output {
302        let mut ast = GLOBAL_AST.lock().unwrap();
303        ast.bin_op_ext(BinOp::Sub, self, rhs)
304    }
305}
306
307impl Mul for ExprExt {
308    type Output = Self;
309
310    fn mul(self, rhs: Self) -> Self::Output {
311        let mut ast = GLOBAL_AST.lock().unwrap();
312        ast.bin_op_ext(BinOp::Mul, self, rhs)
313    }
314}
315
316impl Neg for ExprExt {
317    type Output = Self;
318
319    fn neg(self) -> Self::Output {
320        let mut ast = GLOBAL_AST.lock().unwrap();
321        ast.neg_ext(self)
322    }
323}
324
325impl From<EF> for ExprExt {
326    fn from(f: EF) -> Self {
327        ExprExtRef::ExtConstant(f)
328    }
329}
330
331impl Default for ExprExt {
332    fn default() -> Self {
333        EF::zero().into()
334    }
335}
336
337impl AddAssign for ExprExt {
338    fn add_assign(&mut self, rhs: Self) {
339        *self = *self + rhs;
340    }
341}
342
343impl SubAssign for ExprExt {
344    fn sub_assign(&mut self, rhs: Self) {
345        *self = *self - rhs;
346    }
347}
348
349impl MulAssign for ExprExt {
350    fn mul_assign(&mut self, rhs: Self) {
351        *self = *self * rhs;
352    }
353}
354
355impl Sum for ExprExt {
356    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
357        iter.fold(EF::zero().into(), Add::add)
358    }
359}
360
361impl Product for ExprExt {
362    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
363        iter.fold(EF::one().into(), Mul::mul)
364    }
365}