Skip to main content

midenc_codegen_masm/
stack.rs

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