Arithmetic ops for arrays / slices

I found myself staring at a handwritten fn xor_in_place and wondering if there was anything built into the language to do it for me, and decided what I would really like is if in-place arithmetic were supported for core arrays. Something like:

impl<T: AddAssign, const N: usize> AddAssign for [T; N] {
     fn add_assign(&mut self, rhs: Self) {
        for (a, b) in self.iter_mut().zip(rhs) {
            *a += b;
        }
     }
}

...but for all the core arithmetic ops.

(note: that could probably be more general with respect to Rhs, just an example)

That would make it possible to do:

let mut x = [1, 2, 3];
let y = [4, 5, 6];    
x += y;

It seems like similar functionality could be added to slices as well, although there are a lot more decisions there like how to handle length mismatches.


Edit: I chose to use AddAssign as an example but upon further reflection that might not be the greatest place to start due to questions about e.g. overflow handling.

Bitwise operations seem a little more straightforward and that's really my main use case.

I originally opened this as "in-place" operations but there's no reason that e.g. BitAnd/BitOr/BitXor couldn't be supported as well, so I removed "in-place" from the title.

1 Like

Rust is not Fortran or Matlab; element-wise arithmetic on arrays is not nearly common enough a use case to justify adding these IMO. Arrays just aren’t arithmetic types in Rust, and the meaning of arithmetic ops would be non-obvious. These operations are much more suitable to a third-party vector type that wraps arrays.

3 Likes

Hmm, this already works: https://play.rust-lang.org/?version=stable&mode=debug&edition=2024&gist=a624cf437790f77949f92195edf7c83f

let mut x = [1, 2, 3];
let y = [4, 5, 6];
std::iter::zip(&mut x, &y).for_each(|(a, b)| a.add_assign(b));
dbg!(x);

Which isn't bad at all.

I wonder if there's a place for something to simplify the tuple-to-separate-arguments step, like how there's tuple and untuple in functional languages.

If it was just (since u32: AddAssign<&u32>)

std::iter::zip(&mut x, &y).for_each_pair(AddAssign::add_assign);

or

std::iter::zip(&mut x, &y).for_each(tupled(AddAssign::add_assign));

or something that might be fine.


As a meta-point, the other problem with these operators on arrays is that they're eager. And like how array::map warns about it, you don't really want x * y + z on arrays to be eager that way, because once the array gets long it'll codegen poorly.

What you really want is a way to separate the transformer from the iterator chain, so you can build a transformer that represents the _ * _ + _, then apply that element-wise to the inputs. Ideally in a way that the transformer could be used both for things like arrays but also for iterators too.

(I have no idea what such an API would look like, though.)

6 Likes

This reminds me of what Eigen in C++ does. It basically uses a lot of template black magic to delay the concrete evaluation until an assignment is made to a concrete vector or matrix type. With it being C++ it does of course come with footguns, around self assignments and aliasing (and it is up to the user to keep this in mind).

I haven't had the need to look at linear algebra since I learned rust, so I don't know if the Rust libraries also does things like that (or if Rust generics are even capable of expressing this).

1 Like

Rust's iterators are expression templates (the pattern behind Eigen): the type builds up a series of operations to calculate on the sequence and doesn't make a concrete type until an appropriate API is called. Futures are also expression templates for that matter. However, as with iterators, the problems arise when one needs to name the types on the boundaries of APIs. In C++, this means using Eigen and auto can end up with some…funky error messages (the same problem comes up with auto and the proxy types used by std::vector<bool>).

That is nothing new for C++ though. 1000s of lines of errors due to a small typo is par for the course.

But it sounds like good news, the pattern should work in Rust (if any library implements it for matrices etc I don't know).

I'm afraid that on slices this could make error messages worse, or even cause run-time errors (from mismatched len) when people accidentally mix up scalars and arrays.

The set of operations supported by overloads is limited. There's already array.map. Perhaps array.zip should exist too?

I'm also concerned that eager evaluation (that includes .map and potential .zip) can be a performance footgun.

It would be nice to have an iterator-like interface. Unfortunately the Iterator is not great at working with fixed-length sets, and FromIterator can't make an array. Can Iterator/FromIterator be fixed? Would it be worth adding some kind of FixedLenIterator for arrays?

I think the main point here is that you can't have a fixed length version of anything with a next(&mut self), because it's fundamentally not fixed-length. It needs to be more IntoIterator that has the fixed length, though back in 2018 libs-api wasn't interested on a size hint for the current IntoIterator, at least.

It used to, but was removed: Remove array_zip by workingjubilee · Pull Request #112096 · rust-lang/rust · GitHub

It's not that useful when it's eager, because reifying the array of pairs is almost never what you want. Perhaps a zip_map could work, though.

What I'm imagining is that you still have an iterator-like interface, but you do it over abstract "sources" of some kind, rather than over a particular input iterator. In the scalar version, that's like the classic "compose" combinator that takes T->U and U->V giving T->V, but which arguably is essentially the same as a combining two Iterator::Maps.

So you'd build up some impl Transformer (I have no idea what the trait bounds on it would be). Then you could always pass iterator(s) to that to get out iterators, but also if it implements KnownLengthTransformer then there's a way to pass array(s) to it and get array(s) out.

(Well, it might need to be LengthPreservingTransformer for now, because we can't do [T; N] -> [U; {N * Self::MULTIPLIER}] in the current state of what constant generic expressions are allowed to do.)

I think I may have confused the issue by proposing all arithmetic ops.

The bitwise ops are the ones I actually cared about here and generally sidestep most of the problems that have been raised, especially in regard to working in special arithmetic forms to allow for various forms of lazy evaluation.

(For what its worth I maintain elliptic curve libraries that use projective coordinates to defer inversions as well as lazy field operations that defer reductions for exactly those purposes. For things like linear combinations though, we use dedicated APIs)

TBH then do you actually need generic operations on arrays, or just bigger integer types? Suppose you could just u<512> and do bitops on that...

3 Likes

How about something like element_wise!{x += y}? That would be expressive and could allow more complex expressions. Though implementing it as a macro could be more difficult than expected, at least if more complex expressions should be supported. Something similar could be useful for unchecked_*, wrapping_*, ... methods on integers: unchecked!{a + b - c} * 7 (using a macro to not need a new keyword), making larger chains of these methods more readable: a.unchecked_add(b).unchecked_sub(c) * 7.

Though doing this as a (decl.) macro would be really difficult if other code should be allowed to be part of the block. For example to just set the default way to be used for things like a + b.

1 Like

Yeah, I've had zip_map in mind when I said zip.