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