1 use crate::{isa::reg::Reg, masm::StackSlot}; 2 use smallvec::SmallVec; 3 use wasmparser::{Ieee32, Ieee64}; 4 use wasmtime_environ::WasmValType; 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: WasmValType, 14 } 15 16 impl TypedReg { 17 /// Create a new [`TypedReg`]. 18 pub fn new(ty: WasmValType, 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: WasmValType::I64, 26 reg, 27 } 28 } 29 30 /// Create an i32 [`TypedReg`]. 31 pub fn i32(reg: Reg) -> Self { 32 Self { 33 ty: WasmValType::I32, 34 reg, 35 } 36 } 37 38 /// Create an f64 [`TypedReg`]. 39 pub fn f64(reg: Reg) -> Self { 40 Self { 41 ty: WasmValType::F64, 42 reg, 43 } 44 } 45 46 /// Create an f32 [`TypedReg`]. 47 pub fn f32(reg: Reg) -> Self { 48 Self { 49 ty: WasmValType::F32, 50 reg, 51 } 52 } 53 } 54 55 impl From<TypedReg> for Reg { 56 fn from(tr: TypedReg) -> Self { 57 tr.reg 58 } 59 } 60 61 /// A local value. 62 #[derive(Debug, Eq, PartialEq, Copy, Clone)] 63 pub struct Local { 64 /// The index of the local. 65 pub index: u32, 66 /// The type of the local. 67 pub ty: WasmValType, 68 } 69 70 /// A memory value. 71 #[derive(Debug, Eq, PartialEq, Copy, Clone)] 72 pub struct Memory { 73 /// The type associated with the memory offset. 74 pub ty: WasmValType, 75 /// The stack slot corresponding to the memory value. 76 pub slot: StackSlot, 77 } 78 79 /// Value definition to be used within the shadow stack. 80 #[derive(Debug, Eq, PartialEq, Copy, Clone)] 81 pub(crate) enum Val { 82 /// I32 Constant. 83 I32(i32), 84 /// I64 Constant. 85 I64(i64), 86 /// F32 Constant. 87 F32(Ieee32), 88 /// F64 Constant. 89 F64(Ieee64), 90 /// V128 Constant. 91 V128(i128), 92 /// A register value. 93 Reg(TypedReg), 94 /// A local slot. 95 Local(Local), 96 /// Offset to a memory location. 97 Memory(Memory), 98 } 99 100 impl From<TypedReg> for Val { 101 fn from(tr: TypedReg) -> Self { 102 Val::Reg(tr) 103 } 104 } 105 106 impl From<Local> for Val { 107 fn from(local: Local) -> Self { 108 Val::Local(local) 109 } 110 } 111 112 impl From<Memory> for Val { 113 fn from(mem: Memory) -> Self { 114 Val::Memory(mem) 115 } 116 } 117 118 impl TryFrom<u32> for Val { 119 type Error = anyhow::Error; 120 fn try_from(value: u32) -> Result<Self, Self::Error> { 121 i32::try_from(value).map(Val::i32).map_err(Into::into) 122 } 123 } 124 125 impl Val { 126 /// Create a new I32 constant value. 127 pub fn i32(v: i32) -> Self { 128 Self::I32(v) 129 } 130 131 /// Create a new I64 constant value. 132 pub fn i64(v: i64) -> Self { 133 Self::I64(v) 134 } 135 136 /// Create a new F32 constant value. 137 pub fn f32(v: Ieee32) -> Self { 138 Self::F32(v) 139 } 140 141 pub fn f64(v: Ieee64) -> Self { 142 Self::F64(v) 143 } 144 145 /// Create a new V128 constant value. 146 pub fn v128(v: i128) -> Self { 147 Self::V128(v) 148 } 149 150 /// Create a new Reg value. 151 pub fn reg(reg: Reg, ty: WasmValType) -> Self { 152 Self::Reg(TypedReg { reg, ty }) 153 } 154 155 /// Create a new Local value. 156 pub fn local(index: u32, ty: WasmValType) -> Self { 157 Self::Local(Local { index, ty }) 158 } 159 160 /// Create a Memory value. 161 pub fn mem(ty: WasmValType, slot: StackSlot) -> Self { 162 Self::Memory(Memory { ty, slot }) 163 } 164 165 /// Check whether the value is a register. 166 pub fn is_reg(&self) -> bool { 167 match *self { 168 Self::Reg(_) => true, 169 _ => false, 170 } 171 } 172 173 /// Check whether the value is a memory offset. 174 pub fn is_mem(&self) -> bool { 175 match *self { 176 Self::Memory(_) => true, 177 _ => false, 178 } 179 } 180 181 /// Check whether the value is a constant. 182 pub fn is_const(&self) -> bool { 183 match *self { 184 Val::I32(_) | Val::I64(_) | Val::F32(_) | Val::F64(_) | Val::V128(_) => true, 185 _ => false, 186 } 187 } 188 189 /// Check whether the value is local with a particular index. 190 pub fn is_local_at_index(&self, index: u32) -> bool { 191 match *self { 192 Self::Local(Local { index: i, .. }) if i == index => true, 193 _ => false, 194 } 195 } 196 197 /// Get the register representation of the value. 198 /// 199 /// # Panics 200 /// This method will panic if the value is not a register. 201 pub fn unwrap_reg(&self) -> TypedReg { 202 match self { 203 Self::Reg(tr) => *tr, 204 v => panic!("expected value {:?} to be a register", v), 205 } 206 } 207 208 /// Get the integer representation of the value. 209 /// 210 /// # Panics 211 /// This method will panic if the value is not an i32. 212 pub fn unwrap_i32(&self) -> i32 { 213 match self { 214 Self::I32(v) => *v, 215 v => panic!("expected value {:?} to be i32", v), 216 } 217 } 218 219 /// Get the integer representation of the value. 220 /// 221 /// # Panics 222 /// This method will panic if the value is not an i64. 223 pub fn unwrap_i64(&self) -> i64 { 224 match self { 225 Self::I64(v) => *v, 226 v => panic!("expected value {:?} to be i64", v), 227 } 228 } 229 230 /// Returns the underlying memory value if it is one, panics otherwise. 231 pub fn unwrap_mem(&self) -> Memory { 232 match self { 233 Self::Memory(m) => *m, 234 v => panic!("expected value {:?} to be a Memory", v), 235 } 236 } 237 238 /// Check whether the value is an i32 constant. 239 pub fn is_i32_const(&self) -> bool { 240 match *self { 241 Self::I32(_) => true, 242 _ => false, 243 } 244 } 245 246 /// Check whether the value is an i64 constant. 247 pub fn is_i64_const(&self) -> bool { 248 match *self { 249 Self::I64(_) => true, 250 _ => false, 251 } 252 } 253 254 /// Get the type of the value. 255 pub fn ty(&self) -> WasmValType { 256 match self { 257 Val::I32(_) => WasmValType::I32, 258 Val::I64(_) => WasmValType::I64, 259 Val::F32(_) => WasmValType::F32, 260 Val::F64(_) => WasmValType::F64, 261 Val::V128(_) => WasmValType::V128, 262 Val::Reg(r) => r.ty, 263 Val::Memory(m) => m.ty, 264 Val::Local(l) => l.ty, 265 } 266 } 267 } 268 269 /// The shadow stack used for compilation. 270 #[derive(Default, Debug)] 271 pub(crate) struct Stack { 272 // NB: The 64 is chosen arbitrarily. We can adjust as we see fit. 273 inner: SmallVec<[Val; 64]>, 274 } 275 276 impl Stack { 277 /// Allocate a new stack. 278 pub fn new() -> Self { 279 Self { 280 inner: Default::default(), 281 } 282 } 283 284 /// Returns true if the stack contains a local with the provided index 285 /// except if the only time the local appears is the top element. 286 pub fn contains_latent_local(&self, index: u32) -> bool { 287 self.inner 288 .iter() 289 // Iterate top-to-bottom so we can skip the top element and stop 290 // when we see a memory element. 291 .rev() 292 // The local is not latent if it's the top element because the top 293 // element will be popped next which materializes the local. 294 .skip(1) 295 // Stop when we see a memory element because that marks where we 296 // spilled up to so there will not be any locals past this point. 297 .take_while(|v| !v.is_mem()) 298 .any(|v| v.is_local_at_index(index)) 299 } 300 301 /// Extend the stack with the given elements. 302 pub fn extend(&mut self, values: impl IntoIterator<Item = Val>) { 303 self.inner.extend(values); 304 } 305 306 /// Inserts many values at the given index. 307 pub fn insert_many(&mut self, at: usize, values: &[Val]) { 308 debug_assert!(at <= self.len()); 309 310 if at == self.len() { 311 self.inner.extend_from_slice(values); 312 } else { 313 self.inner.insert_from_slice(at, values); 314 } 315 } 316 317 /// Get the length of the stack. 318 pub fn len(&self) -> usize { 319 self.inner.len() 320 } 321 322 /// Push a value to the stack. 323 pub fn push(&mut self, val: Val) { 324 self.inner.push(val); 325 } 326 327 /// Peek into the top in the stack. 328 pub fn peek(&self) -> Option<&Val> { 329 self.inner.last() 330 } 331 332 /// Returns an iterator referencing the last n items of the stack, 333 /// in bottom-most to top-most order. 334 pub fn peekn(&self, n: usize) -> impl Iterator<Item = &Val> + '_ { 335 let len = self.len(); 336 assert!(n <= len); 337 338 let partition = len - n; 339 self.inner[partition..].into_iter() 340 } 341 342 /// Pops the top element of the stack, if any. 343 pub fn pop(&mut self) -> Option<Val> { 344 self.inner.pop() 345 } 346 347 /// Pops the element at the top of the stack if it is an i32 const; 348 /// returns `None` otherwise. 349 pub fn pop_i32_const(&mut self) -> Option<i32> { 350 match self.peek() { 351 Some(v) => v.is_i32_const().then(|| self.pop().unwrap().unwrap_i32()), 352 _ => None, 353 } 354 } 355 356 /// Pops the element at the top of the stack if it is an i64 const; 357 /// returns `None` otherwise. 358 pub fn pop_i64_const(&mut self) -> Option<i64> { 359 match self.peek() { 360 Some(v) => v.is_i64_const().then(|| self.pop().unwrap().unwrap_i64()), 361 _ => None, 362 } 363 } 364 365 /// Pops the element at the top of the stack if it is a register; 366 /// returns `None` otherwise. 367 pub fn pop_reg(&mut self) -> Option<TypedReg> { 368 match self.peek() { 369 Some(v) => v.is_reg().then(|| self.pop().unwrap().unwrap_reg()), 370 _ => None, 371 } 372 } 373 374 /// Pops the given register if it is at the top of the stack; 375 /// returns `None` otherwise. 376 pub fn pop_named_reg(&mut self, reg: Reg) -> Option<TypedReg> { 377 match self.peek() { 378 Some(v) => { 379 (v.is_reg() && v.unwrap_reg().reg == reg).then(|| self.pop().unwrap().unwrap_reg()) 380 } 381 _ => None, 382 } 383 } 384 385 /// Get a mutable reference to the inner stack representation. 386 pub fn inner_mut(&mut self) -> &mut SmallVec<[Val; 64]> { 387 &mut self.inner 388 } 389 390 /// Get a reference to the inner stack representation. 391 pub fn inner(&self) -> &SmallVec<[Val; 64]> { 392 &self.inner 393 } 394 395 /// Calculates the size of, in bytes, of the top n [Memory] entries 396 /// in the value stack. 397 pub fn sizeof(&self, top: usize) -> u32 { 398 self.peekn(top).fold(0, |acc, v| { 399 if v.is_mem() { 400 acc + v.unwrap_mem().slot.size 401 } else { 402 acc 403 } 404 }) 405 } 406 } 407 408 #[cfg(test)] 409 mod tests { 410 use super::{Stack, Val}; 411 use crate::isa::reg::Reg; 412 use wasmtime_environ::WasmValType; 413 414 #[test] 415 fn test_pop_i32_const() { 416 let mut stack = Stack::new(); 417 stack.push(Val::i32(33i32)); 418 assert_eq!(33, stack.pop_i32_const().unwrap()); 419 420 stack.push(Val::local(10, WasmValType::I32)); 421 assert!(stack.pop_i32_const().is_none()); 422 } 423 424 #[test] 425 fn test_pop_reg() { 426 let mut stack = Stack::new(); 427 let reg = Reg::int(2usize); 428 stack.push(Val::reg(reg, WasmValType::I32)); 429 stack.push(Val::i32(4)); 430 431 assert_eq!(None, stack.pop_reg()); 432 let _ = stack.pop().unwrap(); 433 assert_eq!(reg, stack.pop_reg().unwrap().reg); 434 } 435 436 #[test] 437 fn test_pop_named_reg() { 438 let mut stack = Stack::new(); 439 let reg = Reg::int(2usize); 440 stack.push(Val::reg(reg, WasmValType::I32)); 441 stack.push(Val::reg(Reg::int(4), WasmValType::I32)); 442 443 assert_eq!(None, stack.pop_named_reg(reg)); 444 let _ = stack.pop().unwrap(); 445 assert_eq!(reg, stack.pop_named_reg(reg).unwrap().reg); 446 } 447 } 448