openvm_stark_backend/air_builders/symbolic/
symbolic_variable.rs1use 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 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 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#[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}