xref: /wasmtime-44.0.1/winch/codegen/src/stack.rs (revision 14b39bc2)
1 use crate::{isa::reg::Reg, masm::StackSlot};
2 use std::collections::VecDeque;
3 use wasmparser::{Ieee32, Ieee64};
4 use wasmtime_environ::WasmType;
5 
6 /// A typed register value used to track register values in the value
7 /// stack.
8 #[derive(Debug, Eq, PartialEq, Copy, Clone)]
9 pub struct TypedReg {
10     /// The physical register.
11     pub reg: Reg,
12     /// The type associated to the physical register.
13     pub ty: WasmType,
14 }
15 
16 impl TypedReg {
17     /// Create a new TypedReg.
18     pub fn new(ty: WasmType, reg: Reg) -> Self {
19         Self { ty, reg }
20     }
21 
22     /// Create an i64 TypedReg.
23     pub fn i64(reg: Reg) -> Self {
24         Self {
25             ty: WasmType::I64,
26             reg,
27         }
28     }
29 }
30 
31 impl From<TypedReg> for Reg {
32     fn from(tr: TypedReg) -> Self {
33         tr.reg
34     }
35 }
36 
37 /// A local value.
38 #[derive(Debug, Eq, PartialEq, Copy, Clone)]
39 pub struct Local {
40     /// The index of the local.
41     pub index: u32,
42     /// The type of the local.
43     pub ty: WasmType,
44 }
45 
46 /// A memory value.
47 #[derive(Debug, Eq, PartialEq, Copy, Clone)]
48 pub struct Memory {
49     /// The type associated with the memory offset.
50     pub ty: WasmType,
51     /// The stack slot corresponding to the memory value.
52     pub slot: StackSlot,
53 }
54 
55 /// Value definition to be used within the shadow stack.
56 #[derive(Debug, Eq, PartialEq, Copy, Clone)]
57 pub(crate) enum Val {
58     /// I32 Constant.
59     I32(i32),
60     /// I64 Constant.
61     I64(i64),
62     /// F32 Constant.
63     F32(Ieee32),
64     /// F64 Constant.
65     F64(Ieee64),
66     /// A register value.
67     Reg(TypedReg),
68     /// A local slot.
69     Local(Local),
70     /// Offset to a memory location.
71     Memory(Memory),
72 }
73 
74 impl From<TypedReg> for Val {
75     fn from(tr: TypedReg) -> Self {
76         Val::Reg(tr)
77     }
78 }
79 
80 impl From<Local> for Val {
81     fn from(local: Local) -> Self {
82         Val::Local(local)
83     }
84 }
85 
86 impl From<Memory> for Val {
87     fn from(mem: Memory) -> Self {
88         Val::Memory(mem)
89     }
90 }
91 
92 impl Val {
93     /// Create a new I32 constant value.
94     pub fn i32(v: i32) -> Self {
95         Self::I32(v)
96     }
97 
98     /// Create a new I64 constant value.
99     pub fn i64(v: i64) -> Self {
100         Self::I64(v)
101     }
102 
103     /// Create a new F32 constant value.
104     pub fn f32(v: Ieee32) -> Self {
105         Self::F32(v)
106     }
107 
108     pub fn f64(v: Ieee64) -> Self {
109         Self::F64(v)
110     }
111 
112     /// Create a new Reg value.
113     pub fn reg(reg: Reg, ty: WasmType) -> Self {
114         Self::Reg(TypedReg { reg, ty })
115     }
116 
117     /// Create a new Local value.
118     pub fn local(index: u32, ty: WasmType) -> Self {
119         Self::Local(Local { index, ty })
120     }
121 
122     /// Create a Memory value.
123     pub fn mem(ty: WasmType, slot: StackSlot) -> Self {
124         Self::Memory(Memory { ty, slot })
125     }
126 
127     /// Check whether the value is a register.
128     pub fn is_reg(&self) -> bool {
129         match *self {
130             Self::Reg(_) => true,
131             _ => false,
132         }
133     }
134 
135     /// Check wheter the value is a memory offset.
136     pub fn is_mem(&self) -> bool {
137         match *self {
138             Self::Memory(_) => true,
139             _ => false,
140         }
141     }
142 
143     /// Get the register representation of the value.
144     ///
145     /// # Panics
146     /// This method will panic if the value is not a register.
147     pub fn get_reg(&self) -> TypedReg {
148         match self {
149             Self::Reg(tr) => *tr,
150             v => panic!("expected value {:?} to be a register", v),
151         }
152     }
153 
154     /// Get the integer representation of the value.
155     ///
156     /// # Panics
157     /// This method will panic if the value is not an i32.
158     pub fn get_i32(&self) -> i32 {
159         match self {
160             Self::I32(v) => *v,
161             v => panic!("expected value {:?} to be i32", v),
162         }
163     }
164 
165     /// Get the integer representation of the value.
166     ///
167     /// # Panics
168     /// This method will panic if the value is not an i64.
169     pub fn get_i64(&self) -> i64 {
170         match self {
171             Self::I64(v) => *v,
172             v => panic!("expected value {:?} to be i64", v),
173         }
174     }
175 
176     /// Check whether the value is an i32 constant.
177     pub fn is_i32_const(&self) -> bool {
178         match *self {
179             Self::I32(_) => true,
180             _ => false,
181         }
182     }
183 
184     /// Check whether the value is an i64 constant.
185     pub fn is_i64_const(&self) -> bool {
186         match *self {
187             Self::I64(_) => true,
188             _ => false,
189         }
190     }
191 
192     /// Get the type of the value.
193     pub fn ty(&self) -> WasmType {
194         match self {
195             Val::I32(_) => WasmType::I32,
196             Val::I64(_) => WasmType::I64,
197             Val::F32(_) => WasmType::F32,
198             Val::F64(_) => WasmType::F64,
199             Val::Reg(r) => r.ty,
200             Val::Memory(m) => m.ty,
201             Val::Local(l) => l.ty,
202         }
203     }
204 }
205 
206 /// The shadow stack used for compilation.
207 #[derive(Default, Debug)]
208 pub(crate) struct Stack {
209     inner: VecDeque<Val>,
210 }
211 
212 impl Stack {
213     /// Allocate a new stack.
214     pub fn new() -> Self {
215         Self {
216             inner: Default::default(),
217         }
218     }
219 
220     /// Insert a new value at the specified index.
221     pub fn insert(&mut self, at: usize, val: Val) {
222         self.inner.insert(at, val);
223     }
224 
225     /// Get the length of the stack.
226     pub fn len(&self) -> usize {
227         self.inner.len()
228     }
229 
230     /// Push a value to the stack.
231     pub fn push(&mut self, val: Val) {
232         self.inner.push_back(val);
233     }
234 
235     /// Peek into the top in the stack.
236     pub fn peek(&self) -> Option<&Val> {
237         self.inner.back()
238     }
239 
240     /// Returns an iterator referencing the last n items of the stack,
241     /// in bottom-most to top-most order.
242     pub fn peekn(&self, n: usize) -> impl Iterator<Item = &Val> + '_ {
243         let len = self.len();
244         assert!(n <= len);
245 
246         let partition = len - n;
247         self.inner.range(partition..)
248     }
249 
250     /// Pops the top element of the stack, if any.
251     pub fn pop(&mut self) -> Option<Val> {
252         self.inner.pop_back()
253     }
254 
255     /// Pops the element at the top of the stack if it is an i32 const;
256     /// returns `None` otherwise.
257     pub fn pop_i32_const(&mut self) -> Option<i32> {
258         match self.peek() {
259             Some(v) => v.is_i32_const().then(|| self.pop().unwrap().get_i32()),
260             _ => None,
261         }
262     }
263 
264     /// Pops the element at the top of the stack if it is an i64 const;
265     /// returns `None` otherwise.
266     pub fn pop_i64_const(&mut self) -> Option<i64> {
267         match self.peek() {
268             Some(v) => v.is_i64_const().then(|| self.pop().unwrap().get_i64()),
269             _ => None,
270         }
271     }
272 
273     /// Pops the element at the top of the stack if it is a register;
274     /// returns `None` otherwise.
275     pub fn pop_reg(&mut self) -> Option<TypedReg> {
276         match self.peek() {
277             Some(v) => v.is_reg().then(|| self.pop().unwrap().get_reg()),
278             _ => None,
279         }
280     }
281 
282     /// Pops the given register if it is at the top of the stack;
283     /// returns `None` otherwise.
284     pub fn pop_named_reg(&mut self, reg: Reg) -> Option<TypedReg> {
285         match self.peek() {
286             Some(v) => {
287                 (v.is_reg() && v.get_reg().reg == reg).then(|| self.pop().unwrap().get_reg())
288             }
289             _ => None,
290         }
291     }
292 
293     /// Get a mutable reference to the inner stack representation.
294     pub fn inner_mut(&mut self) -> &mut VecDeque<Val> {
295         &mut self.inner
296     }
297 }
298 
299 #[cfg(test)]
300 mod tests {
301     use super::{Stack, Val};
302     use crate::isa::reg::Reg;
303     use wasmtime_environ::WasmType;
304 
305     #[test]
306     fn test_pop_i32_const() {
307         let mut stack = Stack::new();
308         stack.push(Val::i32(33i32));
309         assert_eq!(33, stack.pop_i32_const().unwrap());
310 
311         stack.push(Val::local(10, WasmType::I32));
312         assert!(stack.pop_i32_const().is_none());
313     }
314 
315     #[test]
316     fn test_pop_reg() {
317         let mut stack = Stack::new();
318         let reg = Reg::int(2usize);
319         stack.push(Val::reg(reg, WasmType::I32));
320         stack.push(Val::i32(4));
321 
322         assert_eq!(None, stack.pop_reg());
323         let _ = stack.pop().unwrap();
324         assert_eq!(reg, stack.pop_reg().unwrap().reg);
325     }
326 
327     #[test]
328     fn test_pop_named_reg() {
329         let mut stack = Stack::new();
330         let reg = Reg::int(2usize);
331         stack.push(Val::reg(reg, WasmType::I32));
332         stack.push(Val::reg(Reg::int(4), WasmType::I32));
333 
334         assert_eq!(None, stack.pop_named_reg(reg));
335         let _ = stack.pop().unwrap();
336         assert_eq!(reg, stack.pop_named_reg(reg).unwrap().reg);
337     }
338 }
339