From 3bba88d10c9334aa6afe45074a2809cff57e7f8e Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Wed, 18 Feb 2026 12:12:39 -0800 Subject: [PATCH] Implement our OOM-handling `SecondaryMap` from scratch I realized we need to adjust its `V: Clone` bound into `V: TryClone` which means that we can no longer actually just wrap an inner `cranelift_entity::SecondaryMap` and need to instead implement our own. This also made me realize that we need `remove` to be fallible because, when the entry being removed is in bounds, it overwrites the entry with the default value, but that default value needs to be `TryClone`d now which is, of course, a fallible operation. --- .../environ/src/collections/secondary_map.rs | 543 ++++++++++++++++-- crates/fuzzing/tests/oom.rs | 38 +- crates/wasmtime/src/runtime/type_registry.rs | 19 +- 3 files changed, 530 insertions(+), 70 deletions(-) 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/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/type_registry.rs b/crates/wasmtime/src/runtime/type_registry.rs index abf9dc59e454..c4fb0a29b300 100644 --- a/crates/wasmtime/src/runtime/type_registry.rs +++ b/crates/wasmtime/src/runtime/type_registry.rs @@ -955,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"); } } @@ -995,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"); } } @@ -1075,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") { @@ -1150,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"); } }