Skip to content

Commit 1e8d17b

Browse files
committed
rust: alloc: add some try_* methods we need
In preparation for enabling `no_global_oom_handling` for `alloc`, we need to add some new methods. They are all marked as: #[stable(feature = "kernel", since = "1.0.0")] for easy identification. Signed-off-by: Miguel Ojeda <[email protected]>
1 parent 8db4a3d commit 1e8d17b

File tree

4 files changed

+373
-3
lines changed

4 files changed

+373
-3
lines changed

rust/alloc/raw_vec.rs

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use crate::collections::TryReserveError::{self, *};
2020
#[cfg(test)]
2121
mod tests;
2222

23-
#[cfg(not(no_global_oom_handling))]
23+
#[allow(dead_code)]
2424
enum AllocInit {
2525
/// The contents of the new memory are uninitialized.
2626
Uninitialized,
@@ -93,6 +93,16 @@ impl<T> RawVec<T, Global> {
9393
Self::with_capacity_in(capacity, Global)
9494
}
9595

96+
/// Tries to create a `RawVec` (on the system heap) with exactly the
97+
/// capacity and alignment requirements for a `[T; capacity]`. This is
98+
/// equivalent to calling `RawVec::new` when `capacity` is `0` or `T` is
99+
/// zero-sized. Note that if `T` is zero-sized this means you will
100+
/// *not* get a `RawVec` with the requested capacity.
101+
#[inline]
102+
pub fn try_with_capacity(capacity: usize) -> Result<Self, TryReserveError> {
103+
Self::try_with_capacity_in(capacity, Global)
104+
}
105+
96106
/// Like `with_capacity`, but guarantees the buffer is zeroed.
97107
#[cfg(not(no_global_oom_handling))]
98108
#[inline]
@@ -144,6 +154,13 @@ impl<T, A: Allocator> RawVec<T, A> {
144154
Self::allocate_in(capacity, AllocInit::Uninitialized, alloc)
145155
}
146156

157+
/// Like `try_with_capacity`, but parameterized over the choice of
158+
/// allocator for the returned `RawVec`.
159+
#[inline]
160+
pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> {
161+
Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc)
162+
}
163+
147164
/// Like `with_capacity_zeroed`, but parameterized over the choice
148165
/// of allocator for the returned `RawVec`.
149166
#[cfg(not(no_global_oom_handling))]
@@ -218,6 +235,29 @@ impl<T, A: Allocator> RawVec<T, A> {
218235
}
219236
}
220237

238+
fn try_allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Result<Self, TryReserveError> {
239+
if mem::size_of::<T>() == 0 {
240+
return Ok(Self::new_in(alloc));
241+
}
242+
243+
let layout = Layout::array::<T>(capacity)?;
244+
alloc_guard(layout.size())?;
245+
let result = match init {
246+
AllocInit::Uninitialized => alloc.allocate(layout),
247+
AllocInit::Zeroed => alloc.allocate_zeroed(layout),
248+
};
249+
let ptr = match result {
250+
Ok(ptr) => ptr,
251+
Err(_) => return Err(TryReserveError::AllocError{ layout, non_exhaustive: () }),
252+
};
253+
254+
Ok(Self {
255+
ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) },
256+
cap: Self::capacity_from_bytes(ptr.len()),
257+
alloc,
258+
})
259+
}
260+
221261
/// Reconstitutes a `RawVec` from a pointer, capacity, and allocator.
222262
///
223263
/// # Safety
@@ -394,6 +434,16 @@ impl<T, A: Allocator> RawVec<T, A> {
394434
pub fn shrink_to_fit(&mut self, amount: usize) {
395435
handle_reserve(self.shrink(amount));
396436
}
437+
438+
/// Tries to shrink the allocation down to the specified amount. If the given amount
439+
/// is 0, actually completely deallocates.
440+
///
441+
/// # Panics
442+
///
443+
/// Panics if the given amount is *larger* than the current capacity.
444+
pub fn try_shrink_to_fit(&mut self, amount: usize) -> Result<(), TryReserveError> {
445+
self.shrink(amount)
446+
}
397447
}
398448

