How to represent an array of only certain sizes in Rust

71 views Asked by At

I would like a data structure that represents an array of values, but which only supports a specific set of sizes and enforces this at compile time. Something like:

struct MyArray<const N: usize>([u8; N]);

But such than N can only be a specific set of values, not just any number representable by usize. For example, I would like a struct can wrap either a [u8; 3], a [u8; 6], or a [u8; 9], but not a [u8; N] for any other N besides 3, 6, and 9. I need the enforcement of this constraint to be at compile time, and ideally part of the type system. Is this possible in Rust? Is there a standard pattern or crate for doing this?

3

There are 3 answers

2
Richard Neumann On BEST ANSWER

Depending on your use case, you can also represent the custom array as an enum:

#[derive(Clone, Debug)]
enum MyArray<T> {
    Three([T; 3]),
    Six([T; 6]),
    Nine([T; 9]),
}

fn main() {
    let mut array = MyArray::<u8>::Six(Default::default());

    if let MyArray::Three(_) = &array {
        println!("Nope.");
    }

    if let MyArray::Six(array) = &mut array {
        array
            .iter_mut()
            .enumerate()
            .for_each(|(index, element)| *element = index as u8);
    }

    println!("{array:?}");
    println!("Size: {}", std::mem::size_of::<MyArray<u8>>());
}

Caveat: This will not be as space efficient as the separate impls, due to the nature of enums (size of largest variant). However, if 9 bytes is the maximum size, I assume that this is not an issue in your case.

3
drewtato On

You can do this by only providing constructors for your desired lengths.

impl MyArray<3> {
    pub fn new(arr: [u8; 3]) -> Self {
        Self(arr)
    }
}

impl MyArray<6> {
    pub fn new(arr: [u8; 6]) -> Self {
        Self(arr)
    }
}

impl MyArray<9> {
    pub fn new(arr: [u8; 9]) -> Self {
        Self(arr)
    }
}

Now outside of the module containing MyArray, you will only be able to make MyArray values with lengths of 3, 6, or 9.

If you have more numbers, a macro will help.

macro_rules! impl_my_array {
    ($($n:expr),* $(,)?) => {$(
        impl MyArray<$n> {
            pub fn new(arr: [u8; $n]) -> Self {
                Self(arr)
            }
        }
    )*}
}

impl_my_array! {
    3, 6, 9
}

The downside to these is that type inference cannot figure out which new to call, so you will need to specify the const generic whenever you construct values.

//                 V
let ma = MyArray::<3>::new([1, 2, 3]);

You may want to give them different names, or provide a higher-level constructor that either picks a specific new or returns an error.


If you need a large amount of numbers, like to span the entire usize range, you'll need to use the const trick.

This one allows every usize divisible by 3.

impl<const N: usize> MyArray<N> {
    const LEN_IS_DIVISIBLE_BY_3: () = if N % 3 != 0 {
        panic!("MyArray values must have length divisible by 3");
    };
    pub fn new(arr: [u8; N]) -> Self {
        let _ = Self::LEN_IS_DIVISIBLE_BY_3;
        Self(arr)
    }
}

This is used in the bytemuck crate here and here.

0
Chayim Friedman On

The answer by @drewtato is good, but has two drawbacks: first, code cannot be generic over the size if it wants to create instances, and second, the code duplication (which can be solved with a macro, but that's still not the best thing).

A better solution (IMHO) that is also used by the standard library for the nightly SIMD module, is to have a trait and implement it only for specific sizes:

trait SupportedArraySize {}

struct MyArray<const N: usize>([u8; N]);

impl SupportedArraySize for MyArray<3> {}
impl SupportedArraySize for MyArray<6> {}
impl SupportedArraySize for MyArray<9> {}

impl<const N: usize> MyArray<N>
where
    MyArray<N>: SupportedArraySize,
{
    fn new() { /* ... */
    }
}