Skip to main content

miden_assembly_syntax/ast/constants/
expr.rs

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// CONSTANT EXPRESSION
16// ================================================================================================
17
18/// Represents a constant expression or value in Miden Assembly syntax.
19#[derive(Clone)]
20#[repr(u8)]
21pub enum ConstantExpr {
22    /// A literal [`Felt`] value.
23    Int(Span<IntValue>),
24    /// A reference to another constant.
25    Var(Span<Arc<Path>>),
26    /// An binary arithmetic operator.
27    BinaryOp {
28        span: SourceSpan,
29        op: ConstantOp,
30        lhs: Box<ConstantExpr>,
31        rhs: Box<ConstantExpr>,
32    },
33    /// A plain spanned string.
34    String(Ident),
35    /// A literal ['WordValue'].
36    Word(Span<WordValue>),
37    /// A spanned string with a [`HashKind`] showing to which type of value the given string should
38    /// be hashed.
39    Hash(HashKind, Ident),
40}
41
42impl ConstantExpr {
43    /// Returns true if this expression is already evaluated to a concrete value
44    pub fn is_value(&self) -> bool {
45        matches!(self, Self::Int(_) | Self::Word(_) | Self::Hash(_, _) | Self::String(_))
46    }
47
48    /// Unwrap an [`IntValue`] from this expression or panic.
49    ///
50    /// This is used in places where we expect the expression to have been folded to an integer,
51    /// otherwise a bug occurred.
52    #[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    /// Unwrap a [`Felt`] value from this expression or panic.
61    ///
62    /// This is used in places where we expect the expression to have been folded to a felt value,
63    /// otherwise a bug occurred.
64    #[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    /// Unwrap a [`Arc<str>`] value from this expression or panic.
73    ///
74    /// This is used in places where we expect the expression to have been folded to a string value,
75    /// otherwise a bug occurred.
76    #[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    /// Unwrap a [ConstantValue] from this expression or panic.
85    ///
86    /// This is used in places where we expect the expression to have been folded to a concrete
87    /// value, otherwise a bug occurred.
88    #[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    /// Try to convert this expression into a [ConstantValue], if the expression is a value.
95    ///
96    /// Returns `Err` if the expression cannot be represented as a [ConstantValue].
97    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    /// Get the [ConstantValue] representation of this expression, if it is a value.
108    ///
109    /// Returns `None` if the expression cannot be represented as a [ConstantValue].
110    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    /// Get any references to other symbols present in this expression
121    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// CONSTANT OPERATION
293// ================================================================================================
294
295/// Represents the set of binary arithmetic operators supported in Miden Assembly syntax.
296#[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        // SAFETY: This is safe because we have given this enum a
344        // primitive representation with #[repr(u8)], with the first
345        // field of the underlying union-of-structs the discriminant
346        //
347        // See the section on "accessing the numeric value of the discriminant"
348        // here: https://doc.rust-lang.org/std/mem/fn.discriminant.html
349        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// HASH KIND
404// ================================================================================================
405
406/// Represents the type of the final value to which some string value should be converted.
407#[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    /// Reduce a string to a word using Blake3 hash function
415    Word,
416    /// Reduce a string to a felt using Blake3 hash function (via 64-bit reduction)
417    Event,
418}
419
420impl HashKind {
421    const fn tag(&self) -> u8 {
422        // SAFETY: This is safe because we have given this enum a
423        // primitive representation with #[repr(u8)], with the first
424        // field of the underlying union-of-structs the discriminant
425        //
426        // See the section on "accessing the numeric value of the discriminant"
427        // here: https://doc.rust-lang.org/std/mem/fn.discriminant.html
428        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}