Commit 5ab560ce authored by Wedson Almeida Filho's avatar Wedson Almeida Filho Committed by Miguel Ojeda
Browse files

rust: alloc: update `VecExt` to take allocation flags 



We also rename the methods by removing the `try_` prefix since the names
are available due to our usage of the `no_global_oom_handling` config
when building the `alloc` crate.

Reviewed-by: default avatarBoqun Feng <boqun.feng@gmail.com>
Signed-off-by: default avatarWedson Almeida Filho <walmeida@microsoft.com>
Reviewed-by: default avatarBenno Lossin <benno.lossin@proton.me>
Link: https://lore.kernel.org/r/20240328013603.206764-8-wedsonaf@gmail.com


Signed-off-by: default avatarMiguel Ojeda <ojeda@kernel.org>
parent 08d3f549

rust/kernel/alloc/vec_ext.rs

+143 −15
Original line number Diff line number Diff line
@@ -2,47 +2,175 @@ 


//! Extensions to [`Vec`] for fallible allocations.



use alloc::{collections::TryReserveError, vec::Vec};

use super::Flags;

use alloc::{alloc::AllocError, vec::Vec};

use core::result::Result;



/// Extensions to [`Vec`].

pub trait VecExt<T>: Sized {

    /// Creates a new [`Vec`] instance with at least the given capacity.

    fn try_with_capacity(capacity: usize) -> Result<Self, TryReserveError>;

    ///

    /// # Examples

    ///

    /// ```

    /// let v = Vec::<u32>::with_capacity(20, GFP_KERNEL)?;

    ///

    /// assert!(v.capacity() >= 20);

    /// # Ok::<(), Error>(())

    /// ```

    fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError>;



    /// Appends an element to the back of the [`Vec`] instance.

    fn try_push(&mut self, v: T) -> Result<(), TryReserveError>;

    ///

    /// # Examples

    ///

    /// ```

    /// let mut v = Vec::new();

    /// v.push(1, GFP_KERNEL)?;

    /// assert_eq!(&v, &[1]);

    ///

    /// v.push(2, GFP_KERNEL)?;

    /// assert_eq!(&v, &[1, 2]);

    /// # Ok::<(), Error>(())

    /// ```

    fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError>;



    /// Pushes clones of the elements of slice into the [`Vec`] instance.

    fn try_extend_from_slice(&mut self, other: &[T]) -> Result<(), TryReserveError>

    ///

    /// # Examples

    ///

    /// ```

    /// let mut v = Vec::new();

    /// v.push(1, GFP_KERNEL)?;

    ///

    /// v.extend_from_slice(&[20, 30, 40], GFP_KERNEL)?;

    /// assert_eq!(&v, &[1, 20, 30, 40]);

    ///

    /// v.extend_from_slice(&[50, 60], GFP_KERNEL)?;

    /// assert_eq!(&v, &[1, 20, 30, 40, 50, 60]);

    /// # Ok::<(), Error>(())

    /// ```

    fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError>

    where

        T: Clone;



    /// Ensures that the capacity exceeds the length by at least `additional` elements.

    ///

    /// # Examples

    ///

    /// ```

    /// let mut v = Vec::new();

    /// v.push(1, GFP_KERNEL)?;

    ///

    /// v.reserve(10, GFP_KERNEL)?;

    /// let cap = v.capacity();

    /// assert!(cap >= 10);

    ///

    /// v.reserve(10, GFP_KERNEL)?;

    /// let new_cap = v.capacity();

    /// assert_eq!(new_cap, cap);

    ///

    /// # Ok::<(), Error>(())

    /// ```

    fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError>;

}



impl<T> VecExt<T> for Vec<T> {

