Support vec![const { ... }; n] syntax

Author: EFanZh

library/alloc/src/macros.rs

@@ -43,6 +43,10 @@ macro_rules! vec {
     () => (
         $crate::vec::Vec::new()
     );
+    (const $elem:block; $n:expr) => (
+        // SAFETY: The `const` keyword asserts the value being the result of a const expression.
+        unsafe { $crate::vec::from_const_elem(const $elem, $n) }
+    );
     ($elem:expr; $n:expr) => (
         $crate::vec::from_elem($elem, $n)
     );

library/alloc/src/vec/mod.rs

@@ -3183,6 +3183,57 @@ pub fn from_elem_in<T: Clone, A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<
     <T as SpecFromElem>::from_elem(elem, n, alloc)
 }
 
+/// # Safety
+///
+/// `elem` must be the result of some const expression.
+#[doc(hidden)]
+#[cfg(not(no_global_oom_handling))]
+#[unstable(feature = "vec_of_const_expr", issue = "none")]
+#[track_caller]
+pub unsafe fn from_const_elem<T>(elem: T, n: usize) -> Vec<T> {
+    // SAFETY: Caller has guaranteed `elem` being the result of some const expression.
+    unsafe { from_const_elem_in(elem, n, Global) }
+}
+
+/// # Safety
+///
+/// `elem` must be the result of some const expression.
+#[doc(hidden)]
+#[cfg(not(no_global_oom_handling))]
+#[unstable(feature = "vec_of_const_expr", issue = "none")]
+#[track_caller]
+unsafe fn from_const_elem_in<T, A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
+    /// # Safety
+    ///
+    /// `value` must points to a valid `T` value that is the result of some const expression.
+    unsafe fn fill_const_value<T>(buffer: &mut [MaybeUninit<T>], value: *const T) {
+        for target in buffer {
+            // SAFETY: If `value` is the result of some const expression, we can make as many
+            // copies as needed.
+            unsafe { target.write(ptr::read(value)) };
+        }
+    }
+
+    // Avoid calling the destructor of `elem`.
+    let elem = ManuallyDrop::new(elem);
+    let elem_ptr = ptr::from_ref(&*elem);
+    let mut result = Vec::<T, A>::with_capacity_in(n, alloc);
+    let buffer_ptr = result.as_mut_ptr().cast::<MaybeUninit<T>>();
+
+    // SAFETY: `with_capacity_in` makes sure the capacity is at least `n`, so we can make a buffer
+    // of length `n` out of it.
+    let buffer = unsafe { slice::from_raw_parts_mut(buffer_ptr, n) };
+
+    // SAFETY: Caller has guaranteed `elem` being the result of some const expression.
+    unsafe { fill_const_value(buffer, elem_ptr) };
+
+    // SAFETY: We have initialized exactly `n` values at the start of the buffer, so we are safe to
+    // set the length accordingly.
+    unsafe { result.set_len(n) };
+
+    result
+}
+
 #[cfg(not(no_global_oom_handling))]
 trait ExtendFromWithinSpec {
     /// # Safety

library/alloc/tests/lib.rs

@@ -36,6 +36,7 @@
 #![cfg_attr(not(bootstrap), feature(strict_provenance_lints))]
 #![feature(drain_keep_rest)]
 #![feature(local_waker)]
+#![feature(vec_of_const_expr)]
 #![feature(vec_pop_if)]
 #![feature(unique_rc_arc)]
 #![feature(macro_metavar_expr_concat)]

library/alloc/tests/vec.rs

@@ -2311,6 +2311,28 @@ fn test_vec_macro_repeat() {
     assert_eq!(vec![el; n], vec![Box::new(1), Box::new(1), Box::new(1)]);
 }
 
+#[test]
+fn test_vec_macro_repeat_const() {
+    #[derive(Eq, PartialEq, Debug)]
+    struct Item {
+        x: u64,
+        y: u32, // Paddings are added to test the uninitialized bytes case.
+    }
+
+    impl Clone for Item {
+        fn clone(&self) -> Self {
+            panic!("no clone should be called");
+        }
+    }
+
+    const ITEM: Item = Item { x: 2, y: 3 };
+
+    assert_eq!(vec![const { ITEM }; 0], vec![ITEM; 0]);
+    assert_eq!(vec![const { ITEM }; 1], vec![ITEM]);
+    assert_eq!(vec![const { ITEM }; 2], vec![ITEM, ITEM]);
+    assert_eq!(vec![const { ITEM }; 3], vec![ITEM, ITEM, ITEM]);
+}
+
 #[test]
 fn test_vec_swap() {
     let mut a: Vec<isize> = vec![0, 1, 2, 3, 4, 5, 6];