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                // A newline outside parentheses terminates a constant declaration.
240                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// CONSTANT OPERATION
289// ================================================================================================
290
291/// Represents the set of binary arithmetic operators supported in Miden Assembly syntax.
292#[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        // SAFETY: This is safe because we have given this enum a
340        // primitive representation with #[repr(u8)], with the first
341        // field of the underlying union-of-structs the discriminant
342        //
343        // See the section on "accessing the numeric value of the discriminant"
344        // here: https://doc.rust-lang.org/std/mem/fn.discriminant.html
345        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// HASH KIND
400// ================================================================================================
401
402/// Represents the type of the final value to which some string value should be converted.
403#[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    /// Reduce a string to a word using Blake3 hash function
411    Word,
412    /// Reduce a string to a felt using Blake3 hash function (via 64-bit reduction)
413    Event,
414}
415
416impl HashKind {
417    const fn tag(&self) -> u8 {
418        // SAFETY: This is safe because we have given this enum a
419        // primitive representation with #[repr(u8)], with the first
420        // field of the underlying union-of-structs the discriminant
421        //
422        // See the section on "accessing the numeric value of the discriminant"
423        // here: https://doc.rust-lang.org/std/mem/fn.discriminant.html
424        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}