openvm_stark_backend/air_builders/symbolic/
symbolic_variable.rs
1use 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 {
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 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#[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}