Carefully remove bounds checks from some chunk iterator functions by thomcc · Pu...

thomcc commented on Jul 8, 2021

So, I was writing code that requires the equivalent of rchunks(N).rev() (which isn't the same as forward chunks(N) — in particular, if the buffer length is not a multiple of N, I must handle the "remainder" first).

I happened to look at the codegen output of the function (I was actually interested in whether or not a nested loop was being unrolled — it was), and noticed that in the outer rchunks(n).rev() loop, LLVM seemed to be unable to remove the bounds checks from the iteration: https://rust.godbolt.org/z/Tnz4MYY8f (this panic was from the split_at in RChunks::next_back).

After doing some experimentation, it seems all of the next_back in the non-exact chunk iterators have the issue: (Chunks::next_back, RChunks::next_back, ChunksMut::next_back, and RChunksMut::next_back)...

Even worse, the forward rchunks iterators sometimes have the issue as well (... but only sometimes). For example https://rust.godbolt.org/z/oGhbqv53r has bounds checks, but if I uncomment the loop body, it manages to remove the check (which is bizarre, since I'd expect the opposite...). I suspect it's highly dependent on the surrounding code, so I decided to remove the bounds checks from them anyway. Overall, this change includes:

• All next_back functions on the non-Exact iterators (e.g. R?Chunks(Mut)?).
• All next functions on the non-exact rchunks iterators (e.g. RChunks(Mut)?).

I wasn't able to catch any of the other chunk iterators failing to remove the bounds checks (I checked iterations over r?chunks(_exact)?(_mut)? with constant chunk sizes under -O3, -Os, and -Oz), which makes sense, since these were the cases where it was harder to prove the bounds check correct to remove...

In fact, it took quite a bit of thinking to convince myself that using unchecked_ here was valid — so I'm not really surprised that LLVM had trouble (although compilers are slightly better at this sort of reasoning than humans). A consequence of that is the fact that the // SAFETY comment for these are... kinda long...

I didn't do this for, or even think about it for, any of the other iteration methods; just next and next_back (where it mattered). If this PR is accepted, I'll file a follow up for someone (possibly me) to look at the others later (in particular, nth/nth_back looked like they had similar logic), but I wanted to do this now, as IMO next/next_back are the most important here, since they're what gets used by the iteration protocol.

Note: While I don't expect this to impact performance directly, the panic is a side effect, which would otherwise not exist in these loops. That is, this could prevent the compiler from being able to move/remove/otherwise rework a loop over these iterators (as an example, it could not delete the code for a loop whose body computes a value which doesn't get used).

Also, some like to be able to have confidence this code has no panicking branches in the optimized code, and "no bounds checks" is kinda part of the selling point of Rust's iterators anyway.