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