pngenc/
lib.rs

1//! uncompressing png encoding crate
2//! ```
3//! # let mut v = vec![];
4//! pngenc::ode(
5//!   pngenc::RGB, // color type
6//!   (2144, 1424), // width and height
7//!   include_bytes!("../benches/cat"), // image to encode
8//! # /*
9//!   &mut std::fs::File::create("hey.png").unwrap(), // output writer
10//! # */
11//! # &mut v,
12//! ).unwrap();
13//! # assert_eq!(crc32fast::hash(&v), 0xd7d47c0e);
14//! ```
15#![allow(incomplete_features)]
16#![feature(generic_const_exprs, test, slice_as_chunks, array_chunks)]
17use atools::prelude::*;
18use std::{
19    io::{self, Write},
20    iter::once,
21};
22
23pub use Color::*;
24
25#[derive(Copy, Debug, Clone)]
26#[repr(u8)]
27/// Color types.
28pub enum Color {
29    /// Grayscale
30    Y = 1,
31    /// Grayscale with alpha
32    YA,
33    /// Red, green, blue
34    RGB,
35    /// RGB with alpha
36    RGBA,
37}
38
39impl Color {
40    /// Color depth (number of channels)
41    #[must_use]
42    pub const fn depth(self) -> u8 {
43        self as u8
44    }
45
46    const fn ty(self) -> u8 {
47        match self {
48            Color::Y => 0,
49            Color::YA => 4,
50            Color::RGB => 2,
51            Color::RGBA => 6,
52        }
53    }
54}
55
56trait W: Write {
57    fn u32(&mut self, x: u32) -> io::Result<()> {
58        self.write_all(&x.to_be_bytes())
59    }
60    fn w(&mut self, x: impl AsRef<[u8]>) -> io::Result<()> {
61        self.write_all(x.as_ref())
62    }
63}
64
65impl<T: Write> W for T {}
66
67const HEADER: &[u8; 8] = b"\x89PNG\x0d\x0a\x1a\x0a";
68
69fn chunk_len(x: usize) -> usize {
70    4 // length
71        + 4 // type
72        + x // data
73        + 4 // crc
74}
75
76/// Get the size of an encoded png. Guaranteed to exactly equal the size of the encoded png.
77pub fn size(color: Color, (width, height): (u32, u32)) -> usize {
78    HEADER.len()
79        + chunk_len(13) // IHDR
80        + chunk_len(deflate_size(((width * color.depth() as u32 + 1) * height) as usize)) // IDAT
81        + chunk_len(1) // sRGB
82        + chunk_len(0) // IEND
83}
84
85#[doc(alias = "encode")]
86/// Encode a png without any compression.
87/// Takes advantage of the [Non-compressed blocks](http://www.zlib.org/rfc-deflate.html#noncompressed) deflate feature.
88///
89/// If you *do* want a compressed image, I recommend the [oxipng](https://docs.rs/oxipng/latest/oxipng/struct.RawImage.html) raw image api.
90///
91/// # Panics
92///
93/// if your width * height * color depth isnt data's length
94pub fn ode(
95    color: Color,
96    (width, height): (u32, u32),
97    data: &[u8],
98    to: &mut impl Write,
99) -> std::io::Result<()> {
100    assert_eq!(
101        (width as usize * height as usize)
102            .checked_mul(color.depth() as usize)
103            .unwrap(),
104        data.len(),
105        "please dont lie to me"
106    );
107    to.w(HEADER)?;
108    chunk(
109        *b"IHDR",
110        &width
111            .to_be_bytes()
112            .couple(height.to_be_bytes())
113            .join(8) // bit depth
114            .join(color.ty())
115            .join(0)
116            .join(0)
117            .join(0),
118        to,
119    )?;
120
121    // removing this allocation is not a performance gain
122    let mut scanned = Vec::<u8>::with_capacity(((width * color.depth() as u32 + 1) * height) as _);
123    let mut out = scanned.as_mut_ptr();
124
125    data.chunks(width as usize * color.depth() as usize)
126        // set filter type for each scanline
127        .flat_map(|x| once(0).chain(x.iter().copied()))
128        .for_each(|x| unsafe {
129            out.write(x);
130            out = out.add(1);
131        });
132    unsafe { scanned.set_len(((width * color.depth() as u32 + 1) * height) as _) };
133
134    let data = deflate(&scanned);
135    chunk(*b"sRGB", &[0], to)?;
136    chunk(*b"IDAT", &data, to)?;
137    chunk(*b"IEND", &[], to)?;
138    Ok(())
139}
140
141fn chunk(ty: [u8; 4], data: &[u8], to: &mut impl Write) -> std::io::Result<()> {
142    to.u32(data.len() as _)?;
143    to.w(ty)?;
144    to.w(data)?;
145    let mut crc = crc32fast::Hasher::new();
146    crc.update(&ty);
147    crc.update(data);
148    to.u32(crc.finalize())?;
149    Ok(())
150}
151
152fn deflate_size(x: usize) -> usize {
153    // 2 header bytes, each header of chunk, and add remainder chunk, along with 4 bytes for adler32
154    2 + 5 * (x / CHUNK_SIZE) + usize::from(x != (x / CHUNK_SIZE) * CHUNK_SIZE || x == 0) + x + 4 + 4
155}
156
157trait P<T: Copy> {
158    unsafe fn put<const N: usize>(&mut self, x: [T; N]);
159}
160
161impl<T: Copy> P<T> for *mut T {
162    #[cfg_attr(debug_assertions, track_caller)]
163    unsafe fn put<const N: usize>(&mut self, x: [T; N]) {
164        self.copy_from(x.as_ptr(), N);
165        *self = self.add(N);
166    }
167}
168
169const CHUNK_SIZE: usize = 0xffff;
170fn deflate(data: &[u8]) -> Vec<u8> {
171    let mut adler = simd_adler32::Adler32::new();
172    let (chunks, remainder) = data.as_chunks::<CHUNK_SIZE>();
173    // SAFETY: deflate_size is very correct.
174    let mut out = Vec::<u8>::with_capacity(deflate_size(data.len()));
175    let mut optr = out.as_mut_ptr();
176    /// return LSB and SLSB
177    fn split(n: u16) -> [u8; 2] {
178        [(n & 0xff) as u8, ((n >> 8) & 0xff) as u8]
179    }
180    // 32k window
181    unsafe { optr.put([0b1_111_000, 1]) };
182    chunks.iter().for_each(|x| unsafe {
183        adler.write(x);
184        // http://www.zlib.org/rfc-deflate.html#noncompressed
185        optr.put(
186            [0b000]
187                // lsb and slsb [255, 255]
188                .couple(split(CHUNK_SIZE as _))
189                // ones complement -- [0, 0]
190                .couple(split(CHUNK_SIZE as _).map(|x| !x)),
191        );
192        optr.put(*x);
193    });
194    unsafe {
195        adler.write(remainder);
196        optr.put(
197            [0b001]
198                .couple(split(CHUNK_SIZE as _))
199                .couple(split(CHUNK_SIZE as _).map(|x| !x)),
200        );
201        optr.copy_from(remainder.as_ptr(), remainder.len());
202        optr = optr.add(remainder.len());
203    };
204    unsafe { optr.put(adler.finish().to_be_bytes()) };
205    unsafe { out.set_len(deflate_size(data.len())) }
206    out
207}
pngenc/lib.rs

pngenc/
lib.rs
