midenc_codegen_masm/
stack.rs

1use core::{
2    fmt,
3    ops::{Index, IndexMut},
4};
5
6use miden_core::{Felt, FieldElement};
7use midenc_hir::{AttributeValue, Immediate, Type, ValueRef};
8use smallvec::{SmallVec, smallvec};
9
10/// This represents a constraint an operand's usage at
11/// a given program point, namely when used as an instruction
12/// or block argument.
13#[derive(Debug, Copy, Clone)]
14pub enum Constraint {
15    /// The operand should be moved, consuming it
16    /// from the stack and making it unavailable for
17    /// further use.
18    Move,
19    /// The operand should be copied, preserving the
20    /// original value for later use.
21    Copy,
22}
23
24/// Represents the type of operand represented on the operand stack
25pub enum OperandType {
26    /// The operand is a literal, unassociated with any value in the IR
27    Const(Box<dyn AttributeValue>),
28    /// The operand is an SSA value of known type
29    Value(ValueRef),
30    /// The operand is an intermediate runtime value of a known type, but
31    /// unassociated with any value in the IR
32    Type(Type),
33}
34impl Clone for OperandType {
35    fn clone(&self) -> Self {
36        match self {
37            Self::Const(value) => Self::Const(value.clone_value()),
38            Self::Value(value) => Self::Value(*value),
39            Self::Type(ty) => Self::Type(ty.clone()),
40        }
41    }
42}
43impl OperandType {
44    /// Get the type representation of this operand
45    pub fn ty(&self) -> Type {
46        match self {
47            Self::Const(imm) => {
48                imm.downcast_ref::<Immediate>().expect("unexpected constant value type").ty()
49            }
50            Self::Value(value) => value.borrow().ty().clone(),
51            Self::Type(ty) => ty.clone(),
52        }
53    }
54}
55impl fmt::Debug for OperandType {
56    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
57        match self {
58            Self::Const(value) => write!(f, "Const({value:?})"),
59            Self::Value(value) => write!(f, "Value({value})"),
60            Self::Type(ty) => write!(f, "Type({ty})"),
61        }
62    }
63}
64impl Eq for OperandType {}
65impl PartialEq for OperandType {
66    fn eq(&self, other: &Self) -> bool {
67        match (self, other) {
68            (Self::Value(a), Self::Value(b)) => a == b,
69            (Self::Value(_), _) | (_, Self::Value(_)) => false,
70            (Self::Const(a), Self::Const(b)) => a == b,
71            (Self::Const(_), _) | (_, Self::Const(_)) => false,
72            (Self::Type(a), Self::Type(b)) => a == b,
73        }
74    }
75}
76impl PartialEq<Type> for OperandType {
77    fn eq(&self, other: &Type) -> bool {
78        match self {
79            Self::Type(a) => a == other,
80            _ => false,
81        }
82    }
83}
84impl PartialEq<Immediate> for OperandType {
85    fn eq(&self, other: &Immediate) -> bool {
86        match self {
87            Self::Const(a) => a.downcast_ref::<Immediate>().is_some_and(|a| a == other),
88            _ => false,
89        }
90    }
91}
92impl PartialEq<dyn AttributeValue> for OperandType {
93    fn eq(&self, other: &dyn AttributeValue) -> bool {
94        match self {
95            Self::Const(a) => a.as_ref() == other,
96            _ => false,
97        }
98    }
99}
100impl PartialEq<ValueRef> for OperandType {
101    fn eq(&self, other: &ValueRef) -> bool {
102        match self {
103            Self::Value(this) => this == other,
104            _ => false,
105        }
106    }
107}
108impl From<ValueRef> for OperandType {
109    fn from(value: ValueRef) -> Self {
110        Self::Value(value)
111    }
112}
113impl From<Type> for OperandType {
114    fn from(ty: Type) -> Self {
115        Self::Type(ty)
116    }
117}
118impl From<Immediate> for OperandType {
119    fn from(value: Immediate) -> Self {
120        Self::Const(Box::new(value))
121    }
122}
123impl From<Box<dyn AttributeValue>> for OperandType {
124    fn from(value: Box<dyn AttributeValue>) -> Self {
125        Self::Const(value)
126    }
127}
128
129/// This type represents a logical operand on the stack, which may consist
130/// of one or more "parts", up to a word in size, on the actual stack.
131///
132/// The [OperandStack] operates in terms of [Operand], but when emitting
133/// Miden Assembly, we must know how to translate operand-oriented operations
134/// into equivalent element-/word-oriented operations. This is accomplished
135/// by tracking the low-level representation of a given operand in this struct.
136#[derive(Debug, Clone, PartialEq, Eq)]
137pub struct Operand {
138    /// The section of stack corresponding to this operand, containing
139    /// up to a full word of elements. No chunk will ever exceed a word
140    /// in size. This field behaves like a miniature [OperandStack], i.e.
141    /// elements are pushed and popped off the end to modify it.
142    ///
143    /// An operand is encoded on this stack in order of lowest
144    /// addressed bytes first. For example, given a struct operand,
145    /// the first field of the struct will be closest to the top of
146    /// the stack.
147    word: SmallVec<[Type; 4]>,
148    /// The high-level operand represented by this item.
149    ///
150    /// If the operand stack is manipulated in such a way that the operand
151    /// is torn apart, say one field of a struct is popped; then this will
152    /// be set to a `Type` operand, representing what high-level information
153    /// we have about the remaining parts of the original operand on the stack.
154    operand: OperandType,
155}
156impl Default for Operand {
157    fn default() -> Self {
158        Self {
159            word: smallvec![Type::Felt],
160            operand: OperandType::Const(Box::new(Immediate::Felt(Felt::ZERO))),
161        }
162    }
163}
164impl PartialEq<ValueRef> for Operand {
165    #[inline(always)]
166    fn eq(&self, other: &ValueRef) -> bool {
167        self.operand.eq(other)
168    }
169}
170impl PartialEq<dyn AttributeValue> for Operand {
171    #[inline(always)]
172    fn eq(&self, other: &dyn AttributeValue) -> bool {
173        self.operand.eq(other)
174    }
175}
176impl PartialEq<Immediate> for Operand {
177    #[inline(always)]
178    fn eq(&self, other: &Immediate) -> bool {
179        self.operand.eq(other)
180    }
181}
182impl PartialEq<Immediate> for &Operand {
183    #[inline(always)]
184    fn eq(&self, other: &Immediate) -> bool {
185        self.operand.eq(other)
186    }
187}
188impl PartialEq<Type> for Operand {
189    #[inline(always)]
190    fn eq(&self, other: &Type) -> bool {
191        self.operand.eq(other)
192    }
193}
194impl PartialEq<Type> for &Operand {
195    #[inline(always)]
196    fn eq(&self, other: &Type) -> bool {
197        self.operand.eq(other)
198    }
199}
200impl From<Immediate> for Operand {
201    #[inline]
202    fn from(imm: Immediate) -> Self {
203        Self::new(imm.into())
204    }
205}
206impl From<u32> for Operand {
207    #[inline]
208    fn from(imm: u32) -> Self {
209        Self::new(Immediate::U32(imm).into())
210    }
211}
212impl TryFrom<&Operand> for ValueRef {
213    type Error = ();
214
215    fn try_from(operand: &Operand) -> Result<Self, Self::Error> {
216        match operand.operand {
217            OperandType::Value(value) => Ok(value),
218            _ => Err(()),
219        }
220    }
221}
222#[cfg(test)]
223impl TryFrom<&Operand> for Immediate {
224    type Error = ();
225
226    fn try_from(operand: &Operand) -> Result<Self, Self::Error> {
227        match &operand.operand {
228            OperandType::Const(value) => value.downcast_ref::<Immediate>().copied().ok_or(()),
229            _ => Err(()),
230        }
231    }
232}
233#[cfg(test)]
234impl TryFrom<&Operand> for Type {
235    type Error = ();
236
237    fn try_from(operand: &Operand) -> Result<Self, Self::Error> {
238        match operand.operand {
239            OperandType::Type(ref ty) => Ok(ty.clone()),
240            _ => Err(()),
241        }
242    }
243}
244#[cfg(test)]
245impl TryFrom<Operand> for Type {
246    type Error = ();
247
248    fn try_from(operand: Operand) -> Result<Self, Self::Error> {
249        match operand.operand {
250            OperandType::Type(ty) => Ok(ty),
251            _ => Err(()),
252        }
253    }
254}
255impl From<Type> for Operand {
256    #[inline]
257    fn from(ty: Type) -> Self {
258        Self::new(OperandType::Type(ty))
259    }
260}
261impl From<ValueRef> for Operand {
262    #[inline]
263    fn from(value: ValueRef) -> Self {
264        Self::new(OperandType::Value(value))
265    }
266}
267impl Operand {
268    pub fn new(operand: OperandType) -> Self {
269        let ty = operand.ty();
270        let mut word = ty.to_raw_parts().expect("invalid operand type");
271        assert!(!word.is_empty(), "invalid operand: must be a sized type");
272        assert!(word.len() <= 4, "invalid operand: must be smaller than or equal to a word");
273        if word.len() > 1 {
274            word.reverse();
275        }
276        Self { word, operand }
277    }
278
279    /// Get the size of this operand in field elements
280    pub fn size(&self) -> usize {
281        self.word.len()
282    }
283
284    /// Get the [OperandType] representing the value of this operand
285    #[inline(always)]
286    pub fn value(&self) -> &OperandType {
287        &self.operand
288    }
289
290    /// Get this operand as a [Value]
291    #[inline]
292    pub fn as_value(&self) -> Option<ValueRef> {
293        self.try_into().ok()
294    }
295
296    /// Get the [Type] of this operand
297    #[inline]
298    pub fn ty(&self) -> Type {
299        self.operand.ty()
300    }
301}
302
303/// This structure emulates the state of the VM's operand stack while
304/// generating code from the SSA representation of a function.
305///
306/// In order to emit efficient and correct stack manipulation code, we must be able to
307/// reason about where values are on the operand stack at a given program point. This
308/// structure tracks what SSA values have been pushed on the operand stack, where they are
309/// on the stack relative to the top, and whether a given stack slot aliases multiple
310/// values.
311///
312/// In addition to the state tracked, this structure also has an API that mimics the
313/// stack manipulation instructions we can emit in the code generator, so that as we
314/// emit instructions and modify this structure at the same time, 1:1.
315#[derive(Clone, PartialEq, Eq)]
316pub struct OperandStack {
317    stack: Vec<Operand>,
318}
319impl Default for OperandStack {
320    fn default() -> Self {
321        Self {
322            stack: Vec::with_capacity(16),
323        }
324    }
325}
326impl OperandStack {
327    /// Renames the `n`th operand from the top of the stack to `value`
328    ///
329    /// The type is assumed to remain unchanged
330    pub fn rename(&mut self, n: usize, value: ValueRef) {
331        match &mut self[n].operand {
332            OperandType::Value(prev_value) => {
333                *prev_value = value;
334            }
335            prev => {
336                *prev = OperandType::Value(value);
337            }
338        }
339    }
340
341    /// Searches for the position on the stack containing the operand corresponding to `value`.
342    ///
343    /// NOTE: This function will panic if `value` is not on the stack
344    pub fn find(&self, value: &ValueRef) -> Option<usize> {
345        self.stack.iter().rev().position(|v| v == value)
346    }
347
348    /// Returns true if the operand stack is empty
349    #[allow(unused)]
350    #[inline(always)]
351    pub fn is_empty(&self) -> bool {
352        self.stack.is_empty()
353    }
354
355    /// Returns the number of field elements on the stack
356    #[inline]
357    pub fn raw_len(&self) -> usize {
358        self.stack.iter().map(|operand| operand.size()).sum()
359    }
360
361    /// Returns the index in the actual runtime stack which corresponds to
362    /// the first element of the operand at `index`.
363    #[track_caller]
364    pub fn effective_index(&self, index: usize) -> usize {
365        assert!(
366            index < self.stack.len(),
367            "expected {} to be less than {}",
368            index,
369            self.stack.len()
370        );
371
372        self.stack.iter().rev().take(index).map(|o| o.size()).sum()
373    }
374
375    /// Returns the index in the actual runtime stack which corresponds to
376    /// the last element of the operand at `index`.
377    #[track_caller]
378    pub fn effective_index_inclusive(&self, index: usize) -> usize {
379        assert!(index < self.stack.len());
380
381        self.stack.iter().rev().take(index + 1).map(|o| o.size()).sum::<usize>() - 1
382    }
383
384    /// Returns the number of operands on the stack
385    #[inline]
386    pub fn len(&self) -> usize {
387        self.stack.len()
388    }
389
390    /// Returns the Some(operand) at the given index, without consuming it.
391    /// If the index is out of bounds, returns None.
392    #[inline]
393    pub fn get(&self, index: usize) -> Option<&Operand> {
394        let effective_len: usize = self.stack.iter().rev().take(index + 1).map(|o| o.size()).sum();
395        assert!(
396            effective_len <= 16,
397            "invalid operand stack index ({index}): requires access to more than 16 elements, \
398             which is not supported in Miden"
399        );
400        let len = self.stack.len();
401        if index >= len {
402            return None;
403        }
404        self.stack.get(len - index - 1)
405    }
406
407    /// Returns the operand on top of the stack, without consuming it
408    #[inline]
409    pub fn peek(&self) -> Option<&Operand> {
410        self.stack.last()
411    }
412
413    /// Pushes an operand on top of the stack
414    #[inline]
415    pub fn push<V: Into<Operand>>(&mut self, value: V) {
416        self.stack.push(value.into());
417    }
418
419    /// Pops the operand on top of the stack
420    #[inline]
421    #[track_caller]
422    pub fn pop(&mut self) -> Option<Operand> {
423        self.stack.pop()
424    }
425
426    /// Drops the top operand on the stack
427    #[allow(clippy::should_implement_trait)]
428    #[track_caller]
429    pub fn drop(&mut self) {
430        self.stack.pop().expect("operand stack is empty");
431    }
432
433    /// Drops the top `n` operands on the stack
434    #[inline]
435    #[track_caller]
436    pub fn dropn(&mut self, n: usize) {
437        let len = self.stack.len();
438        assert!(n <= len, "unable to drop {n} operands, operand stack only has {len}");
439        self.stack.truncate(len - n);
440    }
441
442    /// Duplicates the operand in the `n`th position on the stack
443    ///
444    /// If `n` is 0, duplicates the top of the stack.
445    #[track_caller]
446    pub fn dup(&mut self, n: usize) {
447        let operand = self[n].clone();
448        self.stack.push(operand);
449    }
450
451    /// Swaps the `n`th operand from the top of the stack, with the top of the stack
452    ///
453    /// If `n` is 1, it swaps the first two operands on the stack.
454    ///
455    /// NOTE: This function will panic if `n` is 0, or out of bounds.
456    #[track_caller]
457    pub fn swap(&mut self, n: usize) {
458        assert_ne!(n, 0, "invalid swap, index must be in the range 1..=15");
459        let len = self.stack.len();
460        assert!(n < len, "invalid operand stack index ({n}), only {len} operands are available");
461        let a = len - 1;
462        let b = a - n;
463        self.stack.swap(a, b);
464    }
465
466    /// Moves the `n`th operand to the top of the stack
467    ///
468    /// If `n` is 1, this is equivalent to `swap(1)`.
469    ///
470    /// NOTE: This function will panic if `n` is 0, or out of bounds.
471    pub fn movup(&mut self, n: usize) {
472        assert_ne!(n, 0, "invalid move, index must be in the range 1..=15");
473        let len = self.stack.len();
474        assert!(n < len, "invalid operand stack index ({n}), only {len} operands are available");
475        // Pick the midpoint by counting backwards from the end
476        let mid = len - (n + 1);
477        // Split the stack, and rotate the half that
478        // contains our desired value to place it on top.
479        let (_, r) = self.stack.split_at_mut(mid);
480        r.rotate_left(1);
481    }
482
483    /// Makes the operand on top of the stack, the `n`th operand on the stack
484    ///
485    /// If `n` is 1, this is equivalent to `swap(1)`.
486    ///
487    /// NOTE: This function will panic if `n` is 0, or out of bounds.
488    pub fn movdn(&mut self, n: usize) {
489        assert_ne!(n, 0, "invalid move, index must be in the range 1..=15");
490        let len = self.stack.len();
491        assert!(n < len, "invalid operand stack index ({n}), only {len} operands are available");
492        // Split the stack so that the desired position is in the top half
493        let mid = len - (n + 1);
494        let (_, r) = self.stack.split_at_mut(mid);
495        // Move all elements above the `n`th position up by one, moving the top element to the `n`th
496        // position
497        r.rotate_right(1);
498    }
499
500    #[allow(unused)]
501    #[inline(always)]
502    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &Operand> {
503        self.stack.iter()
504    }
505}
506impl Index<usize> for OperandStack {
507    type Output = Operand;
508
509    #[track_caller]
510    fn index(&self, index: usize) -> &Self::Output {
511        let len = self.stack.len();
512        assert!(
513            index < len,
514            "invalid operand stack index ({index}): only {len} operands are available"
515        );
516        let effective_len: usize = self.stack.iter().rev().take(index + 1).map(|o| o.size()).sum();
517        assert!(
518            effective_len <= 16,
519            "invalid operand stack index ({index}): requires access to more than 16 elements, \
520             which is not supported in Miden"
521        );
522        &self.stack[len - index - 1]
523    }
524}
525impl IndexMut<usize> for OperandStack {
526    #[track_caller]
527    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
528        let len = self.stack.len();
529        assert!(
530            index < len,
531            "invalid operand stack index ({index}): only {len} elements are available"
532        );
533        let effective_len: usize = self.stack.iter().rev().take(index + 1).map(|o| o.size()).sum();
534        assert!(
535            effective_len <= 16,
536            "invalid operand stack index ({index}): requires access to more than 16 elements, \
537             which is not supported in Miden"
538        );
539        &mut self.stack[len - index - 1]
540    }
541}
542
543impl fmt::Debug for OperandStack {
544    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
545        let mut builder = f.debug_list();
546        for (index, value) in self.stack.iter().rev().enumerate() {
547            builder.entry_with(|f| write!(f, "{index} => {value:?}"));
548        }
549        builder.finish()
550    }
551}
552
553#[cfg(test)]
554mod tests {
555    use alloc::rc::Rc;
556
557    use midenc_hir::{ArrayType, BuilderExt, Context, PointerType, StructType};
558
559    use super::*;
560
561    #[test]
562    fn operand_stack_homogenous_operand_sizes_test() {
563        let mut stack = OperandStack::default();
564
565        let zero = Immediate::U32(0);
566        let one = Immediate::U32(1);
567        let two = Immediate::U32(2);
568        let three = Immediate::U32(3);
569
570        #[inline]
571        #[allow(unused)]
572        fn as_imms(word: [Operand; 4]) -> [Immediate; 4] {
573            [
574                (&word[0]).try_into().unwrap(),
575                (&word[1]).try_into().unwrap(),
576                (&word[2]).try_into().unwrap(),
577                (&word[3]).try_into().unwrap(),
578            ]
579        }
580
581        #[inline]
582        fn as_imm(operand: Operand) -> Immediate {
583            (&operand).try_into().unwrap()
584        }
585
586        // push
587        stack.push(zero);
588        stack.push(one);
589        stack.push(two);
590        stack.push(three);
591        assert_eq!(stack.len(), 4);
592        assert_eq!(stack[0], three);
593        assert_eq!(stack[1], two);
594        assert_eq!(stack[2], one);
595        assert_eq!(stack[3], zero);
596
597        // peek
598        assert_eq!(stack.peek().unwrap(), three);
599
600        // dup
601        stack.dup(0);
602        assert_eq!(stack.len(), 5);
603        assert_eq!(stack[0], three);
604        assert_eq!(stack[1], three);
605        assert_eq!(stack[2], two);
606        assert_eq!(stack[3], one);
607        assert_eq!(stack[4], zero);
608
609        stack.dup(3);
610        assert_eq!(stack.len(), 6);
611        assert_eq!(stack[0], one);
612        assert_eq!(stack[1], three);
613        assert_eq!(stack[2], three);
614        assert_eq!(stack[3], two);
615        assert_eq!(stack[4], one);
616        assert_eq!(stack[5], zero);
617
618        // drop
619        stack.drop();
620        assert_eq!(stack.len(), 5);
621        assert_eq!(stack[0], three);
622        assert_eq!(stack[1], three);
623        assert_eq!(stack[2], two);
624        assert_eq!(stack[3], one);
625        assert_eq!(stack[4], zero);
626
627        // swap
628        stack.swap(2);
629        assert_eq!(stack.len(), 5);
630        assert_eq!(stack[0], two);
631        assert_eq!(stack[1], three);
632        assert_eq!(stack[2], three);
633        assert_eq!(stack[3], one);
634        assert_eq!(stack[4], zero);
635
636        stack.swap(1);
637        assert_eq!(stack.len(), 5);
638        assert_eq!(stack[0], three);
639        assert_eq!(stack[1], two);
640        assert_eq!(stack[2], three);
641        assert_eq!(stack[3], one);
642        assert_eq!(stack[4], zero);
643
644        // movup
645        stack.movup(2);
646        assert_eq!(stack.len(), 5);
647        assert_eq!(stack[0], three);
648        assert_eq!(stack[1], three);
649        assert_eq!(stack[2], two);
650        assert_eq!(stack[3], one);
651        assert_eq!(stack[4], zero);
652
653        // movdn
654        stack.movdn(3);
655        assert_eq!(stack.len(), 5);
656        assert_eq!(stack[0], three);
657        assert_eq!(stack[1], two);
658        assert_eq!(stack[2], one);
659        assert_eq!(stack[3], three);
660        assert_eq!(stack[4], zero);
661
662        // pop
663        assert_eq!(stack.pop().map(as_imm), Some(three));
664        assert_eq!(stack.len(), 4);
665        assert_eq!(stack[0], two);
666        assert_eq!(stack[1], one);
667        assert_eq!(stack[2], three);
668        assert_eq!(stack[3], zero);
669
670        // dropn
671        stack.dropn(2);
672        assert_eq!(stack.len(), 2);
673        assert_eq!(stack[0], three);
674        assert_eq!(stack[1], zero);
675    }
676
677    #[test]
678    fn operand_stack_values_test() {
679        use midenc_dialect_hir::Load;
680
681        let mut stack = OperandStack::default();
682        let context = Rc::new(Context::default());
683
684        let ptr_u8 = Type::from(PointerType::new(Type::U8));
685        let array_u8 = Type::from(ArrayType::new(Type::U8, 4));
686        let struct_ty = Type::from(StructType::new([Type::U64, Type::U8]));
687        let block = context.create_block_with_params([ptr_u8, array_u8, Type::U32, struct_ty]);
688        let block = block.borrow();
689        let values = block.arguments();
690
691        let zero = values[0] as ValueRef;
692        let one = values[1] as ValueRef;
693        let two = values[2] as ValueRef;
694        let three = values[3] as ValueRef;
695
696        drop(block);
697
698        // push
699        stack.push(zero);
700        stack.push(one);
701        stack.push(two);
702        stack.push(three);
703        assert_eq!(stack.len(), 4);
704        assert_eq!(stack.raw_len(), 6);
705
706        assert_eq!(stack.find(&zero), Some(3));
707        assert_eq!(stack.find(&one), Some(2));
708        assert_eq!(stack.find(&two), Some(1));
709        assert_eq!(stack.find(&three), Some(0));
710
711        // dup
712        stack.dup(0);
713        assert_eq!(stack.find(&three), Some(0));
714
715        stack.dup(3);
716        assert_eq!(stack.find(&one), Some(0));
717
718        // drop
719        stack.drop();
720        assert_eq!(stack.find(&one), Some(3));
721        assert_eq!(stack.find(&three), Some(0));
722        assert_eq!(stack[1], three);
723
724        // padw
725        stack.push(Immediate::Felt(Felt::ZERO));
726        stack.push(Immediate::Felt(Felt::ZERO));
727        stack.push(Immediate::Felt(Felt::ZERO));
728        stack.push(Immediate::Felt(Felt::ZERO));
729        assert_eq!(stack.find(&one), Some(7));
730        assert_eq!(stack.find(&three), Some(4));
731
732        // rename
733        let four = {
734            let mut builder = midenc_hir::OpBuilder::new(context.clone());
735            let load_builder = builder.create::<Load, _>(Default::default());
736            let load = load_builder(zero).unwrap();
737            load.borrow().result().as_value_ref()
738        };
739        stack.rename(1, four);
740        assert_eq!(stack.find(&four), Some(1));
741        assert_eq!(stack.find(&three), Some(4));
742
743        // pop
744        let top = stack.pop().unwrap();
745        assert_eq!((&top).try_into(), Ok(Immediate::Felt(Felt::ZERO)));
746        assert_eq!(stack.find(&four), Some(0));
747        assert_eq!(stack[1], Immediate::Felt(Felt::ZERO));
748        assert_eq!(stack[2], Immediate::Felt(Felt::ZERO));
749        assert_eq!(stack.find(&three), Some(3));
750
751        // dropn
752        stack.dropn(3);
753        assert_eq!(stack.find(&four), None);
754        assert_eq!(stack.find(&three), Some(0));
755        assert_eq!(stack[1], three);
756        assert_eq!(stack.find(&two), Some(2));
757        assert_eq!(stack.find(&one), Some(3));
758        assert_eq!(stack.find(&zero), Some(4));
759
760        // swap
761        stack.swap(3);
762        assert_eq!(stack.find(&one), Some(0));
763        assert_eq!(stack.find(&three), Some(1));
764        assert_eq!(stack.find(&two), Some(2));
765        assert_eq!(stack[3], three);
766
767        stack.swap(1);
768        assert_eq!(stack.find(&three), Some(0));
769        assert_eq!(stack.find(&one), Some(1));
770        assert_eq!(stack.find(&two), Some(2));
771        assert_eq!(stack.find(&zero), Some(4));
772
773        // movup
774        stack.movup(2);
775        assert_eq!(stack.find(&two), Some(0));
776        assert_eq!(stack.find(&three), Some(1));
777        assert_eq!(stack.find(&one), Some(2));
778        assert_eq!(stack.find(&zero), Some(4));
779
780        // movdn
781        stack.movdn(3);
782        assert_eq!(stack.find(&three), Some(0));
783        assert_eq!(stack.find(&one), Some(1));
784        assert_eq!(stack[2], three);
785        assert_eq!(stack.find(&two), Some(3));
786        assert_eq!(stack.find(&zero), Some(4));
787    }
788
789    #[test]
790    fn operand_stack_heterogenous_operand_sizes_test() {
791        let mut stack = OperandStack::default();
792
793        let zero = Immediate::U32(0);
794        let one = Immediate::U32(1);
795        let two = Type::U64;
796        let three = Type::U64;
797        let struct_a = Type::from(StructType::new([
798            Type::from(PointerType::new(Type::U8)),
799            Type::U16,
800            Type::U32,
801        ]));
802
803        // push
804        stack.push(zero);
805        stack.push(one);
806        stack.push(two.clone());
807        stack.push(three.clone());
808        stack.push(struct_a.clone());
809        assert_eq!(stack.len(), 5);
810        assert_eq!(stack.raw_len(), 9);
811        assert_eq!(stack[0], struct_a);
812        assert_eq!(stack[1], three);
813        assert_eq!(stack[2], two);
814        assert_eq!(stack[3], one);
815        assert_eq!(stack[4], zero);
816
817        // peek
818        assert_eq!(stack.peek().unwrap(), struct_a);
819
820        // dup
821        stack.dup(0);
822        assert_eq!(stack.len(), 6);
823        assert_eq!(stack.raw_len(), 12);
824        assert_eq!(stack[0], struct_a);
825        assert_eq!(stack[1], struct_a);
826        assert_eq!(stack[2], three);
827        assert_eq!(stack[3], two);
828        assert_eq!(stack[4], one);
829        assert_eq!(stack[5], zero);
830        assert_eq!(stack.effective_index(3), 8);
831
832        stack.dup(3);
833        assert_eq!(stack.len(), 7);
834        assert_eq!(stack.raw_len(), 14);
835        assert_eq!(stack[0], two);
836        assert_eq!(stack[1], struct_a);
837        assert_eq!(stack[2], struct_a);
838
839        // drop
840        stack.drop();
841        assert_eq!(stack.len(), 6);
842        assert_eq!(stack.raw_len(), 12);
843        assert_eq!(stack[0], struct_a);
844        assert_eq!(stack[1], struct_a);
845        assert_eq!(stack[2], three);
846        assert_eq!(stack[3], two);
847        assert_eq!(stack[4], one);
848        assert_eq!(stack[5], zero);
849
850        // swap
851        stack.swap(2);
852        assert_eq!(stack.len(), 6);
853        assert_eq!(stack.raw_len(), 12);
854        assert_eq!(stack[0], three);
855        assert_eq!(stack[1], struct_a);
856        assert_eq!(stack[2], struct_a);
857        assert_eq!(stack[3], two);
858        assert_eq!(stack[4], one);
859
860        stack.swap(1);
861        assert_eq!(stack.len(), 6);
862        assert_eq!(stack.raw_len(), 12);
863        assert_eq!(stack[0], struct_a);
864        assert_eq!(stack[1], three);
865        assert_eq!(stack[2], struct_a);
866        assert_eq!(stack[3], two);
867        assert_eq!(stack[4], one);
868        assert_eq!(stack[5], zero);
869
870        // movup
871        stack.movup(4);
872        assert_eq!(stack.len(), 6);
873        assert_eq!(stack.raw_len(), 12);
874        assert_eq!(stack[0], one);
875        assert_eq!(stack[1], struct_a);
876        assert_eq!(stack[2], three);
877        assert_eq!(stack[3], struct_a);
878        assert_eq!(stack[4], two);
879        assert_eq!(stack[5], zero);
880
881        // movdn
882        stack.movdn(3);
883        assert_eq!(stack.len(), 6);
884        assert_eq!(stack.raw_len(), 12);
885        assert_eq!(stack[0], struct_a);
886        assert_eq!(stack[1], three);
887        assert_eq!(stack[2], struct_a);
888        assert_eq!(stack[3], one);
889        assert_eq!(stack[4], two);
890
891        // pop
892        let operand: Type = stack.pop().unwrap().try_into().unwrap();
893        assert_eq!(operand, struct_a);
894        assert_eq!(stack.len(), 5);
895        assert_eq!(stack.raw_len(), 9);
896        assert_eq!(stack[0], three);
897        assert_eq!(stack[1], struct_a);
898        assert_eq!(stack[2], one);
899        assert_eq!(stack[3], two);
900
901        // dropn
902        stack.dropn(2);
903        assert_eq!(stack.len(), 3);
904        assert_eq!(stack.raw_len(), 4);
905        assert_eq!(stack[0], one);
906        assert_eq!(stack[1], two);
907        assert_eq!(stack[2], zero);
908    }
909}