1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use std::io::{self, Error, ErrorKind, Read, Write};

use super::element::IconElement;
use super::icontype::IconType;
use super::image::Image;

/// The first four bytes of an ICNS file:
const ICNS_MAGIC_LITERAL: &[u8; 4] = b"icns";

/// The length of an icon family header, in bytes:
const ICON_FAMILY_HEADER_LENGTH: u32 = 8;

/// A set of icons stored in a single ICNS file.
pub struct IconFamily {
    /// The icon elements stored in the ICNS file.
    pub elements: Vec<IconElement>,
}

impl IconFamily {
    /// Creates a new, empty icon family.
    pub fn new() -> IconFamily {
        IconFamily {
            elements: Vec::new(),
        }
    }

    /// Returns true if the icon family contains no icons nor any other
    /// elements.
    pub fn is_empty(&self) -> bool {
        self.elements.is_empty()
    }

    /// Encodes the image into the family, automatically choosing an
    /// appropriate icon type based on the dimensions of the image.  Returns
    /// an error if there is no supported icon type matching the image
    /// dimensions.
    pub fn add_icon(&mut self, image: &Image) -> io::Result<()> {
        if let Some(icon_type) = IconType::from_pixel_size(image.width(), image.height()) {
            self.add_icon_with_type(image, icon_type)
        } else {
            let msg = format!(
                "no supported icon type has dimensions {}x{}",
                image.width(),
                image.height()
            );
            Err(Error::new(ErrorKind::InvalidInput, msg))
        }
    }

    /// Encodes the image into the family using the given icon type.  If the
    /// selected type has an associated mask type, the image mask will also be
    /// added to the family.  Returns an error if the image has the wrong
    /// dimensions for the selected type.
    pub fn add_icon_with_type(&mut self, image: &Image, icon_type: IconType) -> io::Result<()> {
        self.elements
            .push(IconElement::encode_image_with_type(image, icon_type)?);
        if let Some(mask_type) = icon_type.mask_type() {
            self.elements
                .push(IconElement::encode_image_with_type(image, mask_type)?);
        }
        Ok(())
    }

    /// Returns a list of all (non-mask) icon types for which the icon family
    /// contains the necessary element(s) for a complete icon image (including
    /// alpha channel).  These icon types can be passed to the
    /// [`get_icon_with_type`](#method.get_icon_with_type) method to decode the
    /// icons.
    pub fn available_icons(&self) -> Vec<IconType> {
        let mut result = Vec::new();
        for element in &self.elements {
            if let Some(icon_type) = element.icon_type() {
                if !icon_type.is_mask() {
                    if let Some(mask_type) = icon_type.mask_type() {
                        if self.find_element(mask_type).is_ok() {
                            result.push(icon_type);
                        }
                    } else {
                        result.push(icon_type);
                    }
                }
            }
        }
        result
    }

    /// Determines whether the icon family contains a complete icon with the
    /// given type (including the mask, if the given icon type has an
    /// associated mask type).
    pub fn has_icon_with_type(&self, icon_type: IconType) -> bool {
        if self.find_element(icon_type).is_err() {
            return false;
        } else if let Some(mask_type) = icon_type.mask_type() {
            return self.find_element(mask_type).is_ok();
        }
        true
    }

    /// Decodes an image from the family with the given icon type.  If the
    /// selected type has an associated mask type, the two elements will
    /// decoded together into a single image.  Returns an error if the
    /// element(s) for the selected type are not present in the icon family, or
    /// the if the encoded data is malformed.
    pub fn get_icon_with_type(&self, icon_type: IconType) -> io::Result<Image> {
        let element = self.find_element(icon_type)?;
        if let Some(mask_type) = icon_type.mask_type() {
            let mask = self.find_element(mask_type)?;
            element.decode_image_with_mask(mask)
        } else {
            element.decode_image()
        }
    }

    /// Private helper method.
    fn find_element(&self, icon_type: IconType) -> io::Result<&IconElement> {
        let ostype = icon_type.ostype();
        self.elements
            .iter()
            .find(|el| el.ostype == ostype)
            .ok_or_else(|| {
                let msg = format!(
                    "the icon family does not contain a '{}' \
                               element",
                    ostype
                );
                Error::new(ErrorKind::NotFound, msg)
            })
    }