399449
impl<T, A: Allocator> RawVec<T, A> {
@@ -465,7 +515,7 @@ impl<T, A: Allocator> RawVec<T, A> {
465515
Ok(())
466516
}
467517

468-
#[cfg(not(no_global_oom_handling))]
518+
// TODO: Why this was marked as `#[cfg(not(no_global_oom_handling))]`?
469519
fn shrink(&mut self, amount: usize) -> Result<(), TryReserveError> {
470520
assert!(amount <= self.capacity(), "Tried to shrink to a larger capacity");
471521

rust/alloc/slice.rs

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,11 +95,11 @@ use core::mem::size_of;
9595
use core::ptr;
9696

9797
use crate::alloc::Allocator;
98-
#[cfg(not(no_global_oom_handling))]
9998
use crate::alloc::Global;
10099
#[cfg(not(no_global_oom_handling))]
101100
use crate::borrow::ToOwned;
102101
use crate::boxed::Box;
102+
use crate::collections::TryReserveError;
103103
use crate::vec::Vec;
104104

105105
#[unstable(feature = "slice_range", issue = "76393")]
@@ -472,6 +472,24 @@ impl<T> [T] {
472472
self.to_vec_in(Global)
473473
}
474474

475+
/// Tries to copy `self` into a new `Vec`.
476+
///
477+
/// # Examples
478+
///
479+
/// ```
480+
/// let s = [10, 40, 30];
481+
/// let x = s.try_to_vec().unwrap();
482+
/// // Here, `s` and `x` can be modified independently.
483+
/// ```
484+
#[inline]
485+
#[stable(feature = "kernel", since = "1.0.0")]
486+
pub fn try_to_vec(&self) -> Result<Vec<T>, TryReserveError>
487+
where
488+
T: Clone,
489+
{
490+
self.try_to_vec_in(Global)
491+
}
492+
475493
/// Copies `self` into a new `Vec` with an allocator.
476494
///
477495
/// # Examples
@@ -496,6 +514,36 @@ impl<T> [T] {
496514
hack::to_vec(self, alloc)
497515
}
498516

517+
/// Tries to copy `self` into a new `Vec` with an allocator.
518+
///
519+
/// # Examples
520+
///
521+
/// ```
522+
/// #![feature(allocator_api)]
523+
///
524+
/// use std::alloc::System;
525+
///
526+
/// let s = [10, 40, 30];
527+
/// let x = s.try_to_vec_in(System).unwrap();
528+
/// // Here, `s` and `x` can be modified independently.
529+
/// ```
530+
#[inline]
531+
#[stable(feature = "kernel", since = "1.0.0")]
532+
pub fn try_to_vec_in<A: Allocator>(&self, alloc: A) -> Result<Vec<T, A>, TryReserveError>
533+
where
534+
T: Clone,
535+
{
536+
let mut v = Vec::try_with_capacity_in(self.len(), alloc)?;
537+
// SAFETY:
538+
// allocated above with the capacity of `self`, and initialize to `self.len()` in
539+
// ptr::copy_to_non_overlapping below.
540+
unsafe {
541+
self.as_ptr().copy_to_nonoverlapping(v.as_mut_ptr(), self.len());
542+
v.set_len(self.len());
543+
}
544+
Ok(v)
545+
}
546+
499547
/// Converts `self` into a vector without clones or allocation.
500548
///
501549
/// The resulting vector can be converted back into a box via

rust/alloc/str.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ use core::unicode::conversions;
3838

3939
use crate::borrow::ToOwned;
4040
use crate::boxed::Box;
41+
use crate::collections::TryReserveError;
4142
use crate::slice::{Concat, Join, SliceIndex};
4243
use crate::string::String;
4344
use crate::vec::Vec;
@@ -575,6 +576,22 @@ impl str {
575576
// make_ascii_lowercase() preserves the UTF-8 invariant.
576577
unsafe { String::from_utf8_unchecked(bytes) }
577578
}
579+
580+
/// Tries to create a `String`.
581+
///
582+
/// # Examples
583+
///
584+
/// Basic usage:
585+
///
586+
/// ```
587+
/// let s: &str = "a";
588+
/// let ss: String = s.try_to_owned().unwrap();
589+
/// ```
590+
#[inline]
591+
#[stable(feature = "kernel", since = "1.0.0")]
592+
pub fn try_to_owned(&self) -> Result<String, TryReserveError> {
593+
unsafe { Ok(String::from_utf8_unchecked(self.as_bytes().try_to_vec()?)) }
594+
}
578595
}
579596

580597
/// Converts a boxed slice of bytes to a boxed string slice without checking

0 commit comments

Comments
 (0)