    fn try_with_capacity(capacity: usize) -> Result<Self, TryReserveError> {

    fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError> {

        let mut v = Vec::new();

        v.try_reserve(capacity)?;

        <Self as VecExt<_>>::reserve(&mut v, capacity, flags)?;

        Ok(v)

    }



    fn try_push(&mut self, v: T) -> Result<(), TryReserveError> {

        if let Err(retry) = self.push_within_capacity(v) {

            self.try_reserve(1)?;

            let _ = self.push_within_capacity(retry);

        }

    fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {

        <Self as VecExt<_>>::reserve(self, 1, flags)?;

        let s = self.spare_capacity_mut();

        s[0].write(v);



        // SAFETY: We just initialised the first spare entry, so it is safe to increase the length

        // by 1. We also know that the new length is <= capacity because of the previous call to

        // `reserve` above.

        unsafe { self.set_len(self.len() + 1) };

        Ok(())

    }



    fn try_extend_from_slice(&mut self, other: &[T]) -> Result<(), TryReserveError>

    fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError>

    where

        T: Clone,

    {

        self.try_reserve(other.len())?;

        for item in other {

            self.try_push(item.clone())?;

        <Self as VecExt<_>>::reserve(self, other.len(), flags)?;

        for (slot, item) in core::iter::zip(self.spare_capacity_mut(), other) {

            slot.write(item.clone());

        }



        // SAFETY: We just initialised the `other.len()` spare entries, so it is safe to increase

        // the length by the same amount. We also know that the new length is <= capacity because

        // of the previous call to `reserve` above.

        unsafe { self.set_len(self.len() + other.len()) };

        Ok(())

    }



    #[cfg(any(test, testlib))]

    fn reserve(&mut self, additional: usize, _flags: Flags) -> Result<(), AllocError> {

        Vec::reserve(self, additional);

        Ok(())

    }



    #[cfg(not(any(test, testlib)))]

    fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError> {

        let len = self.len();

        let cap = self.capacity();



        if cap - len >= additional {

            return Ok(());

        }



        if core::mem::size_of::<T>() == 0 {

            // The capacity is already `usize::MAX` for SZTs, we can't go higher.

            return Err(AllocError);

        }



        // We know cap is <= `isize::MAX` because `Layout::array` fails if the resulting byte size

        // is greater than `isize::MAX`. So the multiplication by two won't overflow.

        let new_cap = core::cmp::max(cap * 2, len.checked_add(additional).ok_or(AllocError)?);

        let layout = core::alloc::Layout::array::<T>(new_cap).map_err(|_| AllocError)?;



        let (ptr, len, cap) = destructure(self);



        // SAFETY: `ptr` is valid because it's either NULL or comes from a previous call to

        // `krealloc_aligned`. We also verified that the type is not a ZST.

        let new_ptr = unsafe { super::allocator::krealloc_aligned(ptr.cast(), layout, flags) };

        if new_ptr.is_null() {

            // SAFETY: We are just rebuilding the existing `Vec` with no changes.

            unsafe { rebuild(self, ptr, len, cap) };

            Err(AllocError)

        } else {

            // SAFETY: `ptr` has been reallocated with the layout for `new_cap` elements. New cap

            // is greater than `cap`, so it continues to be >= `len`.

            unsafe { rebuild(self, new_ptr.cast::<T>(), len, new_cap) };

            Ok(())

        }

    }

}



#[cfg(not(any(test, testlib)))]

fn destructure<T>(v: &mut Vec<T>) -> (*mut T, usize, usize) {

    let mut tmp = Vec::new();

    core::mem::swap(&mut tmp, v);

    let mut tmp = core::mem::ManuallyDrop::new(tmp);

    let len = tmp.len();

    let cap = tmp.capacity();

    (tmp.as_mut_ptr(), len, cap)

}



/// Rebuilds a `Vec` from a pointer, length, and capacity.

///

/// # Safety

///

/// The same as [`Vec::from_raw_parts`].

#[cfg(not(any(test, testlib)))]

unsafe fn rebuild<T>(v: &mut Vec<T>, ptr: *mut T, len: usize, cap: usize) {

    // SAFETY: The safety requirements from this function satisfy those of `from_raw_parts`.

    let mut tmp = unsafe { Vec::from_raw_parts(ptr, len, cap) };

    core::mem::swap(&mut tmp, v);

}

rust/kernel/error.rs

+1 −10
Original line number Diff line number Diff line
@@ -6,10 +6,7 @@ 


use crate::str::CStr;



use alloc::{

    alloc::{AllocError, LayoutError},

    collections::TryReserveError,

};

use alloc::alloc::{AllocError, LayoutError};



use core::convert::From;

use core::fmt;
@@ -192,12 +189,6 @@ fn from(_: Utf8Error) -> Error { 
    }

}



impl From<TryReserveError> for Error {

    fn from(_: TryReserveError) -> Error {

        code::ENOMEM

    }

}



impl From<LayoutError> for Error {

    fn from(_: LayoutError) -> Error {

        code::ENOMEM

rust/kernel/lib.rs

+0 −1
Original line number Diff line number Diff line
@@ -18,7 +18,6 @@ 
#![feature(new_uninit)]

#![feature(receiver_trait)]

#![feature(unsize)]

#![feature(vec_push_within_capacity)]



// Ensure conditional compilation based on the kernel configuration works;

// otherwise we may silently break things like initcall handling.

rust/kernel/str.rs

+3 −3
Original line number Diff line number Diff line
@@ -2,7 +2,7 @@ 


//! String representations.



use crate::alloc::vec_ext::VecExt;

use crate::alloc::{flags::*, vec_ext::VecExt};

use alloc::alloc::AllocError;

use alloc::vec::Vec;

use core::fmt::{self, Write};
@@ -807,7 +807,7 @@ pub fn try_from_fmt(args: fmt::Arguments<'_>) -> Result<Self, Error> { 
        let size = f.bytes_written();



        // Allocate a vector with the required number of bytes, and write to it.

        let mut buf = Vec::try_with_capacity(size)?;

        let mut buf = <Vec<_> as VecExt<_>>::with_capacity(size, GFP_KERNEL)?;

        // SAFETY: The buffer stored in `buf` is at least of size `size` and is valid for writes.

        let mut f = unsafe { Formatter::from_buffer(buf.as_mut_ptr(), size) };

        f.write_fmt(args)?;
@@ -856,7 +856,7 @@ impl<'a> TryFrom<&'a CStr> for CString { 
    fn try_from(cstr: &'a CStr) -> Result<CString, AllocError> {

        let mut buf = Vec::new();



        buf.try_extend_from_slice(cstr.as_bytes_with_nul())

        <Vec<_> as VecExt<_>>::extend_from_slice(&mut buf, cstr.as_bytes_with_nul(), GFP_KERNEL)

            .map_err(|_| AllocError)?;



        // INVARIANT: The `CStr` and `CString` types have the same invariants for

rust/kernel/types.rs

+2 −2
Original line number Diff line number Diff line
@@ -157,11 +157,11 @@ unsafe fn from_foreign(_: *const core::ffi::c_void) -> Self {} 
///     let mut vec =

///         ScopeGuard::new_with_data(Vec::new(), |v| pr_info!("vec had {} elements\n", v.len()));

///

///     vec.try_push(10u8)?;

///     vec.push(10u8, GFP_KERNEL)?;

///     if arg {

///         return Ok(());

///     }

///     vec.try_push(20u8)?;

///     vec.push(20u8, GFP_KERNEL)?;

///     Ok(())

/// }

///