Initialize rest of array with a default value

user_7571907 0
Viewed 622

Is there a way in Rust to initialize the first n elements of an array manually, and specify a default value to be used for the rest?

Specifically, when initializing structs, we can specify some fields, and use .. to initialize the remaining fields from another struct, e.g.:

let foo = Foo {
  x: 1,
  y: 2,
  ..Default::default()
};

Is there a similar mechanism for initializing part of an array manually? e.g.

let arr: [i32; 5] = [1, 2, ..3];

to get [1, 2, 3, 3, 3]?

rust
2 Answers

Edit: I realized this can be done on stable. For the original answer, see below.

I had to juggle with the compiler so it will be able to infer the type of the array, but it works:

// A workaround on the same method on `MaybeUninit` being unstable.
// Copy-paste from https://doc.rust-lang.org/stable/src/core/mem/maybe_uninit.rs.html#943-953.
pub unsafe fn maybe_uninit_array_assume_init<T, const N: usize>(
    array: [core::mem::MaybeUninit<T>; N],
) -> [T; N] {
    // SAFETY:
    // * The caller guarantees that all elements of the array are initialized
    // * `MaybeUninit<T>` and T are guaranteed to have the same layout
    // * `MaybeUninit` does not drop, so there are no double-frees
    // And thus the conversion is safe
    (&array as *const _ as *const [T; N]).read()
}

macro_rules! array_with_default {
    (@count) => { 0usize };
    (@count $e:expr, $($rest:tt)*) => { 1usize + array_with_default!(@count $($rest)*) };

    [$($e:expr),* ; $default:expr; $default_size:expr] => {{
        // There is no hygiene for items, so we use unique names here.
        #[allow(non_upper_case_globals)]
        const __array_with_default_EXPRS_LEN: usize = array_with_default!(@count $($e,)*);
        #[allow(non_upper_case_globals)]
        const __array_with_default_DEFAULT_SIZE: usize = $default_size;
        let mut result = unsafe { ::core::mem::MaybeUninit::<
            [::core::mem::MaybeUninit<_>; {
                __array_with_default_EXPRS_LEN + __array_with_default_DEFAULT_SIZE
            }],
        >::uninit().assume_init() };

        let mut dest = result.as_mut_ptr();
        $(
            let expr = $e;
            unsafe {
                ::core::ptr::write((*dest).as_mut_ptr(), expr);
                dest = dest.add(1);
            }
        )*
        for default_value in [$default; __array_with_default_DEFAULT_SIZE] {
            unsafe {
                ::core::ptr::write((*dest).as_mut_ptr(), default_value);
                dest = dest.add(1);
            }
        }

        unsafe { maybe_uninit_array_assume_init(result) }
    }};
}

Playground.

Based on the example from @Denys, here is a macro that works on nightly. Note that I had problems matching the .. syntax (though I'm not entirely sure that's impossible; just didn't put much time into that):

#![feature(generic_const_exprs)]
#![allow(incomplete_features)]

use std::mem::MaybeUninit;

pub fn concat_arrays<T, const N: usize, const M: usize>(a: [T; N], b: [T; M]) -> [T; N + M] {
    unsafe {
        let mut result = MaybeUninit::<[T; N + M]>::uninit();
        let dest = result.as_mut_ptr().cast::<[T; N]>();
        dest.write(a);
        let dest = dest.add(1).cast::<[T; M]>();
        dest.write(b);
        result.assume_init()
    }
}

macro_rules! array_with_default {
    [$($e:expr),* ; $default:expr; $default_size:expr] => {
        concat_arrays([$($e),*], [$default; $default_size])
    };
}

fn main() {
    dbg!(array_with_default![1, 2; 3; 7]);
}

Playground.

user_78843050

As another option, you can build a default filled array and just modify the positions you require in runtime:

#![feature(explicit_generic_args_with_impl_trait)]

fn array_with_default_and_positions<T: Copy, const SIZE: usize>(
    default: T,
    init_values: impl IntoIterator<Item = (usize, T)>,
) -> [T; SIZE] {
    let mut res = [default; SIZE];
    for (i, e) in init_values.into_iter() {
        res[i] = e;
    }
    res
}

Playground

Notice the use of #![feature(explicit_generic_args_with_impl_trait)],which is nightly, it could be replaced by an slice since T and usize are copy:

fn array_with_default_and_positions_v2<T: Copy, const SIZE: usize>(
    default: T,
    init_values: &[(usize, T)],
) -> [T; SIZE] {
    let mut res = [default; SIZE];
    for &(i, e) in init_values.into_iter() {
        res[i] = e;
    }
    res
}
user_16951720
Related