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