user0/rust
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

constify from_fn, try_from_fn, try_map, map

Browse source
This commit is contained in:
bendn 2025-10-12 17:39:28 +07:00
parent eddf2f8c68
commit 1d718e20ac
No known key found for this signature in database
GPG key ID: 0D9D3A2A3B2A93D6
4 changed files with 138 additions and 99 deletions
Download patch file Download diff file Expand all files Collapse all files

134
library/core/src/array/drain.rs
View file

@ -1,76 +1,76 @@
use crate::iter::{TrustedLen, UncheckedIterator};
use crate::assert_unsafe_precondition;
use crate::marker::Destruct;
use crate::mem::ManuallyDrop;
use crate::ptr::drop_in_place;
use crate::slice;
/// A situationally-optimized version of `array.into_iter().for_each(func)`.
///
/// [`crate::array::IntoIter`]s are great when you need an owned iterator, but
/// storing the entire array *inside* the iterator like that can sometimes
/// pessimize code. Notable, it can be more bytes than you really want to move
/// around, and because the array accesses index into it SRoA has a harder time
/// optimizing away the type than it does iterators that just hold a couple pointers.
///
/// Thus this function exists, which gives a way to get *moved* access to the
/// elements of an array using a small iterator -- no bigger than a slice iterator.
///
/// The function-taking-a-closure structure makes it safe, as it keeps callers
/// from looking at already-dropped elements.
pub(crate) fn drain_array_with<T, R, const N: usize>(
array: [T; N],
func: impl for<'a> FnOnce(Drain<'a, T>) -> R,
) -> R {
let mut array = ManuallyDrop::new(array);
// SAFETY: Now that the local won't drop it, it's ok to construct the `Drain` which will.
let drain = Drain(array.iter_mut());
func(drain)
#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
#[unstable(feature = "array_try_map", issue = "79711")]
pub(super) struct Drain<'a, T, U, const N: usize, F: FnMut(T) -> U> {
array: ManuallyDrop<[T; N]>,
moved: usize,
f: &'a mut F,
}
#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
#[unstable(feature = "array_try_map", issue = "79711")]
impl<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Drain<'_, T, U, N, F>
where
F: [const] FnMut(T) -> U,
{
type Output = U;
/// See [`drain_array_with`] -- this is `pub(crate)` only so it's allowed to be
/// mentioned in the signature of that method. (Otherwise it hits `E0446`.)
// INVARIANT: It's ok to drop the remainder of the inner iterator.
pub(crate) struct Drain<'a, T>(slice::IterMut<'a, T>);
impl<T> Drop for Drain<'_, T> {
extern "rust-call" fn call_once(mut self, args: (usize,)) -> Self::Output {
self.call_mut(args)
}
}
#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
#[unstable(feature = "array_try_map", issue = "79711")]
impl<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Drain<'_, T, U, N, F>
where
F: [const] FnMut(T) -> U,
{
extern "rust-call" fn call_mut(&mut self, (i,): (usize,)) -> Self::Output {
// SAFETY: increment moved before moving. if `f` panics, we drop the rest.
self.moved += 1;
assert_unsafe_precondition!(
check_library_ub,
"musnt index array out of bounds", (i: usize = i, size: usize = N) => i < size
);
// SAFETY: the `i` should also always go up, and musnt skip any, else some things will be leaked.
// SAFETY: if it goes down, we will drop freed elements. not good.
// SAFETY: caller guarantees never called with number >= N (see `Drain::new`)
(self.f)(unsafe { self.array.as_ptr().add(i).read() })
}
}
#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
#[unstable(feature = "array_try_map", issue = "79711")]
impl<T: [const] Destruct, U, const N: usize, F: FnMut(T) -> U> const Drop
for Drain<'_, T, U, N, F>
{
fn drop(&mut self) {
// SAFETY: By the type invariant, we're allowed to drop all these.
unsafe { drop_in_place(self.0.as_mut_slice()) }
let mut n = self.moved;
while n != N {
// SAFETY: moved must always be < N
unsafe { self.array.as_mut_ptr().add(n).drop_in_place() };
n += 1;
}
}
}
impl<T> Iterator for Drain<'_, T> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<T> {
let p: *const T = self.0.next()?;
// SAFETY: The iterator was already advanced, so we won't drop this later.
Some(unsafe { p.read() })
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let n = self.len();
(n, Some(n))
}
}
impl<T> ExactSizeIterator for Drain<'_, T> {
#[inline]
fn len(&self) -> usize {
self.0.len()
}
}
// SAFETY: This is a 1:1 wrapper for a slice iterator, which is also `TrustedLen`.
unsafe impl<T> TrustedLen for Drain<'_, T> {}
impl<T> UncheckedIterator for Drain<'_, T> {
unsafe fn next_unchecked(&mut self) -> T {
// SAFETY: `Drain` is 1:1 with the inner iterator, so if the caller promised
// that there's an element left, the inner iterator has one too.
let p: *const T = unsafe { self.0.next_unchecked() };
// SAFETY: The iterator was already advanced, so we won't drop this later.
unsafe { p.read() }
impl<'a, T, U, const N: usize, F: FnMut(T) -> U> Drain<'a, T, U, N, F> {
/// This function returns a function that lets you index the given array in const.
/// As implemented it can optimize better than iterators, and can be constified.
/// It acts like a sort of guard and iterator combined, which can be implemented
/// as it is a struct that implements const fn;
/// in that regard it is somewhat similar to an array::Iter implementing `UncheckedIterator`.
/// The only method you're really allowed to call is `next()`,
/// anything else is more or less UB, hence this function being unsafe.
/// Moved elements will not be dropped.
///
/// Previously this was implemented as a wrapper around a `slice::Iter`, which
/// called `read()` on the returned `&T`; gnarly stuff.
///
/// SAFETY: must be called in order of 0..N, without indexing out of bounds. (see `Drain::call_mut`)
/// Potentially the function could completely disregard the supplied argument, however i think that behaviour would be unintuitive.
// FIXME(const-hack): this is a hack for `let guard = Guard(array); |i| f(guard[i])`.
pub(super) const unsafe fn new(array: [T; N], f: &'a mut F) -> Self {
Self { array: ManuallyDrop::new(array), moved: 0, f }
}
}

61
library/core/src/array/mod.rs
View file

@ -12,6 +12,7 @@ use crate::error::Error;
use crate::hash::{self, Hash};
use crate::intrinsics::transmute_unchecked;
use crate::iter::{UncheckedIterator, repeat_n};
use crate::marker::Destruct;
use crate::mem::{self, MaybeUninit};
use crate::ops::{
ChangeOutputType, ControlFlow, FromResidual, Index, IndexMut, NeverShortCircuit, Residual, Try,
@ -25,7 +26,6 @@ mod drain;
mod equality;
mod iter;
pub(crate) use drain::drain_array_with;
#[stable(feature = "array_value_iter", since = "1.51.0")]
pub use iter::IntoIter;
@ -105,9 +105,10 @@ pub fn repeat<T: Clone, const N: usize>(val: T) -> [T; N] {
/// ```
#[inline]
#[stable(feature = "array_from_fn", since = "1.63.0")]
pub fn from_fn<T, const N: usize, F>(f: F) -> [T; N]
#[rustc_const_unstable(feature = "const_array", issue = "147606")]
pub const fn from_fn<T: [const] Destruct, const N: usize, F>(f: F) -> [T; N]
where
F: FnMut(usize) -> T,
F: [const] FnMut(usize) -> T + [const] Destruct,
{
try_from_fn(NeverShortCircuit::wrap_mut_1(f)).0
}
@ -143,11 +144,15 @@ where
/// ```
#[inline]
#[unstable(feature = "array_try_from_fn", issue = "89379")]
pub fn try_from_fn<R, const N: usize, F>(cb: F) -> ChangeOutputType<R, [R::Output; N]>
#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
pub const fn try_from_fn<R, const N: usize>(
cb: impl [const] FnMut(usize) -> R + [const] Destruct,
) -> ChangeOutputType<R, [R::Output; N]>
where
F: FnMut(usize) -> R,
R: Try,
R::Residual: Residual<[R::Output; N]>,
R: [const] Try<
Residual: [const] Residual<[R::Output; N], TryType: [const] Try>,
Output: [const] Destruct,
>,
{
let mut array = [const { MaybeUninit::uninit() }; N];
match try_from_fn_erased(&mut array, cb) {
@ -549,9 +554,11 @@ impl<T, const N: usize> [T; N] {
/// ```
#[must_use]
#[stable(feature = "array_map", since = "1.55.0")]
pub fn map<F, U>(self, f: F) -> [U; N]
#[rustc_const_unstable(feature = "const_array", issue = "147606")]
pub const fn map<F, U>(self, f: F) -> [U; N]
where
F: FnMut(T) -> U,
F: [const] FnMut(T) -> U + [const] Destruct,
U: [const] Destruct,
{
self.try_map(NeverShortCircuit::wrap_mut_1(f)).0
}
@ -587,11 +594,22 @@ impl<T, const N: usize> [T; N] {
/// assert_eq!(c, Some(a));
/// ```
#[unstable(feature = "array_try_map", issue = "79711")]
pub fn try_map<R>(self, f: impl FnMut(T) -> R) -> ChangeOutputType<R, [R::Output; N]>
#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
pub const fn try_map<R>(
self,
mut f: impl [const] FnMut(T) -> R + [const] Destruct,
) -> ChangeOutputType<R, [R::Output; N]>
where
R: Try<Residual: Residual<[R::Output; N]>>,
R: [const] Try<
Residual: [const] Residual<[R::Output; N], TryType: [const] Try>,
Output: [const] Destruct,
>,
{
drain_array_with(self, |iter| try_from_trusted_iterator(iter.map(f)))
// SAFETY: try_from_fn calls `f` with 0..N.
let mut f = unsafe { drain::Drain::new(self, &mut f) };
let out = try_from_fn(&mut f);
mem::forget(f); // it doesnt like being remembered
out
}
/// Returns a slice containing the entire array. Equivalent to `&s[..]`.
@ -885,13 +903,11 @@ where
/// not optimizing away. So if you give it a shot, make sure to watch what
/// happens in the codegen tests.
#[inline]
fn try_from_fn_erased<T, R>(
buffer: &mut [MaybeUninit<T>],
mut generator: impl FnMut(usize) -> R,
) -> ControlFlow<R::Residual>
where
R: Try<Output = T>,
{
#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
const fn try_from_fn_erased<R: [const] Try<Output: [const] Destruct>>(
buffer: &mut [MaybeUninit<R::Output>],
mut generator: impl [const] FnMut(usize) -> R + [const] Destruct,
) -> ControlFlow<R::Residual> {
let mut guard = Guard { array_mut: buffer, initialized: 0 };
while guard.initialized < guard.array_mut.len() {
@ -930,7 +946,8 @@ impl<T> Guard<'_, T> {
///
/// No more than N elements must be initialized.
#[inline]
pub(crate) unsafe fn push_unchecked(&mut self, item: T) {
#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
pub(crate) const unsafe fn push_unchecked(&mut self, item: T) {
// SAFETY: If `initialized` was correct before and the caller does not
// invoke this method more than N times then writes will be in-bounds
// and slots will not be initialized more than once.
@ -941,11 +958,11 @@ impl<T> Guard<'_, T> {
}
}
impl<T> Drop for Guard<'_, T> {
#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
impl<T: [const] Destruct> const Drop for Guard<'_, T> {
#[inline]
fn drop(&mut self) {
debug_assert!(self.initialized <= self.array_mut.len());
// SAFETY: this slice will contain only initialized objects.
unsafe {
self.array_mut.get_unchecked_mut(..self.initialized).assume_init_drop();

40
library/core/src/ops/try_trait.rs
View file

@ -1,3 +1,4 @@
use crate::marker::{Destruct, PhantomData};
use crate::ops::ControlFlow;
/// The `?` operator and `try {}` blocks.
@ -396,6 +397,25 @@ pub(crate) type ChangeOutputType<T: Try<Residual: Residual<V>>, V> =
/// Not currently planned to be exposed publicly, so just `pub(crate)`.
#[repr(transparent)]
pub(crate) struct NeverShortCircuit<T>(pub T);
// FIXME(const-hack): replace with `|a| NeverShortCircuit(f(a))` when const closures added.
pub(crate) struct Wrapped<T, A, F: FnMut(A) -> T> {
f: F,
p: PhantomData<(T, A)>,
}
#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
impl<T, A, F: [const] FnMut(A) -> T + [const] Destruct> const FnOnce<(A,)> for Wrapped<T, A, F> {
type Output = NeverShortCircuit<T>;
extern "rust-call" fn call_once(mut self, args: (A,)) -> Self::Output {
self.call_mut(args)
}
}
#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
impl<T, A, F: [const] FnMut(A) -> T> const FnMut<(A,)> for Wrapped<T, A, F> {
extern "rust-call" fn call_mut(&mut self, (args,): (A,)) -> Self::Output {
NeverShortCircuit((self.f)(args))
}
}
impl<T> NeverShortCircuit<T> {
/// Wraps a unary function to produce one that wraps the output into a `NeverShortCircuit`.
@ -403,10 +423,11 @@ impl<T> NeverShortCircuit<T> {
/// This is useful for implementing infallible functions in terms of the `try_` ones,
/// without accidentally capturing extra generic parameters in a closure.
#[inline]
pub(crate) fn wrap_mut_1<A>(
mut f: impl FnMut(A) -> T,
) -> impl FnMut(A) -> NeverShortCircuit<T> {
move |a| NeverShortCircuit(f(a))
pub(crate) const fn wrap_mut_1<A, F>(f: F) -> Wrapped<T, A, F>
where
F: [const] FnMut(A) -> T,
{
Wrapped { f, p: PhantomData }
}
#[inline]
@ -417,7 +438,8 @@ impl<T> NeverShortCircuit<T> {
pub(crate) enum NeverShortCircuitResidual {}
impl<T> Try for NeverShortCircuit<T> {
#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
impl<T> const Try for NeverShortCircuit<T> {
type Output = T;
type Residual = NeverShortCircuitResidual;
@ -431,15 +453,15 @@ impl<T> Try for NeverShortCircuit<T> {
NeverShortCircuit(x)
}
}
impl<T> FromResidual for NeverShortCircuit<T> {
#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
impl<T> const FromResidual for NeverShortCircuit<T> {
#[inline]
fn from_residual(never: NeverShortCircuitResidual) -> Self {
match never {}
}
}
impl<T> Residual<T> for NeverShortCircuitResidual {
#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
impl<T: [const] Destruct> const Residual<T> for NeverShortCircuitResidual {
type TryType = NeverShortCircuit<T>;
}

2
library/coretests/tests/lib.rs
View file

@ -17,8 +17,8 @@
#![feature(char_internals)]
#![feature(char_max_len)]
#![feature(clone_to_uninit)]
#![feature(const_cell_traits)]
#![feature(const_array)]
#![feature(const_cell_traits)]
#![feature(const_cmp)]
#![feature(const_convert)]
#![feature(const_destruct)]