    /// Reads an icon family from an ICNS file.
    pub fn read<R: Read>(mut reader: R) -> io::Result<IconFamily> {
        let mut magic = [0u8; 4];
        reader.read_exact(&mut magic)?;
        if magic != *ICNS_MAGIC_LITERAL {
            let msg = "not an icns file (wrong magic literal)";
            return Err(Error::new(ErrorKind::InvalidData, msg));
        }
        let file_length = reader.read_u32::<BigEndian>()?;
        let mut file_position: u32 = ICON_FAMILY_HEADER_LENGTH;
        let mut family = IconFamily::new();
        while file_position < file_length {
            let element = IconElement::read(reader.by_ref())?;
            file_position += element.total_length();
            family.elements.push(element);
        }
        Ok(family)
    }

    /// Writes the icon family to an ICNS file.
    pub fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
        writer.write_all(ICNS_MAGIC_LITERAL)?;
        writer.write_u32::<BigEndian>(self.total_length())?;
        for element in &self.elements {
            element.write(writer.by_ref())?;
        }
        Ok(())
    }

    /// Returns the encoded length of the file, in bytes, including the
    /// length of the header.
    pub fn total_length(&self) -> u32 {
        let mut length = ICON_FAMILY_HEADER_LENGTH;
        for element in &self.elements {
            length += element.total_length();
        }
        length
    }
}

#[cfg(test)]
mod tests {
    use super::super::element::IconElement;
    use super::super::icontype::{IconType, OSType};
    use super::super::image::{Image, PixelFormat};
    use super::*;
    use std::fs::File;
    use std::io::{BufReader, Cursor};

    #[test]
    fn icon_with_type() {
        let mut family = IconFamily::new();
        assert!(!family.has_icon_with_type(IconType::RGB24_16x16));
        let image = Image::new(PixelFormat::Gray, 16, 16);
        family
            .add_icon_with_type(&image, IconType::RGB24_16x16)
            .unwrap();
        assert!(family.has_icon_with_type(IconType::RGB24_16x16));
        assert!(family.get_icon_with_type(IconType::RGB24_16x16).is_ok());
    }

    #[test]
    fn write_empty_icon_family() {
        let family = IconFamily::new();
        assert!(family.is_empty());
        assert_eq!(0, family.elements.len());
        let mut output: Vec<u8> = vec![];
        family.write(&mut output).expect("write failed");
        assert_eq!(b"icns\0\0\0\x08", &output as &[u8]);
    }

    #[test]
    fn read_icon_family_with_fake_elements() {
        let input: Cursor<&[u8]> =
            Cursor::new(b"icns\0\0\0\x1fquux\0\0\0\x0efoobarbaz!\0\0\0\x09#");
        let family = IconFamily::read(input).expect("read failed");
        assert_eq!(2, family.elements.len());
        assert_eq!(OSType(*b"quux"), family.elements[0].ostype);
        assert_eq!(6, family.elements[0].data.len());
        assert_eq!(OSType(*b"baz!"), family.elements[1].ostype);
        assert_eq!(1, family.elements[1].data.len());
    }

    #[test]
    fn write_icon_family_with_fake_elements() {
        let mut family = IconFamily::new();
        family
            .elements
            .push(IconElement::new(OSType(*b"quux"), b"foobar".to_vec()));
        family
            .elements
            .push(IconElement::new(OSType(*b"baz!"), b"#".to_vec()));
        let mut output: Vec<u8> = vec![];
        family.write(&mut output).expect("write failed");
        assert_eq!(
            b"icns\0\0\0\x1fquux\0\0\0\x0efoobarbaz!\0\0\0\x09#",
            &output as &[u8]
        );
    }

    #[test]
    #[cfg(feature = "pngio")]
    fn png_unchanged() {
        let mut icon_family = IconFamily::new();
        let mut png_data = Vec::new();
        BufReader::new(File::open("tests/png/256x256.png").unwrap())
            .read_to_end(&mut png_data)
            .unwrap();
        let image = Image::read_png(png_data.as_slice()).unwrap();
        icon_family.add_icon(&image).unwrap();

        // Save the updated icon family to a new ICNS 'file'.
        let mut out = Vec::new();
        icon_family.write(&mut out).unwrap();

        // Read it in again and check the PNG is untouched.
        let icon_family = IconFamily::read(out.as_slice()).unwrap();
        let image = icon_family
            .get_icon_with_type(IconType::RGBA32_256x256)
            .unwrap();
        assert_eq!(image.data(), png_data.as_slice());
    }
}