Implement `IntoIterator`, changing arrays into structs

Question

I have a basic struct ImageData which, primarily, holds an array of u8s (bytes) that represents a PNG image (layed out as rgbrgbrgb... or rgbargbargba...):

struct ImageData {
    png_type: png::ColorType,
    bytes: Vec<u8>,
    width: u32,
    height: u32,
    is_srgb: bool,
}

I also have a PixelValue struct that will allow me to treat RGB and RGBA similarly - that is, I want to get PixelValues instead of an array of 3 or 4 u8s, because then I would be duplicating a lot of code to deal with those different types.

struct PixelValue {
    r: u8,
    g: u8,
    b: u8,
    a: u8,
}

So I want to implement IntoIterator on ImageData that will return an iterator on PixelValues. It should convert the u8 chunks into PixelValues:

impl<'a> IntoIterator for &'a ImageData {
    type Item = PixelValue;
    type IntoIter = Iter<PixelValue>; // this is definitely wrong

    fn into_iter(self) -> Self::IntoIter {
        let pixel_size = get_pixel_size(self.png_type);  // 3 or 4, depending on RGB or RGBA
        // Need to write something that can take an array of 3 or 4 u8s and returns the PixelValue
        self.bytes.chunks(pixel_size).map(|ar| PixelValue(ar))
    }
}

My questions are:

Am I going in the right direction? Can I change the output value (Item) of the iterator with map? If not, why not?
Is there an obvious "best practice" way to go about this problem (turning an array of pixel rgb or rgba's into structs)?

One limitation with Rust (currently) is that you can't use `impl Trait` in traits, so you end up either having to quite frustratingly define the `IntoIter` associated type rather unwieldily (something like `Map>, fn(&[u8]) -> PixelValue>` and ensure the "closure" doesn't capture any environment so that its type can be "promoted" to a `fn` pointer) or (often more preferably) you define your own struct implementing `Iterator` and return that. — eggyal, May 05 '22 at 18:00
If your pixel value structs are `#[repr(C)]` then you can cast a slice's data pointer to a struct reference (in unsafe code) and avoid the conversion altogether. — cdhowie, May 05 '22 at 18:08
@cdhowie: How would a `[u8; 3]` (in the case of data without the alpha channel) be casted to `PixelValue`? — eggyal, May 05 '22 at 18:09
@eggyal There would need to be an RGB version of the struct to do that. I assumed OP had a different struct for different pixel layouts, as is common when doing this kind of work. — cdhowie, May 05 '22 at 18:13

Netwave · Answer 1 · 2022-05-05T18:26:20.847

You can use dynamic dispatch, wrapping the iterator in a Box:

struct ImageData {
    png_type: u8,
    bytes: Vec<u8>,
    width: u32,
    height: u32,
    is_srgb: bool,
}

struct PixelValue {
    r: u8,
    g: u8,
    b: u8,
    a: u8,
}

impl<'a> IntoIterator for &'a ImageData {
    type Item = PixelValue;
    type IntoIter = Box<dyn Iterator<Item = Self::Item> + 'a>;

    fn into_iter(self) -> Self::IntoIter {
        let size = 3;
        Box::new(self.bytes.chunks(size).map(move |chunk| match size {
            3 => PixelValue {
                r: chunk[0],
                g: chunk[1],
                b: chunk[2],
                a: 255,
            },
            _ => PixelValue {
                r: chunk[0],
                g: chunk[1],
                b: chunk[2],
                a: chunk[3],
            },
        }))
    }
}

Playground

Notice that for the example I changed a few types I did not had access to (png_type, and the get_size method).

thisisrandy · Answer 2 · 2022-05-05T18:39:22.540

You can create an ImageDataIntoPixelValueIterator that implements Iterator and use it as type IntoIter. NB: I've removed png_type from ImageData in order to make this code reproducible and pretended that the alpha channel is always present:

struct ImageData {
    bytes: Vec<u8>,
    width: u32,
    height: u32,
    is_srgb: bool,
}

struct PixelValue {
    r: u8,
    g: u8,
    b: u8,
    a: u8,
}

impl IntoIterator for ImageData {
    type Item = PixelValue;
    type IntoIter = ImageDataIntoPixelValueIterator;

    fn into_iter(self) -> Self::IntoIter {
        ImageDataIntoPixelValueIterator {
            bytes: self.bytes,
            index: 0,
        }
    }
}

struct ImageDataIntoPixelValueIterator {
    bytes: Vec<u8>,
    index: usize,
}

impl Iterator for ImageDataIntoPixelValueIterator {
    type Item = PixelValue;

    fn next(&mut self) -> Option<PixelValue> {
        // assumes that self.bytes.len() is a multiple of four.
        // this should probably be asserted somewhere else
        if self.index >= self.bytes.len() {
            None
        } else {
            let res = PixelValue {
                r: self.bytes[self.index],
                g: self.bytes[self.index + 1],
                b: self.bytes[self.index + 2],
                a: self.bytes[self.index + 3],
            };
            self.index += 4;
            Some(res)
        }
    }
}

Playground

I've basically implemented this answer for your specific needs. Please read that thoroughly, as it provides a wealth of explanatory information to help you understand why certain things will or will not work.

Implement `IntoIterator`, changing arrays into structs

2 Answers2