1mod filter;
2mod inflate;
3#[cfg(test)]
4mod tests;
5
6use crate::media::{self, Error, Result};
7use filter::{Simd, unfilter};
8
9const SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];
10const MAX_PIXELS: u64 = 1 << 28;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Color {
14 Grey,
15 Rgb,
16 Palette,
17 GreyAlpha,
18 Rgba,
19}
20
21impl Color {
22 fn from_byte(b: u8) -> Option<Self> {
23 Some(match b {
24 0 => Color::Grey,
25 2 => Color::Rgb,
26 3 => Color::Palette,
27 4 => Color::GreyAlpha,
28 6 => Color::Rgba,
29 _ => return None,
30 })
31 }
32
33 fn samples(self) -> usize {
34 match self {
35 Color::Grey | Color::Palette => 1,
36 Color::GreyAlpha => 2,
37 Color::Rgb => 3,
38 Color::Rgba => 4,
39 }
40 }
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub struct Info {
45 pub width: u32,
46 pub height: u32,
47 pub depth: u8,
48 pub color: Color,
49 pub interlaced: bool,
50}
51
52fn invalid(what: impl std::fmt::Display) -> Error {
53 Error::Invalid(format!("PNG: {what}"))
54}
55
56fn be32(b: &[u8], at: usize) -> Result<u32> {
57 b.get(at..at + 4)
58 .map(|s| u32::from_be_bytes([s[0], s[1], s[2], s[3]]))
59 .ok_or_else(|| invalid("the file ends inside a chunk header"))
60}
61
62fn header(data: &[u8]) -> Result<Info> {
63 if data.get(..8) != Some(&SIGNATURE[..]) {
64 return Err(invalid("not a PNG file"));
65 }
66 if be32(data, 8)? != 13 || data.get(12..16) != Some(b"IHDR") {
67 return Err(invalid("the first chunk is not a 13-byte IHDR"));
68 }
69 let body = data
70 .get(16..29)
71 .ok_or_else(|| invalid("the file ends inside IHDR"))?;
72 let width = be32(body, 0)?;
73 let height = be32(body, 4)?;
74 let depth = body[8];
75 let color = Color::from_byte(body[9])
76 .ok_or_else(|| invalid(format!("unknown colour type {}", body[9])))?;
77 if width == 0 || height == 0 {
78 return Err(invalid("a zero width or height"));
79 }
80 if u64::from(width) * u64::from(height) > MAX_PIXELS {
81 return Err(invalid(format!("{width}x{height} is too large")));
82 }
83 let depth_ok = match color {
84 Color::Grey => matches!(depth, 1 | 2 | 4 | 8 | 16),
85 Color::Palette => matches!(depth, 1 | 2 | 4 | 8),
86 _ => matches!(depth, 8 | 16),
87 };
88 if !depth_ok {
89 return Err(invalid(format!(
90 "bit depth {depth} with colour type {}",
91 body[9]
92 )));
93 }
94 if body[10] != 0 || body[11] != 0 {
95 return Err(invalid("an unknown compression or filter method"));
96 }
97 let interlaced = match body[12] {
98 0 => false,
99 1 => true,
100 n => return Err(invalid(format!("unknown interlace method {n}"))),
101 };
102 Ok(Info {
103 width,
104 height,
105 depth,
106 color,
107 interlaced,
108 })
109}
110
111pub fn info(data: &[u8]) -> Result<Info> {
112 header(data)
113}
114
115pub fn decode(data: &[u8]) -> Result<media::Rgba> {
116 let mut out = media::Rgba::new(0, 0);
117 Decoder::new().decode_into(data, &mut out)?;
118 Ok(out)
119}
120
121#[derive(Clone, Copy)]
122struct Entry([u8; 4]);
123
124struct Layout {
125 info: Info,
126 palette: [Entry; 256],
127 has_palette: bool,
128 trns: Option<[u16; 3]>,
129}
130
131impl Layout {
132 fn row_bytes(&self, width: usize) -> usize {
133 (width * self.info.color.samples() * usize::from(self.info.depth)).div_ceil(8)
134 }
135
136 fn bpp(&self) -> usize {
137 (self.info.color.samples() * usize::from(self.info.depth)).div_ceil(8)
138 }
139}
140
141pub struct Decoder {
142 idat: Vec<u8>,
143 inflated: Vec<u8>,
144 prev: Vec<u8>,
145 cur: Vec<u8>,
146 unpacked: Vec<u8>,
147 pass: Vec<u8>,
148 simd: Simd,
149}
150
151impl Default for Decoder {
152 fn default() -> Self {
153 Self::new()
154 }
155}
156
157const XORIG: [usize; 7] = [0, 4, 0, 2, 0, 1, 0];
158const YORIG: [usize; 7] = [0, 0, 4, 0, 2, 0, 1];
159const XSPC: [usize; 7] = [8, 8, 4, 4, 2, 2, 1];
160const YSPC: [usize; 7] = [8, 8, 8, 4, 4, 2, 2];
161
162impl Decoder {
163 pub fn new() -> Self {
164 Decoder {
165 idat: Vec::new(),
166 inflated: Vec::new(),
167 prev: Vec::new(),
168 cur: Vec::new(),
169 unpacked: Vec::new(),
170 pass: Vec::new(),
171 simd: Simd::detect(),
172 }
173 }
174
175 pub fn decode_into(&mut self, data: &[u8], out: &mut media::Rgba) -> Result<()> {
176 let info = header(data)?;
177 let mut layout = Layout {
178 info,
179 palette: [Entry([0, 0, 0, 255]); 256],
180 has_palette: false,
181 trns: None,
182 };
183 let mut single: Option<&[u8]> = None;
184 let mut parts = 0usize;
185 self.idat.clear();
186 let mut at = 33usize;
187 let mut ended = false;
188 while at + 8 <= data.len() {
189 let len = be32(data, at)? as usize;
190 let kind = &data[at + 4..at + 8];
191 let body = data
192 .get(at + 8..at + 8 + len)
193 .ok_or_else(|| invalid("a chunk runs past the end of the file"))?;
194 at = at + 12 + len;
195 match kind {
196 b"IDAT" => {
197 parts += 1;
198 match (parts, single) {
199 (1, _) => single = Some(body),
200 (2, Some(first)) => {
201 self.idat.extend_from_slice(first);
202 self.idat.extend_from_slice(body);
203 }
204 _ => self.idat.extend_from_slice(body),
205 }
206 }
207 b"PLTE" => {
208 if !len.is_multiple_of(3) || len > 768 {
209 return Err(invalid("a PLTE chunk of the wrong length"));
210 }
211 for (e, rgb) in layout.palette.iter_mut().zip(body.chunks_exact(3)) {
212 e.0[..3].copy_from_slice(rgb);
213 }
214 layout.has_palette = true;
215 }
216 b"tRNS" => match info.color {
217 Color::Palette => {
218 for (e, a) in layout.palette.iter_mut().zip(body) {
219 e.0[3] = *a;
220 }
221 }
222 Color::Grey if len >= 2 => {
223 let g = u16::from_be_bytes([body[0], body[1]]);
224 layout.trns = Some([g, g, g]);
225 }
226 Color::Rgb if len >= 6 => {
227 let s = |i: usize| u16::from_be_bytes([body[i], body[i + 1]]);
228 layout.trns = Some([s(0), s(2), s(4)]);
229 }
230 _ => {}
231 },
232 b"IEND" => {
233 ended = true;
234 break;
235 }
236 _ => {}
237 }
238 }
239 if parts == 0 {
240 return Err(invalid(if ended {
241 "no image data"
242 } else {
243 "the file ends before any image data"
244 }));
245 }
246 if info.color == Color::Palette && !layout.has_palette {
247 return Err(invalid("a palette image without PLTE"));
248 }
249 let (w, h) = (info.width as usize, info.height as usize);
250 let expected = if info.interlaced {
251 (0..7)
252 .map(|p| {
253 let (x, y) = pass_size(p, w, h);
254 if x == 0 || y == 0 {
255 0
256 } else {
257 (layout.row_bytes(x) + 1) * y
258 }
259 })
260 .sum()
261 } else {
262 (layout.row_bytes(w) + 1) * h
263 };
264 let compressed = match single {
265 Some(s) if parts == 1 => s,
266 _ => &self.idat[..],
267 };
268 let written = inflate::zlib(
269 compressed,
270 &mut self.inflated,
271 expected,
272 expected + 4 * h + 64,
273 )?;
274 if written < expected {
275 return Err(invalid(format!(
276 "the image data ends early ({written} of {expected} bytes)"
277 )));
278 }
279 out.width = info.width;
280 out.height = info.height;
281 out.data.resize(w * h * 4, 0);
282 let rows = std::mem::take(&mut self.inflated);
283 let result = if info.interlaced {
284 self.interlaced(&layout, &rows, out)
285 } else {
286 self.rows(&layout, &rows, w, h, &mut |y, rgba| {
287 out.data[y * w * 4..(y + 1) * w * 4].copy_from_slice(rgba);
288 })
289 };
290 self.inflated = rows;
291 result
292 }
293
294 fn rows(
295 &mut self,
296 layout: &Layout,
297 data: &[u8],
298 width: usize,
299 height: usize,
300 put: &mut dyn FnMut(usize, &[u8]),
301 ) -> Result<()> {
302 let stride = layout.row_bytes(width);
303 let bpp = layout.bpp();
304 self.prev.clear();
305 self.prev.resize(stride, 0);
306 self.cur.clear();
307 self.cur.resize(stride, 0);
308 self.pass.resize(width * 4, 0);
309 let mut pass = std::mem::take(&mut self.pass);
310 let mut result = Ok(());
311 for (y, line) in data.chunks_exact(stride + 1).take(height).enumerate() {
312 if !unfilter(
313 line[0],
314 y == 0,
315 &self.prev,
316 &line[1..],
317 &mut self.cur,
318 bpp,
319 self.simd,
320 ) {
321 result = Err(invalid(format!("unknown filter type {}", line[0])));
322 break;
323 }
324 self.expand(layout, width, &mut pass[..width * 4]);
325 put(y, &pass[..width * 4]);
326 std::mem::swap(&mut self.prev, &mut self.cur);
327 }
328 self.pass = pass;
329 result
330 }
331
332 fn interlaced(&mut self, layout: &Layout, data: &[u8], out: &mut media::Rgba) -> Result<()> {
333 let (w, h) = (out.width as usize, out.height as usize);
334 let mut offset = 0usize;
335 for p in 0..7 {
336 let (x, y) = pass_size(p, w, h);
337 if x == 0 || y == 0 {
338 continue;
339 }
340 let len = (layout.row_bytes(x) + 1) * y;
341 let slice = data
342 .get(offset..offset + len)
343 .ok_or_else(|| invalid("the interlaced data ends early"))?;
344 let pixels = &mut out.data;
345 self.rows(layout, slice, x, y, &mut |j, rgba| {
346 let oy = j * YSPC[p] + YORIG[p];
347 for (i, px) in rgba.chunks_exact(4).enumerate() {
348 let ox = i * XSPC[p] + XORIG[p];
349 let at = (oy * w + ox) * 4;
350 pixels[at..at + 4].copy_from_slice(px);
351 }
352 })?;
353 offset += len;
354 }
355 Ok(())
356 }
357
358 fn expand(&mut self, layout: &Layout, width: usize, out: &mut [u8]) {
359 let info = layout.info;
360 let cur = &self.cur;
361 if info.depth < 8 {
362 let scale = if info.color == Color::Palette {
363 1
364 } else {
365 match info.depth {
366 1 => 0xff,
367 2 => 0x55,
368 _ => 0x11,
369 }
370 };
371 self.unpacked.resize(width, 0);
372 unpack_bits(info.depth, scale, cur, &mut self.unpacked[..width]);
373 let src = &self.unpacked[..width];
374 if info.color == Color::Palette {
375 for (px, &i) in out.chunks_exact_mut(4).zip(src) {
376 px.copy_from_slice(&layout.palette[usize::from(i)].0);
377 }
378 } else {
379 let key = layout.trns.map(|t| {
380 let mask = (1u16 << info.depth) - 1;
381 (t[0] & mask) as u8 * scale
382 });
383 for (px, &g) in out.chunks_exact_mut(4).zip(src) {
384 let a = if key == Some(g) { 0 } else { 255 };
385 px.copy_from_slice(&[g, g, g, a]);
386 }
387 }
388 return;
389 }
390 let sixteen = info.depth == 16;
391 match (info.color, sixteen) {
392 (Color::Rgba, false) => out.copy_from_slice(&cur[..width * 4]),
393 (Color::Rgb, false) => match layout.trns {
394 None => {
395 for (px, rgb) in out.chunks_exact_mut(4).zip(cur.chunks_exact(3)) {
396 px.copy_from_slice(&[rgb[0], rgb[1], rgb[2], 255]);
397 }
398 }
399 Some(t) => {
400 let key = [t[0] as u8, t[1] as u8, t[2] as u8];
401 let key = if t.iter().all(|v| *v < 256) {
402 Some(key)
403 } else {
404 None
405 };
406 for (px, rgb) in out.chunks_exact_mut(4).zip(cur.chunks_exact(3)) {
407 let a = if Some([rgb[0], rgb[1], rgb[2]]) == key {
408 0
409 } else {
410 255
411 };
412 px.copy_from_slice(&[rgb[0], rgb[1], rgb[2], a]);
413 }
414 }
415 },
416 (Color::Palette, _) => {
417 for (px, &i) in out.chunks_exact_mut(4).zip(cur.iter()) {
418 px.copy_from_slice(&layout.palette[usize::from(i)].0);
419 }
420 }
421 (Color::Grey, false) => {
422 let key = layout.trns.and_then(|t| u8::try_from(t[0]).ok());
423 for (px, &g) in out.chunks_exact_mut(4).zip(cur.iter()) {
424 let a = if Some(g) == key { 0 } else { 255 };
425 px.copy_from_slice(&[g, g, g, a]);
426 }
427 }
428 (Color::GreyAlpha, false) => {
429 for (px, ga) in out.chunks_exact_mut(4).zip(cur.chunks_exact(2)) {
430 px.copy_from_slice(&[ga[0], ga[0], ga[0], ga[1]]);
431 }
432 }
433 (Color::Rgba, true) => {
434 for (px, s) in out.chunks_exact_mut(4).zip(cur.chunks_exact(8)) {
435 px.copy_from_slice(&[s[0], s[2], s[4], s[6]]);
436 }
437 }
438 (Color::Rgb, true) => {
439 let key = layout.trns;
440 for (px, s) in out.chunks_exact_mut(4).zip(cur.chunks_exact(6)) {
441 let v = |i: usize| u16::from_be_bytes([s[i], s[i + 1]]);
442 let a = if key == Some([v(0), v(2), v(4)]) {
443 0
444 } else {
445 255
446 };
447 px.copy_from_slice(&[s[0], s[2], s[4], a]);
448 }
449 }
450 (Color::Grey, true) => {
451 let key = layout.trns.map(|t| t[0]);
452 for (px, s) in out.chunks_exact_mut(4).zip(cur.chunks_exact(2)) {
453 let a = if key == Some(u16::from_be_bytes([s[0], s[1]])) {
454 0
455 } else {
456 255
457 };
458 px.copy_from_slice(&[s[0], s[0], s[0], a]);
459 }
460 }
461 (Color::GreyAlpha, true) => {
462 for (px, s) in out.chunks_exact_mut(4).zip(cur.chunks_exact(4)) {
463 px.copy_from_slice(&[s[0], s[0], s[0], s[2]]);
464 }
465 }
466 }
467 }
468}
469
470fn pass_size(p: usize, w: usize, h: usize) -> (usize, usize) {
471 (
472 w.saturating_sub(XORIG[p]).div_ceil(XSPC[p]),
473 h.saturating_sub(YORIG[p]).div_ceil(YSPC[p]),
474 )
475}
476
477fn unpack_bits(depth: u8, scale: u8, input: &[u8], out: &mut [u8]) {
478 let per = usize::from(8 / depth);
479 let mask = (1u8 << depth) - 1;
480 for (bytes, &v) in out.chunks_mut(per).zip(input) {
481 for (k, o) in bytes.iter_mut().enumerate() {
482 let shift = 8 - depth as usize * (k + 1);
483 *o = scale.wrapping_mul((v >> shift) & mask);
484 }
485 }
486}