API Proposal: Slice iterator helpers for taking chunks of data

[ObsidianMinor]

Problem statement

Parsing with slice iterators makes parsing easy! As long as you're going element by element... What if I want to read a whole chunk of data as a slice and advance the iterator?

Motivating examples or use cases

When parsing binary formats, often blobs of data will be prefixed by their length.

let bytes = &[3, 72, 105, 33];
let mut bytes_iter = bytes.iter();
let length = bytes_iter.next()? as usize; // single byte length
// read string all at once... normally this requires converting *back* into a slice, splitting, and reassigning
let (string, remainder) = bytes_iter.as_slice().split_at_checked(length)?;
bytes_iter = remainder.iter();
// now repeat this every single time you want to read a chunk of data.

Or fixed sized data will be somehow tagged, meaning you know exactly how much data you want, like reading an integer.

let bytes = &[55, 1, 0, 0, 0];
let mut bytes_iter = bytes.iter();
let tag = bytes_iter.next()?; // this somehow indicates we expect an LE u32
// read byte array... same thing as before but now we use `split_first_chunk`
let (value_array, remainder) = bytes_iter.as_slice().split_first_chunk<4>()?;
bytes_iter = remainder.iter();
// repeat

Solution sketch

Two simple APIs: take_slice and take_array.

impl<'a, T> core::slice::Iter<'a, T> {
    pub fn take_slice(&mut self, len: usize) -> Option<&'a [T]> {
        // Returns a slice of the data from the current point up to the given length, advancing the iterator.
        // If the iterator is shorter than `len`, this returns None and does not advance the iterator.
    }
    pub fn take_array<const N: usize>(&mut self) -> Option<&'a [T; N]> {
        // Returns an array of the data from the current point up to the given length, advancing the iterator.
        // If the iterator is shorter than `len`, this returns None and does not advance the iterator.
    }
}
impl<'a, T> core::slice::IterMut<'a, T> {
    pub fn take_slice(&mut self, len: usize) -> Option<&'a mut [T]> {
        // Returns a slice of the data from the current point up to the given length, advancing the iterator.
        // If the iterator is shorter than `len`, this returns None and does not advance the iterator.
    }
    pub fn take_array<const N: usize>(&mut self) -> Option<&'a mut [T; N]> {
        // Returns an array of the data from the current point up to the given length, advancing the iterator.
        // If the iterator is shorter than `len`, this returns None and does not advance the iterator.
    }
}

Solution examples

let bytes = &[3, 72, 105, 33];
let mut bytes_iter = bytes.iter();
let length = bytes_iter.next()? as usize;
let string = bytes_iter.take_slice(length)?;

let bytes = &[55, 1, 0, 0, 0];
let mut bytes_iter = bytes.iter();
let tag = bytes_iter.next()?;
let value = u32::from_le_bytes(bytes_iter.take_array()?);

Alternatives

As shown in the examples, you can sort-of do this with the existing core::slice::Iter::as_slice() function, though it's not ideal.

// For slices:
let (new_slice, remainder) = iter.as_slice().split_at_checked(length)?;
iter = remainder.iter();

// For arrays:
let (new_array, remainder) = bytes_iter.as_slice().split_first_chunk<4>()?;
iter = remainder.iter();

It gets more complicated with IterMut.

// For slices:
let (new_slice, remainder) = core::mem::replace(&mut iter, (&mut []).iter_mut()).into_slice().split_at_mut_checked(length)?;
iter = remainter.iter_mut();

// For arrays:
let (new_slice, remainder) = core::mem::replace(&mut iter, (&mut []).iter_mut()).into_slice().split_first_chunk_mut<N>()?;
iter = remainter.iter_mut();

[the8472]

Why would you do this on the iterator instead on a &mut &[T]?

next() is split_off_first(), take_slice(len) is split_off(len)

[ObsidianMinor]

Why would you do this on the iterator instead on a &mut &[T]?

next() is split_off_first(), take_slice(len) is split_off(len)

You could do that, it would be another option, but I don't know how efficient it'd be to constantly call split_off*. The main benefit of using the iterator is you don't need to constantly update your current slice length since the iterator is a start and end pointer. I'll have to check what the generated assembly looks like between the two.

Edit: As expected, the assembly is still having to juggle the extra length field every time it does a split operation. Ideally, this new function wouldn't need to do it since it can reuse the parameter.

[scottmcm]

Why would you do this on the iterator instead on a &mut &[T]?

I would assume for the same reason that slice::Iter isn't just a wrapper on a &[T]: that if you're storing it, rather than consuming it all at ones, there's fewer stores with the iterator than with the &mut &[T].

---

Though maybe if rust-lang/rust#122971 could land then no special APIs would be needed because the as_slice→split→iter approach would optimize perfectly, letting these just be local helpers instead of std ones.

[hanna-kruppe]

The difference between (start ptr, end ptr) and (start ptr, len) matters sometimes, but that cuts both ways. When consuming input one by one, it’s neat that only the start pointer needs to be updated and compared against the unchanging end pointer. But the proposed methods need length checks, so those methods would have to calculate the length from the iterator’s start and end. If you’d be using the proposed methods often enough to worry about their performance, then it’s quite possible that the (start ptr, len) representation is preferable anyway.

At a higher (lower?) level, I don’t think the difference actually matters frequently enough to justify mirroring lots of convenience functionality from slices into slice iterators. If you’re parsing a single sequential byte stream that interleaves control with data, then performance is usually determined by critical path of the parsing loop, not by uarch throughput (amount of uops vs. functional units). In these cases, the length update is generally not on the critical path since the length is only used by very predictable error checks and end-of-buffer checks. Conversely, if you’ve managed to make your parser throughput bound, you probably want to avoid frequent end-of-buffer checks in your hot inner loop even with the (start ptr, end ptr) representation. So there’s not much design space left where the proposed methods seem to make a significant difference.

[the8472]

Edit: As expected, the assembly is still having to juggle the extra length field every time it does a split operation

Do you have code that is bottlenecked on this and would gain a significant speedup?