1 use crate::{ 2 alloc::{TryClone, str_ptr_from_raw_parts, try_realloc}, 3 error::OutOfMemory, 4 }; 5 use core::{fmt, mem, ops}; 6 use std_alloc::{alloc::Layout, boxed::Box, string as inner}; 7 8 /// A newtype wrapper around [`std::string::String`] that only exposes 9 /// fallible-allocation methods. 10 #[derive(Default, Hash, PartialEq, Eq, PartialOrd, Ord)] 11 pub struct String { 12 inner: inner::String, 13 } 14 15 impl TryClone for String { 16 fn try_clone(&self) -> Result<Self, OutOfMemory> { 17 let mut s = Self::new(); 18 s.push_str(self)?; 19 Ok(s) 20 } 21 } 22 23 impl fmt::Debug for String { 24 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 25 fmt::Debug::fmt(&self.inner, f) 26 } 27 } 28 29 impl fmt::Display for String { 30 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 31 fmt::Display::fmt(&self.inner, f) 32 } 33 } 34 35 impl ops::Deref for String { 36 type Target = str; 37 38 #[inline] 39 fn deref(&self) -> &Self::Target { 40 &self.inner 41 } 42 } 43 44 impl ops::DerefMut for String { 45 #[inline] 46 fn deref_mut(&mut self) -> &mut Self::Target { 47 &mut self.inner 48 } 49 } 50 51 impl From<inner::String> for String { 52 #[inline] 53 fn from(inner: inner::String) -> Self { 54 Self { inner } 55 } 56 } 57 58 impl serde::ser::Serialize for String { 59 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> 60 where 61 S: serde::Serializer, 62 { 63 serializer.serialize_str(self) 64 } 65 } 66 67 impl<'de> serde::de::Deserialize<'de> for String { 68 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> 69 where 70 D: serde::Deserializer<'de>, 71 { 72 struct Visitor; 73 74 impl<'de> serde::de::Visitor<'de> for Visitor { 75 type Value = String; 76 77 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { 78 f.write_str("a `wasmtime_core::alloc::String` str") 79 } 80 81 fn visit_str<E>(self, v: &str) -> Result<Self::Value, E> 82 where 83 E: serde::de::Error, 84 { 85 let mut s = String::new(); 86 s.reserve_exact(v.len()).map_err(|oom| E::custom(oom))?; 87 s.push_str(v).expect("reserved capacity"); 88 Ok(s) 89 } 90 } 91 92 // NB: do not use `deserialize_string` as that eagerly allocates the 93 // `String` and does not give us a chance to handle OOM. Instead, use 94 // `deserialize_str` which passes the visitor the borrowed `str`, giving 95 // us a chance to fallibly allocate space. 96 deserializer.deserialize_str(Visitor) 97 } 98 } 99 100 impl String { 101 /// Same as [`std::string::String::new`]. 102 #[inline] 103 pub fn new() -> Self { 104 Self { 105 inner: inner::String::new(), 106 } 107 } 108 109 /// Same as [`std::string::String::with_capacity`] but returns an error on 110 /// allocation failure. 111 #[inline] 112 pub fn with_capacity(capacity: usize) -> Result<Self, OutOfMemory> { 113 let mut s = Self::new(); 114 s.reserve(capacity)?; 115 Ok(s) 116 } 117 118 /// Same as [`std::string::String::capacity`]. 119 #[inline] 120 pub fn capacity(&self) -> usize { 121 self.inner.capacity() 122 } 123 124 /// Same as [`std::string::String::as_str`]. 125 #[inline] 126 pub const fn as_str(&self) -> &str { 127 self.inner.as_str() 128 } 129 130 /// Same as [`std::string::String::reserve`] but returns an error on 131 /// allocation failure. 132 #[inline] 133 pub fn reserve(&mut self, additional: usize) -> Result<(), OutOfMemory> { 134 self.inner 135 .try_reserve(additional) 136 .map_err(|_| OutOfMemory::new(self.len().saturating_add(additional))) 137 } 138 139 /// Same as [`std::string::String::reserve_exact`] but returns an error on 140 /// allocation failure. 141 #[inline] 142 pub fn reserve_exact(&mut self, additional: usize) -> Result<(), OutOfMemory> { 143 self.inner 144 .try_reserve_exact(additional) 145 .map_err(|_| OutOfMemory::new(self.len().saturating_add(additional))) 146 } 147 148 /// Same as [`std::string::String::push`] but returns an error on allocation 149 /// failure. 150 #[inline] 151 pub fn push(&mut self, c: char) -> Result<(), OutOfMemory> { 152 self.reserve(c.len_utf8())?; 153 self.inner.push(c); 154 Ok(()) 155 } 156 157 /// Same as [`std::string::String::push_str`] but returns an error on 158 /// allocation failure. 159 #[inline] 160 pub fn push_str(&mut self, s: &str) -> Result<(), OutOfMemory> { 161 self.reserve(s.len())?; 162 self.inner.push_str(s); 163 Ok(()) 164 } 165 166 /// Same as [`std::string::String::into_raw_parts`]. 167 pub fn into_raw_parts(mut self) -> (*mut u8, usize, usize) { 168 // NB: Can't use `String::into_raw_parts` until our MSRV is >= 1.93. 169 #[cfg(not(miri))] 170 { 171 let ptr = self.as_mut_ptr(); 172 let len = self.len(); 173 let cap = self.capacity(); 174 mem::forget(self); 175 (ptr, len, cap) 176 } 177 // NB: Miri requires using `into_raw_parts`, but always run on nightly, 178 // so it's fine to use there. 179 #[cfg(miri)] 180 { 181 let _ = &mut self; 182 self.inner.into_raw_parts() 183 } 184 } 185 186 /// Same as [`std::string::String::from_raw_parts`]. 187 pub unsafe fn from_raw_parts(buf: *mut u8, length: usize, capacity: usize) -> Self { 188 Self { 189 // Safety: Same as our unsafe contract. 190 inner: unsafe { inner::String::from_raw_parts(buf, length, capacity) }, 191 } 192 } 193 194 /// Same as [`std::string::String::shrink_to_fit`] but returns an error on 195 /// allocation failure. 196 pub fn shrink_to_fit(&mut self) -> Result<(), OutOfMemory> { 197 // If our length is already equal to our capacity, then there is nothing 198 // to shrink. 199 if self.len() == self.capacity() { 200 return Ok(()); 201 } 202 203 // `realloc` requires a non-zero original layout as well as a non-zero 204 // destination layout, so this guard ensures that the sizes below are 205 // all nonzero. This handles a couple cases: 206 // 207 // * If `len == cap == 0` then no allocation has ever been made. 208 // * If `len == 0` and `cap != 0` then this function effectively frees 209 // the memory. 210 // 211 // In both of these cases delegate to the standard library's 212 // `shrink_to_fit` which is guaranteed to not perform a `realloc`. 213 if self.is_empty() { 214 self.inner.shrink_to_fit(); 215 return Ok(()); 216 } 217 218 let (ptr, len, cap) = mem::take(self).into_raw_parts(); 219 debug_assert!(!ptr.is_null()); 220 debug_assert!(len > 0); 221 debug_assert!(cap > len); 222 let old_layout = Layout::array::<u8>(cap).unwrap(); 223 debug_assert_eq!(old_layout.size(), cap); 224 let new_layout = Layout::array::<u8>(len).unwrap(); 225 debug_assert_eq!(old_layout.align(), new_layout.align()); 226 debug_assert_eq!(new_layout.size(), len); 227 228 // SAFETY: `ptr` was previously allocated in the global allocator, 229 // `layout` has a nonzero size and matches the current allocation of 230 // `ptr`, `len` is nonzero, and `len` is a valid array size 231 // for `len` elements given its constructor. 232 let result = unsafe { try_realloc(ptr, old_layout, len) }; 233 234 match result { 235 Ok(ptr) => { 236 // SAFETY: `result` is allocated with the global allocator and 237 // has room for exactly `[u8; len]`. 238 *self = unsafe { Self::from_raw_parts(ptr.as_ptr(), len, len) }; 239 Ok(()) 240 } 241 Err(oom) => { 242 // SAFETY: If reallocation fails then it's guaranteed that the 243 // original allocation is not tampered with, so it's safe to 244 // reassemble the original vector. 245 *self = unsafe { Self::from_raw_parts(ptr, len, cap) }; 246 Err(oom) 247 } 248 } 249 } 250 251 /// Same as [`std::string::String::into_boxed_str`] but returns an error on 252 /// allocation failure. 253 pub fn into_boxed_str(mut self) -> Result<Box<str>, OutOfMemory> { 254 self.shrink_to_fit()?; 255 256 let (ptr, len, cap) = self.into_raw_parts(); 257 debug_assert_eq!(len, cap); 258 let ptr = str_ptr_from_raw_parts(ptr, len); 259 260 // SAFETY: The `ptr` is allocated with the global allocator and points 261 // to a valid block of utf8. 262 let boxed = unsafe { Box::from_raw(ptr) }; 263 264 Ok(boxed) 265 } 266 } 267