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