From 6b0a5e92b5e471dc786871c418854cba5574871e Mon Sep 17 00:00:00 2001 From: Gary Guo Date: Tue, 1 Jun 2021 03:37:01 +0100 Subject: [PATCH] rust: add fallible allocation methods These are all fallible alternatives to all currently used methods that may perform infallible allocation. `assume_fallible` is a semantically marker function to denote that we have proved that the closure passed to it will not perform infallible allocation. Signed-off-by: Gary Guo --- rust/kernel/io_buffer.rs | 6 +- rust/kernel/module_param.rs | 5 +- rust/kernel/prelude.rs | 2 +- rust/kernel/traits.rs | 113 +++++++++++++++++++++++++++++++++++ samples/rust/rust_minimal.rs | 2 +- 5 files changed, 120 insertions(+), 8 deletions(-) diff --git a/rust/kernel/io_buffer.rs b/rust/kernel/io_buffer.rs index 64bf21987cb55f..0a306bd6ac8803 100644 --- a/rust/kernel/io_buffer.rs +++ b/rust/kernel/io_buffer.rs @@ -2,6 +2,7 @@ //! Buffers used in IO. +use crate::traits::VecExt; use crate::Result; use alloc::vec::Vec; use core::mem::{size_of, MaybeUninit}; @@ -29,9 +30,8 @@ pub trait IoBufferReader { /// /// Returns `EFAULT` if the address does not currently point to mapped, readable memory. fn read_all(&mut self) -> Result> { - let mut data = Vec::::new(); - data.try_reserve_exact(self.len())?; - data.resize(self.len(), 0); + let mut data = Vec::try_with_capacity(self.len())?; + data.try_resize(self.len(), 0)?; // SAFETY: The output buffer is valid as we just allocated it. unsafe { self.read_raw(data.as_mut_ptr(), data.len())? }; diff --git a/rust/kernel/module_param.rs b/rust/kernel/module_param.rs index fc6a6b01c588db..da2234105fb6c1 100644 --- a/rust/kernel/module_param.rs +++ b/rust/kernel/module_param.rs @@ -5,6 +5,7 @@ //! C header: [`include/linux/moduleparam.h`](../../../include/linux/moduleparam.h) use crate::str::CStr; +use crate::traits::TryToOwned; use core::fmt::Write; /// Types that can be used for module parameters. @@ -475,9 +476,7 @@ impl ModuleParam for StringParam { let slab_available = unsafe { crate::bindings::slab_is_available() }; arg.and_then(|arg| { if slab_available { - let mut vec = alloc::vec::Vec::new(); - vec.try_reserve_exact(arg.len()).ok()?; - vec.extend_from_slice(arg); + let vec = arg.try_to_owned().ok()?; Some(StringParam::Owned(vec)) } else { Some(StringParam::Ref(arg)) diff --git a/rust/kernel/prelude.rs b/rust/kernel/prelude.rs index 0322b6a2b91f50..91c8f54f2a4aa2 100644 --- a/rust/kernel/prelude.rs +++ b/rust/kernel/prelude.rs @@ -23,4 +23,4 @@ pub use super::static_assert; pub use super::{KernelModule, Result}; -pub use crate::traits::TryPin; +pub use crate::traits::{TryPin, TryToOwned, VecExt as _}; diff --git a/rust/kernel/traits.rs b/rust/kernel/traits.rs index 39a43169bf70e6..e6cd31b990a792 100644 --- a/rust/kernel/traits.rs +++ b/rust/kernel/traits.rs @@ -4,8 +4,19 @@ use core::{ops::Deref, pin::Pin}; +use alloc::collections::TryReserveError; +use alloc::string::String; +use alloc::vec::Vec; use alloc::{alloc::AllocError, sync::Arc}; +type AllocResult = core::result::Result; +type CollectionResult = core::result::Result; + +#[inline] +fn assume_fallible T>(f: F) -> T { + f() +} + /// Trait which provides a fallible version of `pin()` for pointer types. /// /// Common pointer types which implement a `pin()` method include [`Box`](alloc::boxed::Box) and [`Arc`]. @@ -24,3 +35,105 @@ impl TryPin> for Arc { Ok(unsafe { Pin::new_unchecked(Arc::try_new(data)?) }) } } + +/// Faillible alternative to [`alloc::borrow::ToOwned`]. +pub trait TryToOwned { + /// The resulting type after obtaining ownership. + type Owned: core::borrow::Borrow; + + /// Faillible alternative to [`alloc::borrow::ToOwned::to_owned`]. + #[must_use = "cloning is often expensive and is not expected to have side effects"] + fn try_to_owned(&self) -> AllocResult; + + /// Faillible alternative to [`alloc::borrow::ToOwned::clone_into`]. + fn try_clone_into(&self, target: &mut Self::Owned) -> AllocResult { + *target = self.try_to_owned()?; + Ok(()) + } +} + +impl TryToOwned for [T] { + type Owned = Vec; + + fn try_to_owned(&self) -> AllocResult> { + let mut vec = Vec::new(); + self.try_clone_into(&mut vec)?; + Ok(vec) + } + + fn try_clone_into(&self, target: &mut Vec) -> AllocResult { + // Ensure target has enough capacity + target + .try_reserve_exact(self.len().saturating_sub(target.len())) + .map_err(|_| AllocError)?; + + target.clear(); + assume_fallible(|| target.extend_from_slice(self)); + Ok(()) + } +} + +impl TryToOwned for str { + type Owned = String; + + fn try_to_owned(&self) -> AllocResult { + let mut vec = String::new(); + self.try_clone_into(&mut vec)?; + Ok(vec) + } + + fn try_clone_into(&self, target: &mut String) -> AllocResult { + // Ensure target has enough capacity + target + .try_reserve_exact(self.len().saturating_sub(target.len())) + .map_err(|_| AllocError)?; + + target.clear(); + assume_fallible(|| target.push_str(self)); + Ok(()) + } +} + +/// Trait which provides a fallible methods for [`Vec`]. +pub trait VecExt { + /// Faillible alternative to [`Vec::with_capacity`]. + fn try_with_capacity(capacity: usize) -> CollectionResult + where + Self: Sized; + + /// Faillible alternative to [`Vec::extend_from_slice`]. + fn try_extend_from_slice(&mut self, other: &[T]) -> CollectionResult + where + T: Clone; + + /// Faillible alternative to [`Vec::resize`]. + fn try_resize(&mut self, new_len: usize, value: T) -> CollectionResult + where + T: Clone; +} + +impl VecExt for alloc::vec::Vec { + fn try_with_capacity(capacity: usize) -> CollectionResult { + let mut vec = Self::new(); + vec.try_reserve_exact(capacity)?; + Ok(vec) + } + + fn try_extend_from_slice(&mut self, other: &[T]) -> CollectionResult + where + T: Clone, + { + self.try_reserve(other.len())?; + assume_fallible(|| self.extend_from_slice(other)); + Ok(()) + } + + fn try_resize(&mut self, new_len: usize, value: T) -> CollectionResult + where + T: Clone, + { + self.try_reserve(new_len.saturating_sub(self.len()))?; + assume_fallible(|| self.resize(new_len, value)); + Ok(()) + } +} diff --git a/samples/rust/rust_minimal.rs b/samples/rust/rust_minimal.rs index 84a227021aafdd..49cfd8cf3aad69 100644 --- a/samples/rust/rust_minimal.rs +++ b/samples/rust/rust_minimal.rs @@ -25,7 +25,7 @@ impl KernelModule for RustMinimal { pr_info!("Am I built-in? {}\n", !cfg!(MODULE)); Ok(RustMinimal { - message: "on the heap!".to_owned(), + message: "on the heap!".try_to_owned()?, }) } }