1use alloc::{boxed::Box, sync::Arc, vec::Vec};
2use core::fmt;
3
4use miden_core::serde::{
5 ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
6};
7use miden_debug_types::{SourceSpan, Span, Spanned};
8
9use crate::{
10 Felt, Path,
11 ast::{ConstantValue, Ident},
12 parser::{IntValue, WordValue},
13};
14
15#[derive(Clone)]
20#[repr(u8)]
21pub enum ConstantExpr {
22 Int(Span<IntValue>),
24 Var(Span<Arc<Path>>),
26 BinaryOp {
28 span: SourceSpan,
29 op: ConstantOp,
30 lhs: Box<ConstantExpr>,
31 rhs: Box<ConstantExpr>,
32 },
33 String(Ident),
35 Word(Span<WordValue>),
37 Hash(HashKind, Ident),
40}
41
42impl ConstantExpr {
43 pub fn is_value(&self) -> bool {
45 matches!(self, Self::Int(_) | Self::Word(_) | Self::Hash(_, _) | Self::String(_))
46 }
47
48 #[track_caller]
53 pub fn expect_int(&self) -> IntValue {
54 match self {
55 Self::Int(spanned) => spanned.into_inner(),
56 other => panic!("expected constant expression to be a literal, got {other:#?}"),
57 }
58 }
59
60 #[track_caller]
65 pub fn expect_felt(&self) -> Felt {
66 match self {
67 Self::Int(spanned) => Felt::new_unchecked(spanned.inner().as_int()),
68 other => panic!("expected constant expression to be a literal, got {other:#?}"),
69 }
70 }
71
72 #[track_caller]
77 pub fn expect_string(&self) -> Arc<str> {
78 match self {
79 Self::String(spanned) => spanned.clone().into_inner(),
80 other => panic!("expected constant expression to be a string, got {other:#?}"),
81 }
82 }
83
84 #[track_caller]
89 pub fn expect_value(&self) -> ConstantValue {
90 self.as_value()
91 .unwrap_or_else(|| panic!("expected constant expression to be a value, got {self:#?}"))
92 }
93
94 pub fn into_value(self) -> Result<ConstantValue, Self> {
98 match self {
99 Self::Int(value) => Ok(ConstantValue::Int(value)),
100 Self::String(value) => Ok(ConstantValue::String(value)),
101 Self::Word(value) => Ok(ConstantValue::Word(value)),
102 Self::Hash(kind, value) => Ok(ConstantValue::Hash(kind, value)),
103 expr @ (Self::BinaryOp { .. } | Self::Var(_)) => Err(expr),
104 }
105 }
106
107 pub fn as_value(&self) -> Option<ConstantValue> {
111 match self {
112 Self::Int(value) => Some(ConstantValue::Int(*value)),
113 Self::String(value) => Some(ConstantValue::String(value.clone())),
114 Self::Word(value) => Some(ConstantValue::Word(*value)),
115 Self::Hash(kind, value) => Some(ConstantValue::Hash(*kind, value.clone())),
116 Self::BinaryOp { .. } | Self::Var(_) => None,
117 }
118 }
119
120 pub fn references(&self) -> Vec<Span<Arc<Path>>> {
122 use alloc::collections::BTreeSet;
123
124 let mut worklist = smallvec::SmallVec::<[_; 4]>::from_slice(&[self]);
125 let mut references = BTreeSet::new();
126
127 while let Some(ty) = worklist.pop() {
128 match ty {
129 Self::Int(_) | Self::Word(_) | Self::String(_) | Self::Hash(..) => {},
130 Self::Var(path) => {
131 references.insert(path.clone());
132 },
133 Self::BinaryOp { lhs, rhs, .. } => {
134 worklist.push(lhs);
135 worklist.push(rhs);
136 },
137 }
138 }
139
140 references.into_iter().collect()
141 }
142
143 fn render_operand(&self, parent_precedence: u8, is_rhs: bool) -> crate::prettier::Document {
144 use crate::prettier::{PrettyPrint, const_text};
145
146 let rendered = self.render();
147 let Self::BinaryOp { op, .. } = self else {
148 return rendered;
149 };
150
151 let precedence = op.precedence();
152 if precedence < parent_precedence || (is_rhs && precedence == parent_precedence) {
153 const_text("(") + rendered + const_text(")")
154 } else {
155 rendered
156 }
157 }
158}
159
160impl Eq for ConstantExpr {}
161
162impl PartialEq for ConstantExpr {
163 fn eq(&self, other: &Self) -> bool {
164 match (self, other) {
165 (Self::Int(x), Self::Int(y)) => x == y,
166 (Self::Int(_), _) => false,
167 (Self::Word(x), Self::Word(y)) => x == y,
168 (Self::Word(_), _) => false,
169 (Self::Var(x), Self::Var(y)) => x == y,
170 (Self::Var(_), _) => false,
171 (Self::String(x), Self::String(y)) => x == y,
172 (Self::String(_), _) => false,
173 (Self::Hash(x_hk, x_i), Self::Hash(y_hk, y_i)) => x_i == y_i && x_hk == y_hk,
174 (Self::Hash(..), _) => false,
175 (
176 Self::BinaryOp { op: lop, lhs: llhs, rhs: lrhs, .. },
177 Self::BinaryOp { op: rop, lhs: rlhs, rhs: rrhs, .. },
178 ) => lop == rop && llhs == rlhs && lrhs == rrhs,
179 (Self::BinaryOp { .. }, _) => false,
180 }
181 }
182}
183
184impl core::hash::Hash for ConstantExpr {
185 fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
186 core::mem::discriminant(self).hash(state);
187 match self {
188 Self::Int(value) => value.hash(state),
189 Self::Word(value) => value.hash(state),
190 Self::String(value) => value.hash(state),
191 Self::Var(value) => value.hash(state),
192 Self::Hash(hash_kind, string) => {
193 hash_kind.hash(state);
194 string.hash(state);
195 },
196 Self::BinaryOp { op, lhs, rhs, .. } => {
197 op.hash(state);
198 lhs.hash(state);
199 rhs.hash(state);
200 },
201 }
202 }
203}
204
205impl fmt::Debug for ConstantExpr {
206 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
207 match self {
208 Self::Int(lit) => fmt::Debug::fmt(&**lit, f),
209 Self::Word(lit) => fmt::Debug::fmt(&**lit, f),
210 Self::Var(path) => fmt::Debug::fmt(path, f),
211 Self::String(name) => fmt::Debug::fmt(&**name, f),
212 Self::Hash(hash_kind, str) => {
213 f.debug_tuple("Hash").field(hash_kind).field(str).finish()
214 },
215 Self::BinaryOp { op, lhs, rhs, .. } => {
216 f.debug_tuple(op.name()).field(lhs).field(rhs).finish()
217 },
218 }
219 }
220}
221
222impl crate::prettier::PrettyPrint for ConstantExpr {
223 fn render(&self) -> crate::prettier::Document {
224 use crate::prettier::*;
225
226 match self {
227 Self::Int(literal) => literal.render(),
228 Self::Word(literal) => literal.render(),
229 Self::Var(path) => display(path),
230 Self::String(ident) => text(format!("\"{}\"", ident.as_str().escape_debug())),
231 Self::Hash(hash_kind, str) => flatten(
232 display(hash_kind)
233 + const_text("(")
234 + text(format!("\"{}\"", str.as_str().escape_debug()))
235 + const_text(")"),
236 ),
237 Self::BinaryOp { op, lhs, rhs, .. } => {
238 let precedence = op.precedence();
239 lhs.render_operand(precedence, false)
241 + display(op)
242 + rhs.render_operand(precedence, true)
243 },
244 }
245 }
246}
247
248impl Spanned for ConstantExpr {
249 fn span(&self) -> SourceSpan {
250 match self {
251 Self::Int(spanned) => spanned.span(),
252 Self::Word(spanned) => spanned.span(),
253 Self::Hash(_, spanned) => spanned.span(),
254 Self::Var(spanned) => spanned.span(),
255 Self::String(spanned) => spanned.span(),
256 Self::BinaryOp { span, .. } => *span,
257 }
258 }
259}
260
261#[cfg(feature = "arbitrary")]
262impl proptest::arbitrary::Arbitrary for ConstantExpr {
263 type Parameters = ();
264
265 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
266 use proptest::{arbitrary::any, prop_oneof, strategy::Strategy};
267
268 prop_oneof![
269 any::<IntValue>().prop_map(|n| Self::Int(Span::unknown(n))),
270 crate::arbitrary::path::constant_path_random_length(0)
271 .prop_map(|p| Self::Var(Span::unknown(p))),
272 any::<(ConstantOp, IntValue, IntValue)>().prop_map(|(op, lhs, rhs)| Self::BinaryOp {
273 span: SourceSpan::UNKNOWN,
274 op,
275 lhs: Box::new(ConstantExpr::Int(Span::unknown(lhs))),
276 rhs: Box::new(ConstantExpr::Int(Span::unknown(rhs))),
277 }),
278 any::<Ident>().prop_map(Self::String),
279 any::<WordValue>().prop_map(|word| Self::Word(Span::unknown(word))),
280 any::<(HashKind, Ident)>().prop_map(|(kind, s)| Self::Hash(kind, s)),
281 ]
282 .boxed()
283 }
284
285 type Strategy = proptest::prelude::BoxedStrategy<Self>;
286}
287
288#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
293#[repr(u8)]
294#[cfg_attr(
295 all(feature = "arbitrary", test),
296 miden_test_serialization_macros::serialization_test
297)]
298pub enum ConstantOp {
299 Add,
300 Sub,
301 Mul,
302 Div,
303 IntDiv,
304}
305
306impl ConstantOp {
307 const fn name(self) -> &'static str {
308 match self {
309 Self::Add => "Add",
310 Self::Sub => "Sub",
311 Self::Mul => "Mul",
312 Self::Div => "Div",
313 Self::IntDiv => "IntDiv",
314 }
315 }
316
317 pub(crate) const fn precedence(self) -> u8 {
318 match self {
319 Self::Add | Self::Sub => 1,
320 Self::Mul | Self::Div | Self::IntDiv => 2,
321 }
322 }
323}
324
325impl fmt::Display for ConstantOp {
326 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
327 match self {
328 Self::Add => f.write_str("+"),
329 Self::Sub => f.write_str("-"),
330 Self::Mul => f.write_str("*"),
331 Self::Div => f.write_str("/"),
332 Self::IntDiv => f.write_str("//"),
333 }
334 }
335}
336
337impl ConstantOp {
338 const fn tag(&self) -> u8 {
339 unsafe { *(self as *const Self).cast::<u8>() }
346 }
347}
348
349impl Serializable for ConstantOp {
350 fn write_into<W: ByteWriter>(&self, target: &mut W) {
351 target.write_u8(self.tag());
352 }
353}
354
355impl Deserializable for ConstantOp {
356 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
357 const ADD: u8 = ConstantOp::Add.tag();
358 const SUB: u8 = ConstantOp::Sub.tag();
359 const MUL: u8 = ConstantOp::Mul.tag();
360 const DIV: u8 = ConstantOp::Div.tag();
361 const INT_DIV: u8 = ConstantOp::IntDiv.tag();
362
363 match source.read_u8()? {
364 ADD => Ok(Self::Add),
365 SUB => Ok(Self::Sub),
366 MUL => Ok(Self::Mul),
367 DIV => Ok(Self::Div),
368 INT_DIV => Ok(Self::IntDiv),
369 invalid => Err(DeserializationError::InvalidValue(format!(
370 "unexpected ConstantOp tag: '{invalid}'"
371 ))),
372 }
373 }
374}
375
376#[cfg(feature = "arbitrary")]
377impl proptest::arbitrary::Arbitrary for ConstantOp {
378 type Parameters = ();
379
380 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
381 use proptest::{
382 prop_oneof,
383 strategy::{Just, Strategy},
384 };
385
386 prop_oneof![
387 Just(Self::Add),
388 Just(Self::Sub),
389 Just(Self::Mul),
390 Just(Self::Div),
391 Just(Self::IntDiv),
392 ]
393 .boxed()
394 }
395
396 type Strategy = proptest::prelude::BoxedStrategy<Self>;
397}
398
399#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
404#[repr(u8)]
405#[cfg_attr(
406 all(feature = "arbitrary", test),
407 miden_test_serialization_macros::serialization_test
408)]
409pub enum HashKind {
410 Word,
412 Event,
414}
415
416impl HashKind {
417 const fn tag(&self) -> u8 {
418 unsafe { *(self as *const Self).cast::<u8>() }
425 }
426}
427
428impl fmt::Display for HashKind {
429 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
430 match self {
431 Self::Word => f.write_str("word"),
432 Self::Event => f.write_str("event"),
433 }
434 }
435}
436
437#[cfg(feature = "arbitrary")]
438impl proptest::arbitrary::Arbitrary for HashKind {
439 type Parameters = ();
440
441 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
442 use proptest::{
443 prop_oneof,
444 strategy::{Just, Strategy},
445 };
446
447 prop_oneof![Just(Self::Word), Just(Self::Event),].boxed()
448 }
449
450 type Strategy = proptest::prelude::BoxedStrategy<Self>;
451}
452
453impl Serializable for HashKind {
454 fn write_into<W: ByteWriter>(&self, target: &mut W) {
455 target.write_u8(self.tag());
456 }
457}
458
459impl Deserializable for HashKind {
460 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
461 const WORD: u8 = HashKind::Word.tag();
462 const EVENT: u8 = HashKind::Event.tag();
463
464 match source.read_u8()? {
465 WORD => Ok(Self::Word),
466 EVENT => Ok(Self::Event),
467 invalid => Err(DeserializationError::InvalidValue(format!(
468 "unexpected HashKind tag: '{invalid}'"
469 ))),
470 }
471 }
472}