diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index a7cb1023229bb..d0bb59df428fe 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -425,42 +425,62 @@ impl [T; N] { /// assert_eq!(y, [6, 9, 3, 3]); /// ``` #[unstable(feature = "array_map", issue = "75243")] - pub fn map(self, mut f: F) -> [U; N] + pub fn map(self, f: F) -> [U; N] where F: FnMut(T) -> U, { - use crate::mem::MaybeUninit; - struct Guard { - dst: *mut T, - initialized: usize, - } + use crate::mem::{forget, ManuallyDrop, MaybeUninit}; + use crate::ptr; - impl Drop for Guard { + union MaybeUninitArray { + none: (), + partial: ManuallyDrop<[MaybeUninit; N]>, + complete: ManuallyDrop<[T; N]>, + } + struct MapGuard<'a, T, const N: usize> { + arr: &'a mut MaybeUninitArray, + len: usize, + } + impl<'a, T, const N: usize> MapGuard<'a, T, N> { + fn push(&mut self, value: T) { + // SAFETY: Since we know the input size is N, and the output is N, + // this will never exceed the capacity, and MaybeUninit is always in the + // structure of an array. + unsafe { + self.arr.partial[self.len].write(value); + self.len += 1; + } + } + } + impl<'a, T, const N: usize> Drop for MapGuard<'a, T, N> { fn drop(&mut self) { - debug_assert!(self.initialized <= N); - - let initialized_part = - crate::ptr::slice_from_raw_parts_mut(self.dst, self.initialized); - // SAFETY: this raw slice will contain only initialized objects - // that's why, it is allowed to drop it. + //debug_assert!(self.len <= N); + // SAFETY: already pushed `len` elements, but need to drop them now that + // `f` panicked. unsafe { - crate::ptr::drop_in_place(initialized_part); + let p: *mut MaybeUninit = self.arr.partial.as_mut_ptr(); + let slice: *mut [T] = ptr::slice_from_raw_parts_mut(p.cast(), self.len); + ptr::drop_in_place(slice); } } } - let mut dst = MaybeUninit::uninit_array::(); - let mut guard: Guard = - Guard { dst: MaybeUninit::slice_as_mut_ptr(&mut dst), initialized: 0 }; - for (src, dst) in IntoIter::new(self).zip(&mut dst) { - dst.write(f(src)); - guard.initialized += 1; + + fn map_guard_fn( + buffer: &mut MaybeUninitArray, + iter: impl Iterator, + ) { + let mut guard = MapGuard { arr: buffer, len: 0 }; + for v in iter { + guard.push(v); + } + forget(guard); } - // FIXME: Convert to crate::mem::transmute once it works with generics. - // unsafe { crate::mem::transmute::<[MaybeUninit; N], [U; N]>(dst) } - crate::mem::forget(guard); - // SAFETY: At this point we've properly initialized the whole array - // and we just need to cast it to the correct type. - unsafe { crate::mem::transmute_copy::<_, [U; N]>(&dst) } + + let mut buffer = MaybeUninitArray:: { none: () }; + map_guard_fn(&mut buffer, IntoIter::new(self).map(f)); + // SAFETY: all elements have successfully initialized, don't run guard's drop code + // and take completed buffer out of MaybeUninitArray. + unsafe { ManuallyDrop::into_inner(buffer.complete) } } /// Returns a slice containing the entire array. Equivalent to `&s[..]`. diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs index 89c2a969c28bb..497d734f9880b 100644 --- a/library/core/tests/array.rs +++ b/library/core/tests/array.rs @@ -323,26 +323,53 @@ fn array_map() { fn array_map_drop_safety() { use core::sync::atomic::AtomicUsize; use core::sync::atomic::Ordering; - static DROPPED: AtomicUsize = AtomicUsize::new(0); - struct DropCounter; - impl Drop for DropCounter { + static DROPPED: [AtomicUsize; 3] = + [AtomicUsize::new(0), AtomicUsize::new(0), AtomicUsize::new(0)]; + struct DropCounter; + impl Drop for DropCounter { fn drop(&mut self) { - DROPPED.fetch_add(1, Ordering::SeqCst); + DROPPED[N].fetch_add(1, Ordering::SeqCst); } } - let num_to_create = 5; - let success = std::panic::catch_unwind(|| { - let items = [0; 10]; - let mut nth = 0; - items.map(|_| { - assert!(nth < num_to_create); - nth += 1; - DropCounter + { + let num_to_create = 5; + let success = std::panic::catch_unwind(|| { + let items = [0; 10]; + let mut nth = 0; + items.map(|_| { + assert!(nth < num_to_create); + nth += 1; + DropCounter::<0> + }); + }); + assert!(success.is_err()); + assert_eq!(DROPPED[0].load(Ordering::SeqCst), num_to_create); + } + + { + assert_eq!(DROPPED[1].load(Ordering::SeqCst), 0); + let num_to_create = 3; + const TOTAL: usize = 5; + let success = std::panic::catch_unwind(|| { + let items: [DropCounter<1>; TOTAL] = [ + DropCounter::<1>, + DropCounter::<1>, + DropCounter::<1>, + DropCounter::<1>, + DropCounter::<1>, + ]; + let mut nth = 0; + items.map(|_| { + assert!(nth < num_to_create); + nth += 1; + DropCounter::<2> + }); }); - }); - assert!(success.is_err()); - assert_eq!(DROPPED.load(Ordering::SeqCst), num_to_create); + assert!(success.is_err()); + assert_eq!(DROPPED[2].load(Ordering::SeqCst), num_to_create); + assert_eq!(DROPPED[1].load(Ordering::SeqCst), TOTAL); + } panic!("test succeeded") } diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 106c9fe5da3e6..cfdd7f810376a 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -32,6 +32,7 @@ #![feature(raw)] #![feature(sort_internals)] #![feature(slice_partition_at_index)] +#![feature(min_const_generics)] #![feature(min_specialization)] #![feature(step_trait)] #![feature(step_trait_ext)] diff --git a/src/test/codegen/array-map.rs b/src/test/codegen/array-map.rs new file mode 100644 index 0000000000000..f4409cb5d9ca3 --- /dev/null +++ b/src/test/codegen/array-map.rs @@ -0,0 +1,25 @@ +// compile-flags: -C opt-level=3 -Zmir-opt-level=3 +#![crate_type = "lib"] +#![feature(array_map)] + +const SIZE: usize = 4; + +// CHECK-LABEL: @array_cast_to_float +#[no_mangle] +pub fn array_cast_to_float(x: [u32; SIZE]) -> [f32; SIZE] { + // CHECK: cast + // CHECK: @llvm.memcpy + // CHECK: ret + // CHECK-EMPTY + x.map(|v| v as f32) +} + +// CHECK-LABEL: @array_cast_to_u64 +#[no_mangle] +pub fn array_cast_to_u64(x: [u32; SIZE]) -> [u64; SIZE] { + // CHECK: cast + // CHECK: @llvm.memcpy + // CHECK: ret + // CHECK-EMPTY + x.map(|v| v as u64) +}