diff --git a/cranelift/entity/src/packed_option.rs b/cranelift/entity/src/packed_option.rs index 36f97d6d4da2..fd1451802a8f 100644 --- a/cranelift/entity/src/packed_option.rs +++ b/cranelift/entity/src/packed_option.rs @@ -7,8 +7,8 @@ //! This module provides a `PackedOption` for types that have a reserved value that can be used //! to represent `None`. -use core::fmt; -use core::mem; +use core::{fmt, mem}; +use wasmtime_core::{alloc::TryClone, error::OutOfMemory}; #[cfg(feature = "enable-serde")] use serde_derive::{Deserialize, Serialize}; @@ -30,6 +30,15 @@ pub trait ReservedValue { #[repr(transparent)] pub struct PackedOption(T); +impl TryClone for PackedOption +where + T: ReservedValue + TryClone, +{ + fn try_clone(&self) -> Result { + Ok(Self(self.0.try_clone()?)) + } +} + impl PackedOption { /// Returns `true` if the packed option is a `None` value. pub fn is_none(&self) -> bool { diff --git a/crates/core/src/alloc/try_clone.rs b/crates/core/src/alloc/try_clone.rs index c59b62380b5e..85837f94c778 100644 --- a/crates/core/src/alloc/try_clone.rs +++ b/crates/core/src/alloc/try_clone.rs @@ -1,5 +1,6 @@ use crate::error::OutOfMemory; use core::mem; +use std_alloc::sync::Arc; /// A trait for values that can be cloned, but contain owned, heap-allocated /// values whose allocations may fail during cloning. @@ -39,6 +40,40 @@ where } } +impl TryClone for Result +where + T: TryClone, + E: TryClone, +{ + #[inline] + fn try_clone(&self) -> Result { + match self { + Ok(x) => Ok(Ok(x.try_clone()?)), + Err(e) => Ok(Err(e.try_clone()?)), + } + } +} + +impl TryClone for Option +where + T: TryClone, +{ + #[inline] + fn try_clone(&self) -> Result { + match self { + Some(x) => Ok(Some(x.try_clone()?)), + None => Ok(None), + } + } +} + +impl TryClone for Arc { + #[inline] + fn try_clone(&self) -> Result { + Ok(self.clone()) + } +} + macro_rules! impl_try_clone_via_clone { ( $( $ty:ty ),* $(,)? ) => { $( diff --git a/crates/core/src/alloc/vec.rs b/crates/core/src/alloc/vec.rs index 681371f03cc8..c1cd3d3360a3 100644 --- a/crates/core/src/alloc/vec.rs +++ b/crates/core/src/alloc/vec.rs @@ -288,6 +288,11 @@ impl Vec { // use `std`'s `into_boxed_slice` without fear of `realloc`. Ok(self.inner.into_boxed_slice()) } + + /// Same as [`std::vec::Vec::clear`]. + pub fn clear(&mut self) { + self.inner.clear(); + } } impl Deref for Vec { diff --git a/crates/environ/src/collections/secondary_map.rs b/crates/environ/src/collections/secondary_map.rs index cfa1c961c234..11c2c4891b98 100644 --- a/crates/environ/src/collections/secondary_map.rs +++ b/crates/environ/src/collections/secondary_map.rs @@ -1,39 +1,51 @@ -use crate::{collections::Vec, error::OutOfMemory}; -use core::{fmt, ops::Index}; -use cranelift_entity::{EntityRef, SecondaryMap as Inner}; -use serde::ser::SerializeSeq; +use crate::{ + EntityRef, + collections::{TryClone, Vec}, + error::OutOfMemory, +}; +use core::{cmp::Ordering, fmt, marker::PhantomData, mem, ops::Index}; /// Like [`cranelift_entity::SecondaryMap`] but all allocation is fallible. -pub struct SecondaryMap -where - K: EntityRef, - V: Clone, -{ - inner: Inner, +pub struct SecondaryMap { + elems: Vec, + default_value: V, + phantom: PhantomData V>, } impl fmt::Debug for SecondaryMap where K: EntityRef + fmt::Debug, - V: fmt::Debug + Clone, + V: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt::Debug::fmt(&self.inner, f) + struct Entries<'a, K, V>(&'a SecondaryMap); + + impl<'a, K, V> fmt::Debug for Entries<'a, K, V> + where + K: EntityRef + fmt::Debug, + V: fmt::Debug, + { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_map().entries(self.0.iter()).finish() + } + } + + f.debug_struct("SecondaryMap") + .field("entries", &Entries(self)) + .finish() } } -impl SecondaryMap -where - K: EntityRef, - V: Clone, -{ +impl SecondaryMap { /// Same as [`cranelift_entity::SecondaryMap::new`]. pub fn new() -> Self where V: Default, { Self { - inner: Inner::new(), + elems: Vec::new(), + default_value: V::default(), + phantom: PhantomData, } } @@ -43,87 +55,134 @@ where V: Default, { Ok(Self { - inner: Inner::try_with_capacity(capacity)?, + elems: Vec::with_capacity(capacity)?, + default_value: V::default(), + phantom: PhantomData, }) } /// Same as [`cranelift_entity::SecondaryMap::with_default`]. - pub fn with_default(default: V) -> Self { + pub fn with_default(default_value: V) -> Self { Self { - inner: Inner::with_default(default), + elems: Vec::new(), + default_value, + phantom: PhantomData, } } /// Same as [`cranelift_entity::SecondaryMap::capacity`]. pub fn capacity(&self) -> usize { - self.inner.capacity() + self.elems.capacity() } /// Same as [`cranelift_entity::SecondaryMap::get`]. - pub fn get(&self, k: K) -> Option<&V> { - self.inner.get(k) + pub fn get(&self, k: K) -> Option<&V> + where + K: EntityRef, + { + self.elems.get(k.index()) } /// Same as [`cranelift_entity::SecondaryMap::get_mut`]. - pub fn get_mut(&mut self, k: K) -> Option<&mut V> { - self.inner.get_mut(k) + pub fn get_mut(&mut self, k: K) -> Option<&mut V> + where + K: EntityRef, + { + self.elems.get_mut(k.index()) } /// Same as [`cranelift_entity::SecondaryMap::try_insert`]. - pub fn insert(&mut self, k: K, v: V) -> Result, OutOfMemory> { - self.inner.try_insert(k, v) + pub fn insert(&mut self, k: K, v: V) -> Result, OutOfMemory> + where + K: EntityRef, + V: TryClone, + { + if k.index() < self.elems.len() { + Ok(Some(mem::replace(&mut self.elems[k.index()], v))) + } else { + self.resize(k.index() + 1)?; + self.elems[k.index()] = v; + Ok(None) + } } - /// Same as [`cranelift_entity::SecondaryMap::remove`]. - pub fn remove(&mut self, k: K) -> Option { - self.inner.remove(k) + /// Same as [`cranelift_entity::SecondaryMap::remove`] but returns an error + /// if `TryClone`ing the default value fails when overwriting the old entry, + /// if any. + pub fn remove(&mut self, k: K) -> Result, OutOfMemory> + where + K: EntityRef, + V: TryClone, + { + if k.index() < self.elems.len() { + let default = self.default_value.try_clone()?; + Ok(Some(mem::replace(&mut self.elems[k.index()], default))) + } else { + Ok(None) + } } /// Same as [`cranelift_entity::SecondaryMap::is_empty`]. pub fn is_empty(&self) -> bool { - self.inner.is_empty() + self.elems.is_empty() } /// Same as [`cranelift_entity::SecondaryMap::clear`]. pub fn clear(&mut self) { - self.inner.clear() + self.elems.clear(); } /// Same as [`cranelift_entity::SecondaryMap::iter`]. - pub fn iter(&self) -> cranelift_entity::Iter<'_, K, V> { - self.inner.iter() + pub fn iter(&self) -> Iter<'_, K, V> { + Iter { + inner: self.elems.iter().enumerate(), + phantom: PhantomData, + } } /// Same as [`cranelift_entity::SecondaryMap::iter_mut`]. - pub fn iter_mut(&mut self) -> cranelift_entity::IterMut<'_, K, V> { - self.inner.iter_mut() + pub fn iter_mut(&mut self) -> IterMut<'_, K, V> { + IterMut { + inner: self.elems.iter_mut().enumerate(), + phantom: PhantomData, + } } /// Same as [`cranelift_entity::SecondaryMap::keys`]. - pub fn keys(&self) -> cranelift_entity::Keys { - self.inner.keys() + pub fn keys(&self) -> Keys { + Keys { + inner: 0..self.elems.len(), + phantom: PhantomData, + } } /// Same as [`cranelift_entity::SecondaryMap::values`]. pub fn values(&self) -> core::slice::Iter<'_, V> { - self.inner.values() + self.elems.iter() } /// Same as [`cranelift_entity::SecondaryMap::values_mut`]. pub fn values_mut(&mut self) -> core::slice::IterMut<'_, V> { - self.inner.values_mut() + self.elems.iter_mut() } /// Resize the map to have `n` entries by adding default entries as needed. - pub fn resize(&mut self, n: usize) -> Result<(), OutOfMemory> { - self.inner.try_resize(n) + pub fn resize(&mut self, n: usize) -> Result<(), OutOfMemory> + where + V: TryClone, + { + match self.elems.len().cmp(&n) { + Ordering::Less => self.elems.resize(n, self.default_value.try_clone()?)?, + Ordering::Equal => {} + Ordering::Greater => self.elems.truncate(n), + } + Ok(()) } } impl Default for SecondaryMap where - K: EntityRef, - V: Clone + Default, + V: Default, { fn default() -> SecondaryMap { SecondaryMap::new() @@ -135,40 +194,60 @@ where impl Index for SecondaryMap where K: EntityRef, - V: Clone, { type Output = V; fn index(&self, k: K) -> &V { - &self.inner[k] + self.get(k).unwrap_or(&self.default_value) } } impl From> for SecondaryMap where K: EntityRef, - V: Clone + Default, + V: TryClone + Default, { - fn from(values: Vec) -> Self { - let values: alloc::vec::Vec = values.into(); - let inner = Inner::from(values); - Self { inner } + fn from(elems: Vec) -> Self { + Self { + elems, + default_value: V::default(), + phantom: PhantomData, + } } } impl serde::ser::Serialize for SecondaryMap where K: EntityRef, - V: Clone + serde::ser::Serialize, + V: PartialEq + serde::ser::Serialize, { fn serialize(&self, serializer: S) -> Result where S: serde::Serializer, { - let mut seq = serializer.serialize_seq(Some(self.capacity()))?; - for elem in self.values() { - seq.serialize_element(elem)?; + use serde::ser::SerializeSeq as _; + + // Ignore any trailing default values. + let mut len = self.capacity(); + while len > 0 && &self[K::new(len - 1)] == &self.default_value { + len -= 1; } + + // Plus one for the default value. + let mut seq = serializer.serialize_seq(Some(len + 1))?; + + // Always serialize the default value first. + seq.serialize_element(&self.default_value)?; + + for elem in self.values().take(len) { + let elem = if elem == &self.default_value { + None + } else { + Some(elem) + }; + seq.serialize_element(&elem)?; + } + seq.end() } } @@ -176,13 +255,359 @@ where impl<'de, K, V> serde::de::Deserialize<'de> for SecondaryMap where K: EntityRef, - V: Clone + Default + serde::de::Deserialize<'de>, + V: TryClone + serde::de::Deserialize<'de>, { fn deserialize(deserializer: D) -> Result where D: serde::Deserializer<'de>, { - let values: Vec = serde::de::Deserialize::deserialize(deserializer)?; - Ok(Self::from(values)) + struct Visitor(core::marker::PhantomData SecondaryMap>) + where + K: EntityRef, + V: TryClone; + + impl<'de, K, V> serde::de::Visitor<'de> for Visitor + where + K: EntityRef, + V: TryClone + serde::de::Deserialize<'de>, + { + type Value = SecondaryMap; + + fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str("struct SecondaryMap") + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: serde::de::SeqAccess<'de>, + { + // Minus one to account for the default element, which is always + // the first in the sequence. + let size_hint = seq.size_hint().and_then(|n| n.checked_sub(1)); + + let Some(default) = seq.next_element::()? else { + return Err(serde::de::Error::custom("Default value required")); + }; + + let mut map = SecondaryMap::::with_default( + default + .try_clone() + .map_err(|oom| serde::de::Error::custom(oom))?, + ); + + if let Some(n) = size_hint { + map.resize(n).map_err(|oom| serde::de::Error::custom(oom))?; + } + + let mut idx = 0; + while let Some(val) = seq.next_element::>()? { + let key = K::new(idx); + let val = match val { + None => default + .try_clone() + .map_err(|oom| serde::de::Error::custom(oom))?, + Some(val) => val, + }; + + map.insert(key, val) + .map_err(|oom| serde::de::Error::custom(oom))?; + + idx += 1; + } + + Ok(map) + } + } + + deserializer.deserialize_seq(Visitor(core::marker::PhantomData)) + } +} + +/// A shared iterator over a `SecondaryMap`. +pub struct Iter<'a, K, V> { + inner: core::iter::Enumerate>, + phantom: PhantomData K>, +} + +impl<'a, K, V> Iterator for Iter<'a, K, V> +where + K: EntityRef, +{ + type Item = (K, &'a V); + + fn next(&mut self) -> Option { + let (i, v) = self.inner.next()?; + Some((K::new(i), v)) + } +} + +/// An exclusive iterator over a `SecondaryMap`. +pub struct IterMut<'a, K, V> { + inner: core::iter::Enumerate>, + phantom: PhantomData K>, +} + +impl<'a, K, V> Iterator for IterMut<'a, K, V> +where + K: EntityRef, +{ + type Item = (K, &'a mut V); + + fn next(&mut self) -> Option { + let (i, v) = self.inner.next()?; + Some((K::new(i), v)) + } +} + +/// An iterator over the keys in a `SecondaryMap`. +pub struct Keys { + inner: core::ops::Range, + phantom: PhantomData K>, +} + +impl Iterator for Keys +where + K: EntityRef, +{ + type Item = K; + + fn next(&mut self) -> Option { + Some(K::new(self.inner.next()?)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::Result; + use alloc::vec; + use alloc::vec::Vec; + + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] + struct K(u32); + crate::entity_impl!(K); + + fn k(i: usize) -> K { + K::new(i) + } + + #[test] + fn with_capacity() -> Result<()> { + let map = SecondaryMap::::with_capacity(100)?; + assert!(map.capacity() >= 100); + Ok(()) + } + + #[test] + fn with_default() -> Result<()> { + let map = SecondaryMap::::with_default(42); + assert_eq!(map[k(99)], 42); + Ok(()) + } + + #[test] + fn get() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(10), 99)?; + assert_eq!(map.get(k(0)).copied(), Some(0)); + assert_eq!(map.get(k(10)).copied(), Some(99)); + assert!(map.get(k(100)).is_none()); + Ok(()) + } + + #[test] + fn get_mut() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(10), 99)?; + *map.get_mut(k(0)).unwrap() += 1; + *map.get_mut(k(10)).unwrap() += 1; + assert_eq!(map[k(0)], 1); + assert_eq!(map[k(10)], 100); + assert!(map.get_mut(k(100)).is_none()); + Ok(()) + } + + #[test] + fn insert() -> Result<()> { + let mut map = SecondaryMap::new(); + assert_eq!(map[k(3)], 0); + map.insert(k(3), 99)?; + assert_eq!(map[k(3)], 99); + Ok(()) + } + + #[test] + fn remove() -> Result<()> { + let mut map = SecondaryMap::new(); + + let old = map.remove(k(3))?; + assert!(old.is_none()); + + map.insert(k(3), 99)?; + + let old = map.remove(k(3))?; + assert_eq!(old, Some(99)); + assert_eq!(map[k(3)], 0); + + let old = map.remove(k(3))?; + assert_eq!(old, Some(0)); + + Ok(()) + } + + #[test] + fn is_empty() -> Result<()> { + let mut map = SecondaryMap::new(); + assert!(map.is_empty()); + map.insert(k(0), 1)?; + assert!(!map.is_empty()); + Ok(()) + } + + #[test] + fn clear() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(0), 1)?; + map.clear(); + assert!(map.is_empty()); + assert_eq!(map[k(0)], 0); + Ok(()) + } + + #[test] + fn iter() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(0), 'a')?; + map.insert(k(1), 'b')?; + map.insert(k(5), 'c')?; + assert_eq!( + map.iter().collect::>(), + vec![ + (k(0), &'a'), + (k(1), &'b'), + (k(2), &char::default()), + (k(3), &char::default()), + (k(4), &char::default()), + (k(5), &'c'), + ], + ); + Ok(()) + } + + #[test] + fn iter_mut() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(0), 'a')?; + map.insert(k(1), 'b')?; + map.insert(k(5), 'c')?; + assert_eq!( + map.iter_mut().collect::>(), + vec![ + (k(0), &mut 'a'), + (k(1), &mut 'b'), + (k(2), &mut char::default()), + (k(3), &mut char::default()), + (k(4), &mut char::default()), + (k(5), &mut 'c'), + ], + ); + Ok(()) + } + + #[test] + fn keys() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(2), 9)?; + assert_eq!(map.keys().collect::>(), vec![k(0), k(1), k(2)]); + Ok(()) + } + + #[test] + fn values() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(2), 9)?; + assert_eq!(map.values().collect::>(), vec![&0, &0, &9]); + Ok(()) + } + + #[test] + fn values_mut() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(2), 9)?; + assert_eq!( + map.values_mut().collect::>(), + vec![&mut 0, &mut 0, &mut 9] + ); + Ok(()) + } + + #[test] + fn resize() -> Result<()> { + let mut map = SecondaryMap::::new(); + assert!(map.is_empty()); + + // Grow via resize. + map.resize(2)?; + assert!(!map.is_empty()); + assert!(map.get(k(0)).is_some()); + assert!(map.get(k(1)).is_some()); + assert!(map.get(k(2)).is_none()); + + // Resize to same size. + map.resize(2)?; + assert!(!map.is_empty()); + assert!(map.get(k(0)).is_some()); + assert!(map.get(k(1)).is_some()); + assert!(map.get(k(2)).is_none()); + + // Shrink via resize. + map.resize(1)?; + assert!(!map.is_empty()); + assert!(map.get(k(0)).is_some()); + assert!(map.get(k(1)).is_none()); + + Ok(()) + } + + #[test] + fn index() -> Result<()> { + let mut map = SecondaryMap::new(); + map.insert(k(0), 55)?; + assert_eq!(map[k(0)], 55); + assert_eq!(map[k(999)], 0); + Ok(()) + } + + #[test] + fn from_vec() -> Result<()> { + let v = crate::collections::vec![10, 20, 30]?; + let map = SecondaryMap::from(v); + assert_eq!(map[k(0)], 10); + assert_eq!(map[k(1)], 20); + assert_eq!(map[k(2)], 30); + assert_eq!(map[k(3)], 0); + Ok(()) + } + + #[test] + fn serialize_deserialize() -> Result<()> { + let mut map = SecondaryMap::::with_default(99); + map.insert(k(0), 33)?; + map.insert(k(1), 44)?; + map.insert(k(2), 55)?; + map.insert(k(3), 99)?; + map.insert(k(4), 99)?; + + let bytes = postcard::to_allocvec(&map)?; + let map2: SecondaryMap = postcard::from_bytes(&bytes)?; + + for i in 0..10 { + assert_eq!(map[k(i)], map2[k(i)]); + } + + // Trailing default entries were omitted from the serialization. + assert_eq!(map2.keys().collect::>(), vec![k(0), k(1), k(2)]); + + Ok(()) } } diff --git a/crates/environ/src/gc.rs b/crates/environ/src/gc.rs index f15e2b8cb93b..68c5ebf5ba22 100644 --- a/crates/environ/src/gc.rs +++ b/crates/environ/src/gc.rs @@ -16,10 +16,13 @@ pub mod drc; pub mod null; use crate::{ - WasmArrayType, WasmCompositeInnerType, WasmCompositeType, WasmStorageType, WasmStructType, - WasmValType, + WasmArrayType, WasmCompositeInnerType, WasmCompositeType, WasmExnType, WasmStorageType, + WasmStructType, WasmValType, + collections::{self, TryClone}, + error::OutOfMemory, + prelude::*, }; -use crate::{WasmExnType, prelude::*}; +use alloc::sync::Arc; use core::alloc::Layout; /// Discriminant to check whether GC reference is an `i31ref` or not. @@ -122,7 +125,9 @@ fn common_struct_or_exn_layout( fields: &[crate::WasmFieldType], header_size: u32, header_align: u32, -) -> (u32, u32, Vec) { +) -> (u32, u32, collections::Vec) { + use crate::PanicOnOom as _; + // Process each field, aligning it to its natural alignment. // // We don't try and do any fancy field reordering to minimize padding (yet?) @@ -143,7 +148,8 @@ fn common_struct_or_exn_layout( let is_gc_ref = f.element_type.is_vmgcref_type_and_not_i31(); GcStructLayoutField { offset, is_gc_ref } }) - .collect(); + .try_collect::, _>() + .panic_on_oom(); // Ensure that the final size is a multiple of the alignment, for // simplicity. @@ -228,12 +234,12 @@ pub trait GcTypeLayouts { assert!(!ty.shared); match &ty.inner { WasmCompositeInnerType::Array(ty) => Some(self.array_layout(ty).into()), - WasmCompositeInnerType::Struct(ty) => Some(self.struct_layout(ty).into()), + WasmCompositeInnerType::Struct(ty) => Some(Arc::new(self.struct_layout(ty)).into()), WasmCompositeInnerType::Func(_) => None, WasmCompositeInnerType::Cont(_) => { unimplemented!("Stack switching feature not compatible with GC, yet") } - WasmCompositeInnerType::Exn(ty) => Some(self.exn_layout(ty).into()), + WasmCompositeInnerType::Exn(ty) => Some(Arc::new(self.exn_layout(ty)).into()), } } @@ -254,7 +260,7 @@ pub enum GcLayout { Array(GcArrayLayout), /// The layout of a GC-managed struct or exception object. - Struct(GcStructLayout), + Struct(Arc), } impl From for GcLayout { @@ -263,16 +269,22 @@ impl From for GcLayout { } } -impl From for GcLayout { - fn from(layout: GcStructLayout) -> Self { +impl From> for GcLayout { + fn from(layout: Arc) -> Self { Self::Struct(layout) } } +impl TryClone for GcLayout { + fn try_clone(&self) -> core::result::Result { + Ok(self.clone()) + } +} + impl GcLayout { /// Get the underlying `GcStructLayout`, or panic. #[track_caller] - pub fn unwrap_struct(&self) -> &GcStructLayout { + pub fn unwrap_struct(&self) -> &Arc { match self { Self::Struct(s) => s, _ => panic!("GcLayout::unwrap_struct on non-struct GC layout"), @@ -368,12 +380,23 @@ pub struct GcStructLayout { /// The fields of this struct. The `i`th entry contains information about /// the `i`th struct field's layout. - pub fields: Vec, + pub fields: collections::Vec, /// Whether this is an exception object layout. pub is_exception: bool, } +impl TryClone for GcStructLayout { + fn try_clone(&self) -> Result { + Ok(GcStructLayout { + size: self.size, + align: self.align, + fields: self.fields.try_clone()?, + is_exception: self.is_exception, + }) + } +} + impl GcStructLayout { /// Get a `core::alloc::Layout` for a struct of this type. pub fn layout(&self) -> Layout { @@ -397,6 +420,12 @@ pub struct GcStructLayoutField { pub is_gc_ref: bool, } +impl TryClone for GcStructLayoutField { + fn try_clone(&self) -> Result { + Ok(*self) + } +} + /// The kind of an object in a GC heap. /// /// Note that this type is accessed from Wasm JIT code. diff --git a/crates/environ/src/module_types.rs b/crates/environ/src/module_types.rs index 56999f895dff..0da4fb721d59 100644 --- a/crates/environ/src/module_types.rs +++ b/crates/environ/src/module_types.rs @@ -1,8 +1,9 @@ use crate::{ - ModuleInternedRecGroupIndex, ModuleInternedTypeIndex, PrimaryMap, TypeTrace, WasmSubType, + ModuleInternedRecGroupIndex, ModuleInternedTypeIndex, PanicOnOom as _, TypeTrace, WasmSubType, + collections::{PrimaryMap, SecondaryMap}, + packed_option::PackedOption, }; use core::ops::{Index, Range}; -use cranelift_entity::{SecondaryMap, packed_option::PackedOption}; use serde_derive::{Deserialize, Serialize}; /// All types used in a core wasm module. @@ -77,7 +78,7 @@ impl ModuleTypes { /// Adds a new type to this interned list of types. pub fn push(&mut self, ty: WasmSubType) -> ModuleInternedTypeIndex { - self.wasm_types.push(ty) + self.wasm_types.push(ty).panic_on_oom() } /// Iterate over the trampoline function types that this module requires. @@ -131,7 +132,9 @@ impl ModuleTypes { .is_trampoline_type() ); - self.trampoline_types[for_ty] = Some(trampoline_ty).into(); + self.trampoline_types + .insert(for_ty, Some(trampoline_ty).into()) + .panic_on_oom(); } /// Adds a new rec group to this interned list of types. @@ -139,12 +142,12 @@ impl ModuleTypes { &mut self, range: Range, ) -> ModuleInternedRecGroupIndex { - self.rec_groups.push(range) + self.rec_groups.push(range).panic_on_oom() } /// Reserves space for `amt` more types. pub fn reserve(&mut self, amt: usize) { - self.wasm_types.reserve(amt) + self.wasm_types.reserve(amt).panic_on_oom() } /// Returns the next return value of `push_rec_group`. diff --git a/crates/fuzzing/tests/oom.rs b/crates/fuzzing/tests/oom.rs index 960d5ab2f816..fc71e22d1137 100644 --- a/crates/fuzzing/tests/oom.rs +++ b/crates/fuzzing/tests/oom.rs @@ -9,7 +9,7 @@ use std::{ use wasmtime::{error::OutOfMemory, *}; use wasmtime_core::alloc::TryCollect; use wasmtime_environ::{ - PrimaryMap, SecondaryMap, + EntityRef, PrimaryMap, collections::{self, *}, }; use wasmtime_fuzzing::oom::{OomTest, OomTestAllocator}; @@ -169,27 +169,51 @@ fn primary_map_try_reserve_exact() -> Result<()> { } #[test] -fn secondary_map_try_with_capacity() -> Result<()> { +fn secondary_map_with_capacity() -> Result<()> { OomTest::new().test(|| { - let _map = SecondaryMap::::try_with_capacity(32)?; + let _map = SecondaryMap::::with_capacity(32)?; Ok(()) }) } #[test] -fn secondary_map_try_resize() -> Result<()> { +fn secondary_map_resize() -> Result<()> { OomTest::new().test(|| { let mut map = SecondaryMap::::new(); - map.try_resize(100)?; + map.resize(100)?; Ok(()) }) } #[test] -fn secondary_map_try_insert() -> Result<()> { +fn secondary_map_insert() -> Result<()> { OomTest::new().test(|| { let mut map = SecondaryMap::::new(); - map.try_insert(Key::from_u32(42), 100)?; + map.insert(Key::from_u32(42), 100)?; + Ok(()) + }) +} + +#[test] +fn secondary_map_remove() -> Result<()> { + OomTest::new().test(|| { + let mut default = Vec::new(); + default.push(3)?; + default.push(3)?; + default.push(3)?; + + let mut map = collections::SecondaryMap::>::with_default(default); + assert_eq!(&*map[Key::new(0)], &[3, 3, 3]); + + map.insert(Key::new(0), Vec::new())?; + assert!(map[Key::new(0)].is_empty()); + + // This may fail because it requires `TryClone`ing the default value. + let old = map.remove(Key::new(0))?; + + assert!(old.unwrap().is_empty()); + assert_eq!(&*map[Key::new(0)], &[3, 3, 3]); + Ok(()) }) } diff --git a/crates/wasmtime/src/runtime/gc/enabled/exnref.rs b/crates/wasmtime/src/runtime/gc/enabled/exnref.rs index 91ac8ad54b53..a456a1b2fe54 100644 --- a/crates/wasmtime/src/runtime/gc/enabled/exnref.rs +++ b/crates/wasmtime/src/runtime/gc/enabled/exnref.rs @@ -11,6 +11,7 @@ use crate::{ store::{AutoAssertNoGc, StoreOpaque}, }; use crate::{ExnType, FieldType, GcHeapOutOfMemory, StoreContextMut, Tag, prelude::*}; +use alloc::sync::Arc; use core::mem; use core::mem::MaybeUninit; use wasmtime_environ::{GcLayout, GcStructLayout, VMGcKind, VMSharedTypeIndex}; @@ -563,7 +564,7 @@ impl ExnRef { Ok(gc_ref.as_exnref_unchecked()) } - fn layout(&self, store: &AutoAssertNoGc<'_>) -> Result { + fn layout(&self, store: &AutoAssertNoGc<'_>) -> Result> { assert!(self.comes_from_same_store(&store)); let type_index = self.type_index(store)?; let layout = store diff --git a/crates/wasmtime/src/runtime/gc/enabled/structref.rs b/crates/wasmtime/src/runtime/gc/enabled/structref.rs index 144fd31355a8..b09afd2c5b42 100644 --- a/crates/wasmtime/src/runtime/gc/enabled/structref.rs +++ b/crates/wasmtime/src/runtime/gc/enabled/structref.rs @@ -12,6 +12,7 @@ use crate::{ prelude::*, store::{AutoAssertNoGc, StoreContextMut, StoreOpaque, StoreResourceLimiter}, }; +use alloc::sync::Arc; use core::mem::{self, MaybeUninit}; use wasmtime_environ::{GcLayout, GcStructLayout, VMGcKind, VMSharedTypeIndex}; @@ -515,7 +516,7 @@ impl StructRef { Ok(gc_ref.as_structref_unchecked()) } - fn layout(&self, store: &AutoAssertNoGc<'_>) -> Result { + fn layout(&self, store: &AutoAssertNoGc<'_>) -> Result> { assert!(self.comes_from_same_store(&store)); let type_index = self.type_index(store)?; let layout = store diff --git a/crates/wasmtime/src/runtime/type_registry.rs b/crates/wasmtime/src/runtime/type_registry.rs index bed00a226e6e..0fea5bc10721 100644 --- a/crates/wasmtime/src/runtime/type_registry.rs +++ b/crates/wasmtime/src/runtime/type_registry.rs @@ -26,7 +26,7 @@ use wasmtime_core::slab::{Id as SlabId, Slab}; use wasmtime_environ::{ EngineOrModuleTypeIndex, EntityRef, GcLayout, ModuleInternedTypeIndex, ModuleTypes, TypeTrace, Undo, VMSharedTypeIndex, WasmRecGroup, WasmSubType, - collections::{HashSet, PrimaryMap, SecondaryMap, TryClone as _, Vec}, + collections::{HashSet, PrimaryMap, SecondaryMap, TryClone, Vec}, iter_entity_range, packed_option::{PackedOption, ReservedValue}, }; @@ -285,6 +285,12 @@ impl Clone for RegisteredType { } } +impl TryClone for RegisteredType { + fn try_clone(&self) -> Result { + Ok(self.clone()) + } +} + impl Drop for RegisteredType { fn drop(&mut self) { self.engine.signatures().debug_assert_contains(self.index); @@ -457,6 +463,12 @@ impl RegisteredType { #[derive(Clone)] struct RecGroupEntry(Arc); +impl TryClone for RecGroupEntry { + fn try_clone(&self) -> Result { + Ok(self.clone()) + } +} + impl Debug for RecGroupEntry { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { struct FormatAsPtr<'a, P>(&'a P); @@ -943,7 +955,9 @@ impl TypeRegistryInner { fn remove_entry_rec_groups(&mut self, entry: &RecGroupEntry) { for &ty in &entry.0.shared_type_indices { debug_assert!(ty.index() < self.type_to_rec_group.capacity()); - self.type_to_rec_group.remove(ty); + self.type_to_rec_group + .remove(ty) + .expect("`None.try_clone()` cannot fail"); } } @@ -983,7 +997,9 @@ impl TypeRegistryInner { } for &ty in &entry.0.shared_type_indices { - self.type_to_supertypes.remove(ty); + self.type_to_supertypes + .remove(ty) + .expect("`None.try_clone()` cannot fail"); } } @@ -1063,7 +1079,12 @@ impl TypeRegistryInner { } for &ty in &entry.0.shared_type_indices { - if let Some(tramp_ty) = self.type_to_trampoline.remove(ty).and_then(|x| x.expand()) { + if let Some(tramp_ty) = self + .type_to_trampoline + .remove(ty) + .expect("`PackedOption::default().try_clone()` cannot fail") + .and_then(|x| x.expand()) + { self.debug_assert_registered(tramp_ty); let tramp_entry = self.type_to_rec_group[tramp_ty].as_ref().unwrap(); if tramp_entry.decref("dropping rec group's trampoline-type references") { @@ -1107,10 +1128,10 @@ impl TypeRegistryInner { gc_runtime.layouts().array_layout(a).into() } wasmtime_environ::WasmCompositeInnerType::Struct(s) => { - gc_runtime.layouts().struct_layout(s).into() + try_new::>(gc_runtime.layouts().struct_layout(s))?.into() } wasmtime_environ::WasmCompositeInnerType::Exn(e) => { - gc_runtime.layouts().exn_layout(e).into() + try_new::>(gc_runtime.layouts().exn_layout(e))?.into() } wasmtime_environ::WasmCompositeInnerType::Cont(_) => continue, // FIXME: #10248 stack switching support. }; @@ -1138,7 +1159,9 @@ impl TypeRegistryInner { } for ty in &entry.0.shared_type_indices { - self.type_to_gc_layout.remove(*ty); + self.type_to_gc_layout + .remove(*ty) + .expect("`None.try_clone()` cannot fail"); } }