xref: /wasmtime-44.0.1/winch/codegen/src/stack.rs (revision 2329ecc3)
1 use crate::isa::reg::Reg;
2 use std::collections::VecDeque;
3 
4 /// Value definition to be used within the shadow stack.
5 #[derive(Debug, Eq, PartialEq)]
6 pub(crate) enum Val {
7     /// I32 Constant.
8     I32(i32),
9     /// A register.
10     Reg(Reg),
11     /// A local slot.
12     Local(u32),
13     /// Offset to a memory location.
14     Memory(u32),
15 }
16 
17 impl Val {
18     /// Create a new I32 constant value.
19     pub fn i32(v: i32) -> Self {
20         Self::I32(v)
21     }
22 
23     /// Create a new Reg value.
24     pub fn reg(r: Reg) -> Self {
25         Self::Reg(r)
26     }
27 
28     /// Create a new Local value.
29     pub fn local(index: u32) -> Self {
30         Self::Local(index)
31     }
32 
33     /// Check whether the value is a register.
34     pub fn is_reg(&self) -> bool {
35         match *self {
36             Self::Reg(_) => true,
37             _ => false,
38         }
39     }
40 
41     /// Get the register representation of the value.
42     ///
43     /// # Panics
44     /// This method will panic if the value is not a register.
45     pub fn get_reg(&self) -> Reg {
46         match self {
47             Self::Reg(r) => *r,
48             v => panic!("expected value {:?} to be a register", v),
49         }
50     }
51 
52     /// Get the integer representation of the value.
53     ///
54     /// # Panics
55     /// This method will panic if the value is not an i32.
56     pub fn get_i32(&self) -> i32 {
57         match self {
58             Self::I32(v) => *v,
59             v => panic!("expected value {:?} to be i32", v),
60         }
61     }
62 
63     /// Check whether the value is an i32 constant.
64     pub fn is_i32_const(&self) -> bool {
65         match *self {
66             Self::I32(_) => true,
67             _ => false,
68         }
69     }
70 }
71 
72 /// The shadow stack used for compilation.
73 #[derive(Default, Debug)]
74 pub(crate) struct Stack {
75     inner: VecDeque<Val>,
76 }
77 
78 impl Stack {
79     /// Allocate a new stack.
80     pub fn new() -> Self {
81         Self {
82             inner: Default::default(),
83         }
84     }
85 
86     /// Push a value to the stack.
87     pub fn push(&mut self, val: Val) {
88         self.inner.push_back(val);
89     }
90 
91     /// Peek into the top in the stack.
92     pub fn peek(&mut self) -> Option<&Val> {
93         self.inner.back()
94     }
95 
96     /// Pops the top element of the stack, if any.
97     pub fn pop(&mut self) -> Option<Val> {
98         self.inner.pop_back()
99     }
100 
101     /// Pops the element at the top of the stack if it is a const;
102     /// returns `None` otherwise.
103     pub fn pop_i32_const(&mut self) -> Option<i32> {
104         match self.peek() {
105             Some(v) => v.is_i32_const().then(|| self.pop().unwrap().get_i32()),
106             _ => None,
107         }
108     }
109 
110     /// Pops the element at the top of the stack if it is a register;
111     /// returns `None` otherwise.
112     pub fn pop_reg(&mut self) -> Option<Reg> {
113         match self.peek() {
114             Some(v) => v.is_reg().then(|| self.pop().unwrap().get_reg()),
115             _ => None,
116         }
117     }
118 
119     /// Pops the given register if it is at the top of the stack;
120     /// returns `None` otherwise.
121     pub fn pop_named_reg(&mut self, reg: Reg) -> Option<Reg> {
122         match self.peek() {
123             Some(v) => (v.is_reg() && v.get_reg() == reg).then(|| self.pop().unwrap().get_reg()),
124             _ => None,
125         }
126     }
127 
128     /// Get a mutable reference to the inner stack representation.
129     pub fn inner_mut(&mut self) -> &mut VecDeque<Val> {
130         &mut self.inner
131     }
132 }
133 
134 #[cfg(test)]
135 mod tests {
136     use super::{Stack, Val};
137     use crate::isa::reg::Reg;
138 
139     #[test]
140     fn test_pop_i32_const() {
141         let mut stack = Stack::new();
142         stack.push(Val::i32(33i32));
143         assert_eq!(33, stack.pop_i32_const().unwrap());
144 
145         stack.push(Val::local(10));
146         assert!(stack.pop_i32_const().is_none());
147     }
148 
149     #[test]
150     fn test_pop_reg() {
151         let mut stack = Stack::new();
152         let reg = Reg::int(2usize);
153         stack.push(Val::reg(reg));
154         stack.push(Val::i32(4));
155 
156         assert_eq!(None, stack.pop_reg());
157         let _ = stack.pop().unwrap();
158         assert_eq!(reg, stack.pop_reg().unwrap());
159     }
160 
161     #[test]
162     fn test_pop_named_reg() {
163         let mut stack = Stack::new();
164         let reg = Reg::int(2usize);
165         stack.push(Val::reg(reg));
166         stack.push(Val::reg(Reg::int(4)));
167 
168         assert_eq!(None, stack.pop_named_reg(reg));
169         let _ = stack.pop().unwrap();
170         assert_eq!(reg, stack.pop_named_reg(reg).unwrap());
171     }
172 }
173