Making elementwise operations using const generics more ergonomic? (Possible new trait or FromIterator impl)

The Problem

I was experimenting with const generics and I noticed a pattern that cropped up quite a bit. I was using a wrapper around a [T; N] called Vector to do some linear algebra:

struct Vector<T, const N: usize>([T; N])

I wanted to implement elementwise addition and subtraction to this struct.

let a = Vector([1, 2, 3)];
let b = Vector([1, 1, 1]);

assert_eq(a + b, Vector([2, 3, 4]));

In order to implement elementwise addition with safe rust I used this impl:

impl<T: Add<Output = T> + Copy + Default, const N: usize> Add<T> for Vector<T, { N }> {
    type Output = Self;

    fn add(self, other: T) -> Vector<T, N> {
        let mut out: [T; N] = [Default::default(); N];
        for x in 0..N {
            out[x] = self.0[x] + other;
        }
        Self(out)
    }
}

There is a problem here that could be fixed, the Default impl is technically unnecessary and simply used to create an initialized array. With unsafe code it is possible to remove the Default impl.

impl<T: Add<Output = T> + Copy, const N: usize> Add for Vector<T, { N }> {
    type Output = Self;
    fn add(self, other: Vector<T, N>) -> Vector<T, N> {
        let mut out: [MaybeUninit<T>; N] = unsafe { MaybeUninit::uninit().assume_init()};
        for x in 0..N {
            out[x] = MaybeUninit::new(self.0[x] + other.0[x]);
        }
        let arr = unsafe { mem::transmute_copy::<[MaybeUninit<T>; N],[T; N]>(&out) };
        Self(arr)
    }
}

The transmute_copy is due to transmute currently not working with const generics. I don't write much unsafe code, so I'm not sure about this safe.

More generally, there would be quite an obvious desire to make it ergonomic to write functions with a signature like this:

fn example<const N: usize>() -> [usize; N]

where there would be some elementwise operation occuring inside the function.

How can we improve this

(1) Allow initially uninitialed arrays (but initialized before returning) to compile.

Now the cleanest way of allowing these kinds of functions to be implemented is to simply allow a function like this compile:

fn add<const N: usize>(a: [usize; N], b: [usize; N])  -> [usize; N]{
    let arr:[usize; N];
    for idx in 0..N {
        arr[idx] = a[idx] + b[idx];
    }
    arr
}

Currently this does not compile due to arr possibly not being initialized. This seems like it would be a lot of work to implement in the compiler. A slight alteration would be to create a special syntax for doing this kind of operation.

(2) Trait level support for collecting on a array.

Iterators already are an abstraction over performing elementwise operations. Therefore another option would be to provide some trait level for producing an array from support to abstract the idea of "collecting" into an array from an iterator. This could be a FromIterator impl (although this feels like a bit of an abuse of the trait.) Another option would be to add a new trait (maybe ConstFromIterator?) to abstract over converting an iterator to an array. We could use ExactSizeIterator to aid in this.

This could allow the same interface to be used with other const-generic collections (such as collecting into my Vector.)

For example, the addition of the FromIterator trait would make the my Add impl look like this:

impl<T: Add<Output = T> + Copy, const N: usize> Add for Vector<T, { N }> {
    type Output = Self;
    fn add(self, other: Vector<T, N>) -> Vector<T, N> {
        self.0
            .iter()
            .zip(other.0.iter())
            .map(|(a, b)| *a + *b)
            .collect()
    }
}

Notice that there is no unsafe and no Default impl. it also looks a lot cleaner. A demo of the ForIterator impl would be:

impl<T, const N: usize> FromIterator<T> for Vector<T, N> {
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = T>,
    {
        let mut out: [MaybeUninit<T>; N] = unsafe { MaybeUninit::uninit().assume_init() };
        let mut iter = iter.into_iter();
        for idx in 0..N {
            match iter.next() {
                Some(x) => out[idx] = MaybeUninit::new(x),
                None => panic!(),
            }
        }
        unsafe { mem::transmute_copy::<[MaybeUninit<T>; N], [T; N]>(&out) }
    }
}

This is horrendously unsafe (again) but it shows that the impl shields the library author from unsafe.

(3) Free function/macro/crate support

Free function or macro support. The third option is not to require any trait and instead use a free function to convert an iterator to an array, perhaps a function with the signature:

fn const_collect<I, T, const N: usize>(iter: I) -> [T; N]
where
    I: IntoIterator<Item = T>,
    T: Copy

This would essentially do the same as the FromIterator impl.

Conclusion

The option of a new trait and a free function/macro do not absolutlely require std support so it may not be a libs/lang issue, but it feels like there should be std/lang support for this eventually.

It might be possible to do this in safe code by iterating on a slice, then using the TryFrom<&'a [T]> for [T; N] impl. Even if this is safe this is quite ugly.

Obviously this is far down the line from a const-generics MVP, but I thought it would be cool to think about how const-generics could actually be used and make them ergonomic.

1 Like

I think something like this could probably be placed into a library. Personally I've written something like option 2 several times, as it's a good abstraction for array operations.

The "unsafe code around an uninitialized array that tracks which parts are initialized" is essentially arrayvec. And having the length be tracked is nice in that it doesn't force a collect to throw the elements on the floor if it turns out to be shorter.

I'd be quite happy to have a const-generics-using ArrayVec type thing in core -- and we in fact already almost have an ArrayDeque in core even; it's just called core::array::IntoIter.