openvm_stark_backend/air_builders/symbolic/
symbolic_variable.rs

1// Copied from uni-stark/src/symbolic_variable.rs.
2
3use core::{
4    marker::PhantomData,
5    ops::{Add, Mul, Sub},
6};
7
8use p3_field::Field;
9use serde::{Deserialize, Serialize};
10
11use super::symbolic_expression::SymbolicExpression;
12
13#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
14#[repr(C)]
15pub enum Entry {
16    Preprocessed {
17        offset: usize,
18    },
19    /// Main may be partitioned
20    Main {
21        part_index: usize,
22        offset: usize,
23    },
24    Permutation {
25        offset: usize,
26    },
27    Public,
28    Challenge,
29    Exposed,
30}
31
32impl Entry {
33    pub fn offset(&self) -> Option<usize> {
34        match self {
35            Entry::Preprocessed { offset } => Some(*offset),
36            Entry::Main { offset, .. } => Some(*offset),
37            Entry::Permutation { offset } => Some(*offset),
38            Entry::Public => None,
39            Entry::Challenge => None,
40            Entry::Exposed => None,
41        }
42    }
43
44    /// Advance the internal offset of the entry by the given `offset`.
45    pub fn rotate(self, offset: usize) -> Self {
46        match self {
47            Entry::Preprocessed { offset: old_offset } => Entry::Preprocessed {
48                offset: old_offset + offset,
49            },
50            Entry::Main {
51                part_index,
52                offset: old_offset,
53            } => Entry::Main {
54                part_index,
55                offset: old_offset + offset,
56            },
57            Entry::Permutation { offset: old_offset } => Entry::Permutation {
58                offset: old_offset + offset,
59            },
60            Entry::Public | Entry::Challenge | Entry::Exposed => self,
61        }
62    }
63
64    pub fn next(self) -> Self {
65        self.rotate(1)
66    }
67}
68
69/// A variable within the evaluation window, i.e. a column in either the local or next row.
70#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Serialize, Deserialize)]
71#[repr(C)]
72pub struct SymbolicVariable<F> {
73    pub entry: Entry,
74    pub index: usize,
75    pub(crate) _phantom: PhantomData<F>,
76}
77
78impl<F: Field> SymbolicVariable<F> {
79    pub const fn new(entry: Entry, index: usize) -> Self {
80        Self {
81            entry,
82            index,
83            _phantom: PhantomData,
84        }
85    }
86
87    pub const fn degree_multiple(&self) -> usize {
88        match self.entry {
89            Entry::Preprocessed { .. } | Entry::Main { .. } | Entry::Permutation { .. } => 1,
90            Entry::Public | Entry::Challenge | Entry::Exposed => 0,
91        }
92    }
93
94    pub fn rotate(self, offset: usize) -> Self {
95        Self {
96            entry: self.entry.rotate(offset),
97            index: self.index,
98            _phantom: PhantomData,
99        }
100    }
101
102    pub fn next(self) -> Self {
103        self.rotate(1)
104    }
105}
106
107impl<F: Field> From<SymbolicVariable<F>> for SymbolicExpression<F> {
108    fn from(value: SymbolicVariable<F>) -> Self {
109        SymbolicExpression::Variable(value)
110    }
111}
112
113impl<F: Field> Add for SymbolicVariable<F> {
114    type Output = SymbolicExpression<F>;
115
116    fn add(self, rhs: Self) -> Self::Output {
117        SymbolicExpression::from(self) + SymbolicExpression::from(rhs)
118    }
119}
120
121impl<F: Field> Add<F> for SymbolicVariable<F> {
122    type Output = SymbolicExpression<F>;
123
124    fn add(self, rhs: F) -> Self::Output {
125        SymbolicExpression::from(self) + SymbolicExpression::from(rhs)
126    }
127}
128
129impl<F: Field> Add<SymbolicExpression<F>> for SymbolicVariable<F> {
130    type Output = SymbolicExpression<F>;
131
132    fn add(self, rhs: SymbolicExpression<F>) -> Self::Output {
133        SymbolicExpression::from(self) + rhs
134    }
135}
136
137impl<F: Field> Add<SymbolicVariable<F>> for SymbolicExpression<F> {
138    type Output = Self;
139
140    fn add(self, rhs: SymbolicVariable<F>) -> Self::Output {
141        self + Self::from(rhs)
142    }
143}
144
145impl<F: Field> Sub for SymbolicVariable<F> {
146    type Output = SymbolicExpression<F>;
147
148    fn sub(self, rhs: Self) -> Self::Output {
149        SymbolicExpression::from(self) - SymbolicExpression::from(rhs)
150    }
151}
152
153impl<F: Field> Sub<F> for SymbolicVariable<F> {
154    type Output = SymbolicExpression<F>;
155
156    fn sub(self, rhs: F) -> Self::Output {
157        SymbolicExpression::from(self) - SymbolicExpression::from(rhs)
158    }
159}
160
161impl<F: Field> Sub<SymbolicExpression<F>> for SymbolicVariable<F> {
162    type Output = SymbolicExpression<F>;
163
164    fn sub(self, rhs: SymbolicExpression<F>) -> Self::Output {
165        SymbolicExpression::from(self) - rhs
166    }
167}
168
169impl<F: Field> Sub<SymbolicVariable<F>> for SymbolicExpression<F> {
170    type Output = Self;
171
172    fn sub(self, rhs: SymbolicVariable<F>) -> Self::Output {
173        self - Self::from(rhs)
174    }
175}
176
177impl<F: Field> Mul for SymbolicVariable<F> {
178    type Output = SymbolicExpression<F>;
179
180    fn mul(self, rhs: Self) -> Self::Output {
181        SymbolicExpression::from(self) * SymbolicExpression::from(rhs)
182    }
183}
184
185impl<F: Field> Mul<F> for SymbolicVariable<F> {
186    type Output = SymbolicExpression<F>;
187
188    fn mul(self, rhs: F) -> Self::Output {
189        SymbolicExpression::from(self) * SymbolicExpression::from(rhs)
190    }
191}
192
193impl<F: Field> Mul<SymbolicExpression<F>> for SymbolicVariable<F> {
194    type Output = SymbolicExpression<F>;
195
196    fn mul(self, rhs: SymbolicExpression<F>) -> Self::Output {
197        SymbolicExpression::from(self) * rhs
198    }
199}
200
201impl<F: Field> Mul<SymbolicVariable<F>> for SymbolicExpression<F> {
202    type Output = Self;
203
204    fn mul(self, rhs: SymbolicVariable<F>) -> Self::Output {
205        self * Self::from(rhs)
206    }
207}