Implement `IntoIterator`, changing arrays into structs

Viewed 90

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)?
2 Answers

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).

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.

Related