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 let single_line = lhs.render_operand(precedence, false)
240 + display(op)
241 + rhs.render_operand(precedence, true);
242 let multi_line = lhs.render_operand(precedence, false)
243 + nl()
244 + display(op)
245 + rhs.render_operand(precedence, true);
246 single_line | multi_line
247 },
248 }
249 }
250}
251
252impl Spanned for ConstantExpr {
253 fn span(&self) -> SourceSpan {
254 match self {
255 Self::Int(spanned) => spanned.span(),
256 Self::Word(spanned) => spanned.span(),
257 Self::Hash(_, spanned) => spanned.span(),
258 Self::Var(spanned) => spanned.span(),
259 Self::String(spanned) => spanned.span(),
260 Self::BinaryOp { span, .. } => *span,
261 }
262 }
263}
264
265#[cfg(feature = "arbitrary")]
266impl proptest::arbitrary::Arbitrary for ConstantExpr {
267 type Parameters = ();
268
269 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
270 use proptest::{arbitrary::any, prop_oneof, strategy::Strategy};
271
272 prop_oneof![
273 any::<IntValue>().prop_map(|n| Self::Int(Span::unknown(n))),
274 crate::arbitrary::path::constant_path_random_length(0)
275 .prop_map(|p| Self::Var(Span::unknown(p))),
276 any::<(ConstantOp, IntValue, IntValue)>().prop_map(|(op, lhs, rhs)| Self::BinaryOp {
277 span: SourceSpan::UNKNOWN,
278 op,
279 lhs: Box::new(ConstantExpr::Int(Span::unknown(lhs))),
280 rhs: Box::new(ConstantExpr::Int(Span::unknown(rhs))),
281 }),
282 any::<Ident>().prop_map(Self::String),
283 any::<WordValue>().prop_map(|word| Self::Word(Span::unknown(word))),
284 any::<(HashKind, Ident)>().prop_map(|(kind, s)| Self::Hash(kind, s)),
285 ]
286 .boxed()
287 }
288
289 type Strategy = proptest::prelude::BoxedStrategy<Self>;
290}
291
292#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
297#[repr(u8)]
298#[cfg_attr(
299 all(feature = "arbitrary", test),
300 miden_test_serialization_macros::serialization_test
301)]
302pub enum ConstantOp {
303 Add,
304 Sub,
305 Mul,
306 Div,
307 IntDiv,
308}
309
310impl ConstantOp {
311 const fn name(self) -> &'static str {
312 match self {
313 Self::Add => "Add",
314 Self::Sub => "Sub",
315 Self::Mul => "Mul",
316 Self::Div => "Div",
317 Self::IntDiv => "IntDiv",
318 }
319 }
320
321 pub(crate) const fn precedence(self) -> u8 {
322 match self {
323 Self::Add | Self::Sub => 1,
324 Self::Mul | Self::Div | Self::IntDiv => 2,
325 }
326 }
327}
328
329impl fmt::Display for ConstantOp {
330 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
331 match self {
332 Self::Add => f.write_str("+"),
333 Self::Sub => f.write_str("-"),
334 Self::Mul => f.write_str("*"),
335 Self::Div => f.write_str("/"),
336 Self::IntDiv => f.write_str("//"),
337 }
338 }
339}
340
341impl ConstantOp {
342 const fn tag(&self) -> u8 {
343 unsafe { *(self as *const Self).cast::<u8>() }
350 }
351}
352
353impl Serializable for ConstantOp {
354 fn write_into<W: ByteWriter>(&self, target: &mut W) {
355 target.write_u8(self.tag());
356 }
357}
358
359impl Deserializable for ConstantOp {
360 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
361 const ADD: u8 = ConstantOp::Add.tag();
362 const SUB: u8 = ConstantOp::Sub.tag();
363 const MUL: u8 = ConstantOp::Mul.tag();
364 const DIV: u8 = ConstantOp::Div.tag();
365 const INT_DIV: u8 = ConstantOp::IntDiv.tag();
366
367 match source.read_u8()? {
368 ADD => Ok(Self::Add),
369 SUB => Ok(Self::Sub),
370 MUL => Ok(Self::Mul),
371 DIV => Ok(Self::Div),
372 INT_DIV => Ok(Self::IntDiv),
373 invalid => Err(DeserializationError::InvalidValue(format!(
374 "unexpected ConstantOp tag: '{invalid}'"
375 ))),
376 }
377 }
378}
379
380#[cfg(feature = "arbitrary")]
381impl proptest::arbitrary::Arbitrary for ConstantOp {
382 type Parameters = ();
383
384 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
385 use proptest::{
386 prop_oneof,
387 strategy::{Just, Strategy},
388 };
389
390 prop_oneof![
391 Just(Self::Add),
392 Just(Self::Sub),
393 Just(Self::Mul),
394 Just(Self::Div),
395 Just(Self::IntDiv),
396 ]
397 .boxed()
398 }
399
400 type Strategy = proptest::prelude::BoxedStrategy<Self>;
401}
402
403#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
408#[repr(u8)]
409#[cfg_attr(
410 all(feature = "arbitrary", test),
411 miden_test_serialization_macros::serialization_test
412)]
413pub enum HashKind {
414 Word,
416 Event,
418}
419
420impl HashKind {
421 const fn tag(&self) -> u8 {
422 unsafe { *(self as *const Self).cast::<u8>() }
429 }
430}
431
432impl fmt::Display for HashKind {
433 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
434 match self {
435 Self::Word => f.write_str("word"),
436 Self::Event => f.write_str("event"),
437 }
438 }
439}
440
441#[cfg(feature = "arbitrary")]
442impl proptest::arbitrary::Arbitrary for HashKind {
443 type Parameters = ();
444
445 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
446 use proptest::{
447 prop_oneof,
448 strategy::{Just, Strategy},
449 };
450
451 prop_oneof![Just(Self::Word), Just(Self::Event),].boxed()
452 }
453
454 type Strategy = proptest::prelude::BoxedStrategy<Self>;
455}
456
457impl Serializable for HashKind {
458 fn write_into<W: ByteWriter>(&self, target: &mut W) {
459 target.write_u8(self.tag());
460 }
461}
462
463impl Deserializable for HashKind {
464 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
465 const WORD: u8 = HashKind::Word.tag();
466 const EVENT: u8 = HashKind::Event.tag();
467
468 match source.read_u8()? {
469 WORD => Ok(Self::Word),
470 EVENT => Ok(Self::Event),
471 invalid => Err(DeserializationError::InvalidValue(format!(
472 "unexpected HashKind tag: '{invalid}'"
473 ))),
474 }
475 }
476}