wasmtime_internal_core/alloc/
string.rs

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

wasmtime_internal_core/alloc/
string.rs
