Skip to main content

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