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