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 /// 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 get_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 get_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 get_i64(&self) -> i64 { 193 match self { 194 Self::I64(v) => *v, 195 v => panic!("expected value {:?} to be i64", v), 196 } 197 } 198 199 /// Check whether the value is an i32 constant. 200 pub fn is_i32_const(&self) -> bool { 201 match *self { 202 Self::I32(_) => true, 203 _ => false, 204 } 205 } 206 207 /// Check whether the value is an i64 constant. 208 pub fn is_i64_const(&self) -> bool { 209 match *self { 210 Self::I64(_) => true, 211 _ => false, 212 } 213 } 214 215 /// Get the type of the value. 216 pub fn ty(&self) -> WasmType { 217 match self { 218 Val::I32(_) => WasmType::I32, 219 Val::I64(_) => WasmType::I64, 220 Val::F32(_) => WasmType::F32, 221 Val::F64(_) => WasmType::F64, 222 Val::Reg(r) => r.ty, 223 Val::Memory(m) => m.ty, 224 Val::Local(l) => l.ty, 225 } 226 } 227 } 228 229 /// The shadow stack used for compilation. 230 #[derive(Default, Debug)] 231 pub(crate) struct Stack { 232 inner: VecDeque<Val>, 233 } 234 235 impl Stack { 236 /// Allocate a new stack. 237 pub fn new() -> Self { 238 Self { 239 inner: Default::default(), 240 } 241 } 242 243 /// Returns true if the stack contains a local with the provided index 244 /// except if the only time the local appears is the top element. 245 pub fn contains_latent_local(&self, index: u32) -> bool { 246 self.inner 247 .iter() 248 // Iterate top-to-bottom so we can skip the top element and stop 249 // when we see a memory element. 250 .rev() 251 // The local is not latent if it's the top element because the top 252 // element will be popped next which materializes the local. 253 .skip(1) 254 // Stop when we see a memory element because that marks where we 255 // spilled up to so there will not be any locals past this point. 256 .take_while(|v| !v.is_mem()) 257 .any(|v| v.is_local_at_index(index)) 258 } 259 260 /// Extend the stack with the given elements. 261 pub fn extend(&mut self, values: impl IntoIterator<Item = Val>) { 262 self.inner.extend(values); 263 } 264 265 /// Inserts many values at the given index. 266 pub fn insert_many(&mut self, at: usize, values: impl IntoIterator<Item = Val>) { 267 debug_assert!(at <= self.len()); 268 // If last, simply extend. 269 if at == self.inner.len() { 270 self.inner.extend(values); 271 } else { 272 let mut tail = self.inner.split_off(at); 273 self.inner.extend(values); 274 self.inner.append(&mut tail); 275 } 276 } 277 278 /// Get the length of the stack. 279 pub fn len(&self) -> usize { 280 self.inner.len() 281 } 282 283 /// Push a value to the stack. 284 pub fn push(&mut self, val: Val) { 285 self.inner.push_back(val); 286 } 287 288 /// Peek into the top in the stack. 289 pub fn peek(&self) -> Option<&Val> { 290 self.inner.back() 291 } 292 293 /// Returns an iterator referencing the last n items of the stack, 294 /// in bottom-most to top-most order. 295 pub fn peekn(&self, n: usize) -> impl Iterator<Item = &Val> + '_ { 296 let len = self.len(); 297 assert!(n <= len); 298 299 let partition = len - n; 300 self.inner.range(partition..) 301 } 302 303 /// Pops the top element of the stack, if any. 304 pub fn pop(&mut self) -> Option<Val> { 305 self.inner.pop_back() 306 } 307 308 /// Pops the element at the top of the stack if it is an i32 const; 309 /// returns `None` otherwise. 310 pub fn pop_i32_const(&mut self) -> Option<i32> { 311 match self.peek() { 312 Some(v) => v.is_i32_const().then(|| self.pop().unwrap().get_i32()), 313 _ => None, 314 } 315 } 316 317 /// Pops the element at the top of the stack if it is an i64 const; 318 /// returns `None` otherwise. 319 pub fn pop_i64_const(&mut self) -> Option<i64> { 320 match self.peek() { 321 Some(v) => v.is_i64_const().then(|| self.pop().unwrap().get_i64()), 322 _ => None, 323 } 324 } 325 326 /// Pops the element at the top of the stack if it is a register; 327 /// returns `None` otherwise. 328 pub fn pop_reg(&mut self) -> Option<TypedReg> { 329 match self.peek() { 330 Some(v) => v.is_reg().then(|| self.pop().unwrap().get_reg()), 331 _ => None, 332 } 333 } 334 335 /// Pops the given register if it is at the top of the stack; 336 /// returns `None` otherwise. 337 pub fn pop_named_reg(&mut self, reg: Reg) -> Option<TypedReg> { 338 match self.peek() { 339 Some(v) => { 340 (v.is_reg() && v.get_reg().reg == reg).then(|| self.pop().unwrap().get_reg()) 341 } 342 _ => None, 343 } 344 } 345 346 /// Get a mutable reference to the inner stack representation. 347 pub fn inner_mut(&mut self) -> &mut VecDeque<Val> { 348 &mut self.inner 349 } 350 } 351 352 #[cfg(test)] 353 mod tests { 354 use super::{Stack, Val}; 355 use crate::isa::reg::Reg; 356 use wasmtime_environ::WasmType; 357 358 #[test] 359 fn test_pop_i32_const() { 360 let mut stack = Stack::new(); 361 stack.push(Val::i32(33i32)); 362 assert_eq!(33, stack.pop_i32_const().unwrap()); 363 364 stack.push(Val::local(10, WasmType::I32)); 365 assert!(stack.pop_i32_const().is_none()); 366 } 367 368 #[test] 369 fn test_pop_reg() { 370 let mut stack = Stack::new(); 371 let reg = Reg::int(2usize); 372 stack.push(Val::reg(reg, WasmType::I32)); 373 stack.push(Val::i32(4)); 374 375 assert_eq!(None, stack.pop_reg()); 376 let _ = stack.pop().unwrap(); 377 assert_eq!(reg, stack.pop_reg().unwrap().reg); 378 } 379 380 #[test] 381 fn test_pop_named_reg() { 382 let mut stack = Stack::new(); 383 let reg = Reg::int(2usize); 384 stack.push(Val::reg(reg, WasmType::I32)); 385 stack.push(Val::reg(Reg::int(4), WasmType::I32)); 386 387 assert_eq!(None, stack.pop_named_reg(reg)); 388 let _ = stack.pop().unwrap(); 389 assert_eq!(reg, stack.pop_named_reg(reg).unwrap().reg); 390 } 391 } 392