Iterating over the contents of an Option, or over a specific value - rust

Question

Let's say that we have the following C-code (assume that srclen == dstlen and the length is divisible by 64).
void stream(uint8_t *dst, uint8_t *src, size_t dstlen) {
int i;
uint8_t block[64];
while (dstlen > 64) {
some_function_that_initializes_block(block);
for (i=0; i<64; i++) {
dst[i] = ((src != NULL)?src[i]:0) ^ block[i];
}
dst += 64;
dstlen -= 64;
if (src != NULL) { src += 64; }
}
}
That is a function that takes a source and a destination and xors source with some value that
the function computes. When source is set to a NULL-pointer dst is just the computed value.
In rust it is quite simple to do this when src cannot be null, we can do something like:
fn stream(dst: &mut [u8], src: &[u8]) {
let mut block = [0u8, ..64];
for (dstchunk, srcchunk) in dst.chunks_mut(64).zip(src.chunks(64)) {
some_function_that_initializes_block(block);
for (d, (&s, &b)) in dstchunk.iter_mut().zip(srcchunk.iter().zip(block.iter())) {
*d = s ^ b;
}
}
}
However let us assume that we want to be able to mimic the original C-function. Then we would like to do something like:
fn stream(dst: &mut[u8], osrc: Option<&[u8]>) {
let srciter = match osrc {
None => repeat(0),
Some(src) => src.iter()
};
// the rest of the code as above
}
Alas, this won't work since repeat(0) and src.iter() have different types. However it doesn't seem possible to solve this by using a trait object since we get a compiler error saying cannot convert to a trait object because trait 'core::iter::Iterator' is not object safe. (also there is no function in the standard library that chunks an iterator).
Is there any nice way to solve this, or should I just duplicate the code in each arm of the match statement?

Answer 1

Instead of repeating the code in each arm, you can call a generic inner function:
fn stream(dst: &mut[u8], osrc: Option<&[u8]>) {
fn inner<T>(dst: &mut[u8], srciter: T) where T: Iterator<u8> {
let mut block = [0u8, ..64];
//...
}
match osrc {
None => inner(dst, repeat(0)),
Some(src) => inner(dst, src.iter().map(|a| *a))
}
}
Note the additional map to make both iterators compatible (Iterator<u8>).
As you mentioned, Iterator doesn't have a built-in way to do chunking. Let's incorporate Vladimir's solution and use an iterator over chunks:
fn stream(dst: &mut[u8], osrc: Option<&[u8]>) {
const CHUNK_SIZE: uint = 64;
fn inner<'a, T>(dst: &mut[u8], srciter: T) where T: Iterator<&'a [u8]> {
let mut block = [0u8, ..CHUNK_SIZE];
for (dstchunk, srcchunk) in dst.chunks_mut(CHUNK_SIZE).zip(srciter) {
some_function_that_initializes_block(block);
for (d, (&s, &b)) in dstchunk.iter_mut().zip(srcchunk.iter().zip(block.iter())) {
*d = s ^ b;
}
}
}
static ZEROES: &'static [u8] = &[0u8, ..CHUNK_SIZE];
match osrc {
None => inner(dst, repeat(ZEROES)),
Some(src) => inner(dst, src.chunks(CHUNK_SIZE))
}
}

Answer 2

Unfortunately, it is impossible to use different iterators directly or with trait objects (which have recently been changed to disallow instantiation of trait objects with inappropriate methods i.e. ones which use Self type in their signature). There is a workaround for your particular case, however. Just use enums:
fn stream(dst: &mut [u8], src: Option<&[u8]>) {
static EMPTY: &'static [u8] = &[0u8, ..64]; // '
enum DifferentIterators<'a> { // '
FromSlice(std::slice::Chunks<'a, u8>), // '
FromRepeat(std::iter::Repeat<&'a [u8]>) // '
}
impl<'a> Iterator<&'a [u8]> for DifferentIterators<'a> { // '
#[inline]
fn next(&mut self) -> Option<&'a [u8]> { // '
match *self {
FromSlice(ref mut i) => i.next(),
FromRepeat(ref mut i) => i.next()
}
}
}
let srciter = match src {
None => FromRepeat(repeat(EMPTY)),
Some(src) => FromSlice(src.chunks(64))
};
let mut block = [0u8, ..64];
for (dstchunk, srcchunk) in dst.chunks_mut(64).zip(srciter) {
some_function_that_initializes_block(block);
for (d, (&s, &b)) in dstchunk.iter_mut().zip(srcchunk.iter().zip(block.iter())) {
*d = s ^ b;
}
}
}
This is a lot of code, unfortunately, but in return it is more safe and less error-prone than the C version. It is also possible to optimize it in order not to require repeat() at all:
fn stream(dst: &mut [u8], src: Option<&[u8]>) {
static EMPTY: &'static [u8] = &[0u8, ..64]; // '
enum DifferentIterators<'a> { // '
FromSlice(std::slice::Chunks<'a, u8>), // '
AlwaysZeros
}
impl<'a> Iterator<&'a [u8]> for DifferentIterators<'a> { // '
#[inline]
fn next(&mut self) -> Option<&'a [u8]> { // '
match *self {
FromSlice(ref mut i) => i.next(),
AlwaysZeros => Some(STATIC),
}
}
}
let srciter = match src {
None => AlwaysZeros,
Some(src) => FromSlice(src.chunks(64))
};
let mut block = [0u8, ..64];
for (dstchunk, srcchunk) in dst.chunks_mut(64).zip(srciter) {
some_function_that_initializes_block(block);
for (d, (&s, &b)) in dstchunk.iter_mut().zip(srcchunk.iter().zip(block.iter())) {
*d = s ^ b;
}
}
}